@Test public void testLoadbalancingStrategyRemovesTasksWhenDatastreamIsDeleted() { String[] instances = new String[]{"instance1", "instance2", "instance3"}; List<Datastream> datastreams = Arrays.asList(DatastreamTestUtils.createDatastreams(DummyConnector.CONNECTOR_TYPE, "ds1", "ds2")); datastreams.forEach(x -> x.getSource().setPartitions(12)); datastreams.forEach( x -> x.getMetadata().put(DatastreamMetadataConstants.TASK_PREFIX, DatastreamTaskImpl.getTaskPrefix(x))); List<DatastreamGroup> dgs = datastreams.stream().map(x -> new DatastreamGroup(Collections.singletonList(x))).collect(Collectors.toList()); LoadbalancingStrategy strategy = new LoadbalancingStrategy(); ZkAdapter adapter = createMockAdapter(); Map<String, Set<DatastreamTask>> assignment = strategy.assign(dgs, Arrays.asList(instances), new HashMap<>()); for (String instance : instances) { Assert.assertEquals(assignment.get(instance).size(), 4); assignment.get(instance).forEach(t -> ((DatastreamTaskImpl) t).setZkAdapter(adapter)); } dgs = new ArrayList<>(dgs); dgs.remove(1); Map<String, Set<DatastreamTask>> newAssignment = strategy.assign(dgs, Arrays.asList(instances), assignment); for (String instance : instances) { Assert.assertEquals(newAssignment.get(instance).size(), 2); } }
@Override public Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { LOG.info("Assign called with datastreams: {}, instances: {}, currentAssignment: {}", datastreams, instances, currentAssignment); if (instances.size() == 0) { return new HashMap<>(); } int maxTasksPerDatastream = Math.max(instances.size() * overPartitioningFactor, minTasks); List<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); datastreamTasks.sort(Comparator.comparing(DatastreamTask::getDatastreamTaskName)); instances.sort(String::compareTo); Map<String, Set<DatastreamTask>> assignment = new HashMap<>(); instances.forEach(i -> assignment.put(i, new HashSet<>())); int datastreamTaskLength = datastreamTasks.size(); for (int i = 0; i < datastreamTaskLength; i++) { String instanceName = instances.get(i % instances.size()); assignment.get(instanceName).add(datastreamTasks.get(i)); } LOG.info("Datastream Groups: {}, instances: {}, currentAssignment: {}, NewAssignment: {}", datastreams.stream().map(DatastreamGroup::getTaskPrefix).collect(Collectors.toList()), instances, currentAssignment, assignment); return assignment; }
LoadbalancingStrategy implements AssignmentStrategy { @Override public Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { LOG.info("Assign called with datastreams: {}, instances: {}, currentAssignment: {}", datastreams, instances, currentAssignment); if (instances.size() == 0) { return new HashMap<>(); } int maxTasksPerDatastream = Math.max(instances.size() * overPartitioningFactor, minTasks); List<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); datastreamTasks.sort(Comparator.comparing(DatastreamTask::getDatastreamTaskName)); instances.sort(String::compareTo); Map<String, Set<DatastreamTask>> assignment = new HashMap<>(); instances.forEach(i -> assignment.put(i, new HashSet<>())); int datastreamTaskLength = datastreamTasks.size(); for (int i = 0; i < datastreamTaskLength; i++) { String instanceName = instances.get(i % instances.size()); assignment.get(instanceName).add(datastreamTasks.get(i)); } LOG.info("Datastream Groups: {}, instances: {}, currentAssignment: {}, NewAssignment: {}", datastreams.stream().map(DatastreamGroup::getTaskPrefix).collect(Collectors.toList()), instances, currentAssignment, assignment); return assignment; } }
LoadbalancingStrategy implements AssignmentStrategy { @Override public Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { LOG.info("Assign called with datastreams: {}, instances: {}, currentAssignment: {}", datastreams, instances, currentAssignment); if (instances.size() == 0) { return new HashMap<>(); } int maxTasksPerDatastream = Math.max(instances.size() * overPartitioningFactor, minTasks); List<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); datastreamTasks.sort(Comparator.comparing(DatastreamTask::getDatastreamTaskName)); instances.sort(String::compareTo); Map<String, Set<DatastreamTask>> assignment = new HashMap<>(); instances.forEach(i -> assignment.put(i, new HashSet<>())); int datastreamTaskLength = datastreamTasks.size(); for (int i = 0; i < datastreamTaskLength; i++) { String instanceName = instances.get(i % instances.size()); assignment.get(instanceName).add(datastreamTasks.get(i)); } LOG.info("Datastream Groups: {}, instances: {}, currentAssignment: {}, NewAssignment: {}", datastreams.stream().map(DatastreamGroup::getTaskPrefix).collect(Collectors.toList()), instances, currentAssignment, assignment); return assignment; } LoadbalancingStrategy();  LoadbalancingStrategy(Properties properties); }
LoadbalancingStrategy implements AssignmentStrategy { @Override public Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { LOG.info("Assign called with datastreams: {}, instances: {}, currentAssignment: {}", datastreams, instances, currentAssignment); if (instances.size() == 0) { return new HashMap<>(); } int maxTasksPerDatastream = Math.max(instances.size() * overPartitioningFactor, minTasks); List<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); datastreamTasks.sort(Comparator.comparing(DatastreamTask::getDatastreamTaskName)); instances.sort(String::compareTo); Map<String, Set<DatastreamTask>> assignment = new HashMap<>(); instances.forEach(i -> assignment.put(i, new HashSet<>())); int datastreamTaskLength = datastreamTasks.size(); for (int i = 0; i < datastreamTaskLength; i++) { String instanceName = instances.get(i % instances.size()); assignment.get(instanceName).add(datastreamTasks.get(i)); } LOG.info("Datastream Groups: {}, instances: {}, currentAssignment: {}, NewAssignment: {}", datastreams.stream().map(DatastreamGroup::getTaskPrefix).collect(Collectors.toList()), instances, currentAssignment, assignment); return assignment; } LoadbalancingStrategy();  LoadbalancingStrategy(Properties properties); @Override Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances,
      Map<String, Set<DatastreamTask>> currentAssignment); }
LoadbalancingStrategy implements AssignmentStrategy { @Override public Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances, Map<String, Set<DatastreamTask>> currentAssignment) { LOG.info("Assign called with datastreams: {}, instances: {}, currentAssignment: {}", datastreams, instances, currentAssignment); if (instances.size() == 0) { return new HashMap<>(); } int maxTasksPerDatastream = Math.max(instances.size() * overPartitioningFactor, minTasks); List<DatastreamTask> datastreamTasks = getDatastreamTasks(datastreams, currentAssignment, maxTasksPerDatastream); datastreamTasks.sort(Comparator.comparing(DatastreamTask::getDatastreamTaskName)); instances.sort(String::compareTo); Map<String, Set<DatastreamTask>> assignment = new HashMap<>(); instances.forEach(i -> assignment.put(i, new HashSet<>())); int datastreamTaskLength = datastreamTasks.size(); for (int i = 0; i < datastreamTaskLength; i++) { String instanceName = instances.get(i % instances.size()); assignment.get(instanceName).add(datastreamTasks.get(i)); } LOG.info("Datastream Groups: {}, instances: {}, currentAssignment: {}, NewAssignment: {}", datastreams.stream().map(DatastreamGroup::getTaskPrefix).collect(Collectors.toList()), instances, currentAssignment, assignment); return assignment; } LoadbalancingStrategy();  LoadbalancingStrategy(Properties properties); @Override Map<String, Set<DatastreamTask>> assign(List<DatastreamGroup> datastreams, List<String> instances,
      Map<String, Set<DatastreamTask>> currentAssignment); static final int DEFAULT_OVER_PARTITIONING_FACTOR; static final String CFG_OVER_PARTITIONING_FACTOR; static final String CFG_MIN_TASKS; static final int DEFAULT_MIN_TASKS; }