ZkAdapter { public Map<String, Set<String>> getPartitionMovement(String connectorType, String datastreamGroupName, long timestamp) { Map<String, Set<String>> result = new HashMap<>(); String path = KeyBuilder.getTargetAssignmentPath(_cluster, connectorType, datastreamGroupName); if (_zkclient.exists(path)) { List<String> nodes = _zkclient.getChildren(path); Collections.sort(nodes, Collections.reverseOrder()); Map<String, String> hostInstanceMap = getHostInstanceMap(); Set<String> processedPartitions = new HashSet<>(); for (String node : nodes) { if (Long.parseLong(node) > timestamp) { continue; } String content = _zkclient.readData(path + '/' + node); HostTargetAssignment assignment = HostTargetAssignment.fromJson(content); List<String> hostPartitions = new ArrayList<>(assignment.getPartitionNames()); hostPartitions.removeAll(processedPartitions); processedPartitions.addAll(assignment.getPartitionNames()); if (hostInstanceMap.containsKey(assignment.getTargetHost()) && assignment.getPartitionNames() != null) { LOG.info("Added assignment {} {}", assignment.getTargetHost(), hostPartitions); if (result.containsKey(hostInstanceMap.get(assignment.getTargetHost()))) { result.get(hostInstanceMap.get(assignment.getTargetHost())).addAll(hostPartitions); } else { result.put(hostInstanceMap.get(assignment.getTargetHost()), new HashSet<>(hostPartitions)); } } else { LOG.warn("assignment target host {} not found from the live instances {}", assignment.getTargetHost(), hostInstanceMap.keySet()); } } } return result; } 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 testGetPartitionMovement() throws Exception { String testCluster = "testGetPartitionMovement"; String connectorType = "connectorType"; ZkClient zkClient = new ZkClient(_zkConnectionString); ZkAdapter adapter1 = createZkAdapter(testCluster); Datastream[] datastreams = DatastreamTestUtils.createAndStoreDatastreams(zkClient, testCluster, connectorType, "datastream1"); DatastreamGroup datastreamGroup = new DatastreamGroup(Arrays.asList(datastreams)); adapter1.connect(); PollUtils.poll(adapter1::isLeader, 50, 5000); List<String> instances = adapter1.getLiveInstances(); String hostName = instances.get(0).substring(0, instances.get(0).lastIndexOf('-')); long current = System.currentTimeMillis(); String path = KeyBuilder.getTargetAssignmentPath(testCluster, connectorType, datastreamGroup.getName()); HostTargetAssignment targetAssignment = new HostTargetAssignment(ImmutableList.of("t-0", "t-1"), hostName); zkClient.ensurePath(path); if (zkClient.exists(path)) { String json = targetAssignment.toJson(); zkClient.ensurePath(path + '/' + current); zkClient.writeData(path + '/' + current, json); } long current2 = System.currentTimeMillis(); Map<String, Set<String>> newAssignment = adapter1.getPartitionMovement(datastreamGroup.getConnectorName(), datastreamGroup.getName(), current2); Assert.assertEquals(newAssignment.get(instances.get(0)), ImmutableSet.of("t-0", "t-1")); adapter1.disconnect(); zkClient.close(); }