Source:
public Map<String, Set<DatastreamTask>> assign(List<Datastream> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { int maxTasksPerDatastream = instances.size() * PARTITIONING_FACTOR; Set<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); if (instances.size() == 0) { return new HashMap<>(); } List<DatastreamTask> allTasks = new ArrayList<>(); allTasks.addAll(datastreamTasks); Collections.shuffle(allTasks); int numberOfTasksPerInstance = (int) ceil(datastreamTasks.size() * 1.0 /instances.size()); Map<String, Set<DatastreamTask>> newAssignment = new HashMap<>(); Collections.sort(instances); for(String instance : instances) { if(currentAssignment.containsKey(instance)) { int addedTasksCount = 0; Set<DatastreamTask> newAssignmentForInstance = new HashSet<>(); for(DatastreamTask task : currentAssignment.get(instance)) { if(addedTasksCount < numberOfTasksPerInstance) { addedTasksCount++; newAssignmentForInstance.add(task); datastreamTasks.remove(task); } } newAssignment.put(instance, newAssignmentForInstance); } else { newAssignment.put(instance, new HashSet<>()); } } int unAssignedTaskIndex = 0; for(String instance : instances) { Set<DatastreamTask> newAssignmentForInstance = newAssignment.get(instance); if(newAssignmentForInstance.size() == numberOfTasksPerInstance) { continue; } int numberOfTasksToAdd = numberOfTasksPerInstance - newAssignmentForInstance.size(); for(int index = 0; index < numberOfTasksToAdd; index++) { unAssignedTaskIndex = getNextUnAssignedTask(allTasks, unAssignedTaskIndex, datastreamTasks); if(unAssignedTaskIndex >= allTasks.size()) { break; } newAssignmentForInstance.add(allTasks.get(unAssignedTaskIndex)); unAssignedTaskIndex++; } } return newAssignment; }

Target:
public Map<String, Set<DatastreamTask>> assign(List<Datastream> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { int numberOfInstances = instances.size(); int maxTasksPerDatastream = instances.size() * PARTITIONING_FACTOR; Set<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); if (instances.size() == 0) { return new HashMap<>(); } int minTasksPerInstance = (int) floor(datastreamTasks.size() * 1.0 / instances.size()); int maxTasksPerInstance = datastreamTasks.size() % instances.size() > 0 ? minTasksPerInstance + 1 : minTasksPerInstance; Map<String, Set<DatastreamTask>> newAssignment = new HashMap<>(); Collections.sort(instances); for(String instance : instances) { if(currentAssignment.containsKey(instance)) { int addedTasksCount = 0; Set<DatastreamTask> newAssignmentForInstance = new HashSet<>(); for(DatastreamTask task : currentAssignment.get(instance)) { if(addedTasksCount < minTasksPerInstance) { addedTasksCount++; newAssignmentForInstance.add(task); datastreamTasks.remove(task); } } newAssignment.put(instance, newAssignmentForInstance); } else { newAssignment.put(instance, new HashSet<>()); } } int index = 0; for(DatastreamTask datastreamTask : datastreamTasks) { String instance = instances.get(index%numberOfInstances); if(newAssignment.get(instance).size() >= maxTasksPerInstance) { continue; } else { newAssignment.get(instance).add(datastreamTask); index++; } } LOG.info(String.format("Datastreams: %s, instances: %s, currentAssignment: %s, NewAssignment: %s", datastreams.stream().map(x -> x.getName()).collect(Collectors.toList()), instances, currentAssignment, newAssignment)); return newAssignment; }