Source:
<code>public ScheduleResult schedule() { List<RemoteTask> newTasks = ImmutableList.of(); if (!scheduledTasks) { <START> OptionalInt totalPartitions = OptionalInt.of(stageNodeList.size()); <END> newTasks = Streams.mapWithIndex( stageNodeList.stream(), (node, taskId) -> stage.scheduleTask(node, toIntExact(taskId), totalPartitions)) .collect(toImmutableList()); scheduledTasks = true; } boolean allBlocked = true; List<ListenableFuture<?>> blocked = new ArrayList<>(); BlockedReason blockedReason = BlockedReason.NO_ACTIVE_DRIVER_GROUP; if (groupedLifespanScheduler.isPresent()) { blocked.add(groupedLifespanScheduler.get().schedule(sourceSchedulers.get(0))); } int splitsScheduled = 0; Iterator<SourceScheduler> schedulerIterator = sourceSchedulers.iterator(); List<Lifespan> driverGroupsToStart = ImmutableList.of(); boolean shouldInvokeNoMoreDriverGroups = false; while (schedulerIterator.hasNext()) { SourceScheduler sourceScheduler = schedulerIterator.next(); for (Lifespan lifespan : driverGroupsToStart) { sourceScheduler.startLifespan(lifespan, partitionHandleFor(lifespan)); } if (shouldInvokeNoMoreDriverGroups) { sourceScheduler.noMoreLifespans(); } ScheduleResult schedule = sourceScheduler.schedule(); splitsScheduled += schedule.getSplitsScheduled(); if (schedule.getBlockedReason().isPresent()) { blocked.add(schedule.getBlocked()); blockedReason = blockedReason.combineWith(schedule.getBlockedReason().get()); } else { verify(schedule.getBlocked().isDone(), "blockedReason not provided when scheduler is blocked"); allBlocked = false; } driverGroupsToStart = sourceScheduler.drainCompletedLifespans(); if (schedule.isFinished()) { stage.schedulingComplete(sourceScheduler.getPlanNodeId()); schedulerIterator.remove(); sourceScheduler.close(); shouldInvokeNoMoreDriverGroups = true; } else { shouldInvokeNoMoreDriverGroups = false; } } if (allBlocked) { return new ScheduleResult(sourceSchedulers.isEmpty(), newTasks, whenAnyComplete(blocked), blockedReason, splitsScheduled); } else { return new ScheduleResult(sourceSchedulers.isEmpty(), newTasks, splitsScheduled); } }</code><technical_language>totalPartitions number of nodes? term partition is overloaded confuse readers me</technical_language>

Target:
public ScheduleResult schedule() { List<RemoteTask> newTasks = ImmutableList.of(); if (!scheduledTasks) { OptionalInt totalPartitions = OptionalInt.of(nodes.size()); newTasks = Streams.mapWithIndex( nodes.stream(), (node, id) -> stage.scheduleTask(node, toIntExact(id), totalPartitions)) .collect(toImmutableList()); scheduledTasks = true; } boolean allBlocked = true; List<ListenableFuture<?>> blocked = new ArrayList<>(); BlockedReason blockedReason = BlockedReason.NO_ACTIVE_DRIVER_GROUP; if (groupedLifespanScheduler.isPresent()) { blocked.add(groupedLifespanScheduler.get().schedule(sourceSchedulers.get(0))); } int splitsScheduled = 0; Iterator<SourceScheduler> schedulerIterator = sourceSchedulers.iterator(); List<Lifespan> driverGroupsToStart = ImmutableList.of(); boolean shouldInvokeNoMoreDriverGroups = false; while (schedulerIterator.hasNext()) { SourceScheduler sourceScheduler = schedulerIterator.next(); for (Lifespan lifespan : driverGroupsToStart) { sourceScheduler.startLifespan(lifespan, partitionHandleFor(lifespan)); } if (shouldInvokeNoMoreDriverGroups) { sourceScheduler.noMoreLifespans(); } ScheduleResult schedule = sourceScheduler.schedule(); splitsScheduled += schedule.getSplitsScheduled(); if (schedule.getBlockedReason().isPresent()) { blocked.add(schedule.getBlocked()); blockedReason = blockedReason.combineWith(schedule.getBlockedReason().get()); } else { verify(schedule.getBlocked().isDone(), "blockedReason not provided when scheduler is blocked"); allBlocked = false; } driverGroupsToStart = sourceScheduler.drainCompletedLifespans(); if (schedule.isFinished()) { stage.schedulingComplete(sourceScheduler.getPlanNodeId()); schedulerIterator.remove(); sourceScheduler.close(); shouldInvokeNoMoreDriverGroups = true; } else { shouldInvokeNoMoreDriverGroups = false; } } if (allBlocked) { return new ScheduleResult(sourceSchedulers.isEmpty(), newTasks, whenAnyComplete(blocked), blockedReason, splitsScheduled); } else { return new ScheduleResult(sourceSchedulers.isEmpty(), newTasks, splitsScheduled); } }