ZkAdapter { public void updateAllAssignments(Map<String, List<DatastreamTask>> assignmentsByInstance) { Map<String, DatastreamTask> assignmentsMap = assignmentsByInstance.values() .stream() .flatMap(List::stream) .collect(Collectors.toMap(DatastreamTask::getDatastreamTaskName, Function.identity())); Map<String, Set<String>> nodesToRemove = new HashMap<>(); Map<String, Set<String>> nodesToAdd = new HashMap<>(); diffAssignmentNodes(assignmentsByInstance, nodesToRemove, nodesToAdd); for (String instance : nodesToAdd.keySet()) { Set<String> added = nodesToAdd.get(instance); if (added.size() > 0) { LOG.info("Instance: {}, adding assignments: {}", instance, added); for (String name : added) { addTaskNodes(instance, (DatastreamTaskImpl) assignmentsMap.get(name)); } } } for (String instance : nodesToRemove.keySet()) { Set<String> removed = nodesToRemove.get(instance); if (removed.size() > 0) { LOG.info("Instance: {}, removing assignments: {}", instance, removed); for (String name : removed) { removeTaskNodes(instance, name); } } } _liveTaskMap = new HashMap<>(); for (String instance : nodesToAdd.keySet()) { _liveTaskMap.put(instance, new HashSet<>(assignmentsByInstance.get(instance))); } } ZkAdapter(String zkServers, String cluster, String defaultTransportProviderName, int sessionTimeoutMs,
      int connectionTimeoutMs, int operationRetryTimeoutMs, long debounceTimerMs, ZkAdapterListener listener); @VisibleForTesting  ZkAdapter(String zkServers, String cluster, String defaultTransportProviderName, int sessionTimeoutMs,
      int connectionTimeoutMs, long debounceTimerMs, ZkAdapterListener listener); synchronized boolean isLeader(); String getInstanceName(); void disconnect(); void connect(); boolean updateDatastream(Datastream datastream); void deleteTasksWithPrefix(String connector, String taskPrefix); void deleteDatastream(String datastreamName); List<String> getAllInstances(); void touchAllInstanceAssignments(); List<String> getLiveInstances(); List<String> getInstanceAssignment(String instance); Map<String, Set<DatastreamTask>> getAllAssignedDatastreamTasks(); DatastreamTaskImpl getAssignedDatastreamTask(String instance, String taskName); void updateAllAssignments(Map<String, List<DatastreamTask>> assignmentsByInstance); void addConnectorType(String connectorType); void zkSaveInstanceError(String message); String getDatastreamTaskStateForKey(DatastreamTask datastreamTask, String key); void setDatastreamTaskStateForKey(DatastreamTask datastreamTask, String key, String value); void cleanUpDeadInstanceDataAndOtherUnusedTasks(
      Map<String, Set<DatastreamTask>> previousAssignmentByInstance,
      Map<String, List<DatastreamTask>> newAssignmentsByInstance,
      List<String> liveInstances); void cleanUpOldUnusedTasksFromConnector(Map<String, Set<DatastreamTask>> previousAssignmentByInstance,
      Map<String, List<DatastreamTask>> newAssignmentsByInstance); int cleanUpOrphanConnectorTasks(boolean cleanUpOrphanTasksInConnector); int cleanUpOrphanConnectorTaskLocks(boolean cleanUpOrphanTaskLocksInConnector); void acquireTask(DatastreamTaskImpl task, Duration timeout); boolean checkIsTaskLocked(String connectorType, String taskPrefix, String taskName); void waitForDependencies(DatastreamTaskImpl task, Duration timeout); void cleanUpPartitionMovement(String connectorType, String datastreamGroupName, long timestamp); List<String> getDatastreamsNeedPartitionMovement(String connectorType); Map<String, Set<String>> getPartitionMovement(String connectorType, String datastreamGroupName,
      long timestamp); void releaseTask(DatastreamTaskImpl task); static String parseHostnameFromZkInstance(String instance);  }
@Test public void testUpdateAllAssignments() throws Exception { String testCluster = "testUpdateAllInstanceAssignment"; String connectorType = "connectorType"; ZkClient zkClient = new ZkClient(_zkConnectionString); ZkAdapter adapter1 = createZkAdapter(testCluster); ZkAdapter adapter2 = createZkAdapter(testCluster); adapter1.connect(); adapter2.connect(); DatastreamTestUtils.createAndStoreDatastreams(zkClient, testCluster, connectorType, "task1", "task2", "task3"); DatastreamTaskImpl task1 = new DatastreamTaskImpl(); task1.setTaskPrefix("task1"); task1.setConnectorType(connectorType); DatastreamTaskImpl task2 = new DatastreamTaskImpl(); task2.setTaskPrefix("task2"); task2.setConnectorType(connectorType); DatastreamTaskImpl task3 = new DatastreamTaskImpl(); task3.setTaskPrefix("task3"); task3.setConnectorType(connectorType); Map<String, List<DatastreamTask>> assignmentsByInstance = new HashMap<>(); assignmentsByInstance.put(adapter1.getInstanceName(), Collections.singletonList(task1)); assignmentsByInstance.put(adapter2.getInstanceName(), Arrays.asList(task2, task3)); adapter1.updateAllAssignments(assignmentsByInstance); List<String> assignment1 = zkClient.getChildren(KeyBuilder.instanceAssignments(testCluster, adapter1.getInstanceName())); Assert.assertTrue(PollUtils.poll((assn) -> assn.size() == 1, 100, 5000, assignment1)); Assert.assertEquals(assignment1.get(0), "task1"); validateConnectorTask(testCluster, connectorType, "task1", zkClient); List<String> assignment2 = zkClient.getChildren(KeyBuilder.instanceAssignments(testCluster, adapter2.getInstanceName())); Assert.assertTrue(PollUtils.poll((assn) -> assn.size() == 2, 100, 5000, assignment2)); Assert.assertEquals(assignment2.get(0), "task2"); Assert.assertEquals(assignment2.get(1), "task3"); zkClient.close(); }