Source:
public Set<ContainerModel> balance(Set<TaskModel> tasks, LocalityManager localityManager) { validateTasks(tasks); if (localityManager == null) { log.info("Locality manager is null. Cannot read or write task assignments. Invoking grouper."); return group(tasks); } TaskAssignmentManager taskAssignmentManager = localityManager.getTaskAssignmentManager(); try { List<TaskGroup> containers = getPreviousContainers(taskAssignmentManager, tasks.size()); if (containers == null || containers.size() == 1 || containerCount == 1) { log.info("Balancing does not apply. Invoking grouper."); Set<ContainerModel> models = group(tasks); saveTaskAssignments(models, taskAssignmentManager); return models; } int prevContainerCount = containers.size(); int containerDelta = containerCount - prevContainerCount; if (containerDelta == 0) { log.info("Container count has not changed. Reusing previous container models."); return buildContainerModels(tasks, containers); } log.info("Container count changed from to {}. Balancing tasks.", prevContainerCount, containerCount); int[] expectedTaskCountPerContainer = calculateTaskCountPerContainer(tasks.size(), prevContainerCount, containerCount); List<String> taskNamesToReassign = new LinkedList<>(); for (int i = 0; i < prevContainerCount; i++) { TaskGroup taskGroup = containers.get(i); while (taskGroup.size() > expectedTaskCountPerContainer[i]) { taskNamesToReassign.add(taskGroup.removeTask()); } } if (containerDelta > 0) { List<TaskGroup> newContainers = createContainers(prevContainerCount, containerCount); containers.addAll(newContainers); } else { containers = containers.subList(0, containerCount); } assignTasksToContainers(expectedTaskCountPerContainer, taskNamesToReassign, containers); Set<ContainerModel> models = buildContainerModels(tasks, containers); saveTaskAssignments(models, taskAssignmentManager); return models; } finally { taskAssignmentManager.close(); } }

Target:
public Set<ContainerModel> balance(Set<TaskModel> tasks, LocalityManager localityManager) { validateTasks(tasks); if (localityManager == null) { log.info("Locality manager is null. Cannot read or write task assignments. Invoking grouper."); return group(tasks); } TaskAssignmentManager taskAssignmentManager = new TaskAssignmentManager(config, new MetricsRegistryMap()); taskAssignmentManager.init(); try { List<TaskGroup> containers = getPreviousContainers(taskAssignmentManager, tasks.size()); if (containers == null || containers.size() == 1 || containerCount == 1) { log.info("Balancing does not apply. Invoking grouper."); Set<ContainerModel> models = group(tasks); saveTaskAssignments(models, taskAssignmentManager); return models; } int prevContainerCount = containers.size(); int containerDelta = containerCount - prevContainerCount; if (containerDelta == 0) { log.info("Container count has not changed. Reusing previous container models."); return buildContainerModels(tasks, containers); } log.info("Container count changed from to {}. Balancing tasks.", prevContainerCount, containerCount); int[] expectedTaskCountPerContainer = calculateTaskCountPerContainer(tasks.size(), prevContainerCount, containerCount); List<String> taskNamesToReassign = new LinkedList<>(); for (int i = 0; i < prevContainerCount; i++) { TaskGroup taskGroup = containers.get(i); while (taskGroup.size() > expectedTaskCountPerContainer[i]) { taskNamesToReassign.add(taskGroup.removeTask()); } } if (containerDelta > 0) { List<TaskGroup> newContainers = createContainers(prevContainerCount, containerCount); containers.addAll(newContainers); } else { containers = containers.subList(0, containerCount); } assignTasksToContainers(expectedTaskCountPerContainer, taskNamesToReassign, containers); Set<ContainerModel> models = buildContainerModels(tasks, containers); saveTaskAssignments(models, taskAssignmentManager); return models; } finally { taskAssignmentManager.close(); } }