Source:
private boolean assignTasksToClients(final Set<String> allSourceTopics, final Map<TaskId, Set<TopicPartition>> partitionsForTask, final Map<Integer, TopicsInfo> topicGroups, final Map<UUID, ClientMetadata> clientMetadataMap, final Cluster fullMetadata) { final Map<TopicPartition, TaskId> taskForPartition = new HashMap<>(); final Map<Integer, Set<TaskId>> tasksForTopicGroup = new HashMap<>(); populateTasksForMaps(taskForPartition, tasksForTopicGroup, allSourceTopics, partitionsForTask, fullMetadata); final Map<TaskId, Set<TopicPartition>> changelogsByStatefulTask = prepareChangelogTopics(topicGroups, tasksForTopicGroup); final Map<UUID, ClientState> clientStates = new HashMap<>(); final boolean lagComputationSuccessful = populateClientStatesMap(clientStates, clientMetadataMap, taskForPartition, changelogsByStatefulTask); final Set<TaskId> allTasks = partitionsForTask.keySet(); final Set<TaskId> statefulTasks = changelogsByStatefulTask.keySet(); log.debug("Assigning tasks to clients with number of replicas {}", allTasks, clientStates, numStandbyReplicas()); final TaskAssignor taskAssignor; if (!lagComputationSuccessful) { log.info("Failed to fetch end offsets for changelogs, will return previous assignment to clients and " + "trigger another rebalance to retry."); setAssignmentErrorCode(AssignorError.REBALANCE_NEEDED.code()); taskAssignor = new PriorTaskAssignor(); } else { taskAssignor = this.taskAssignor.get(); } final boolean followupRebalanceNeeded = taskAssignor.assign(clientStates, allTasks, statefulTasks, assignmentConfigs); log.info("Assigned tasks to clients as {}{}.", Utils.NL, clientStates.entrySet().stream().map(Map.Entry::toString).collect(Collectors.joining(Utils.NL))); return followupRebalanceNeeded; }

Target:
private boolean assignTasksToClients(final Set<String> allSourceTopics, final Map<TaskId, Set<TopicPartition>> partitionsForTask, final Map<Integer, TopicsInfo> topicGroups, final Map<UUID, ClientMetadata> clientMetadataMap, final Cluster fullMetadata) { final Map<TopicPartition, TaskId> taskForPartition = new HashMap<>(); final Map<Integer, Set<TaskId>> tasksForTopicGroup = new HashMap<>(); populateTasksForMaps(taskForPartition, tasksForTopicGroup, allSourceTopics, partitionsForTask, fullMetadata); final Map<TaskId, Set<TopicPartition>> changelogsByStatefulTask = prepareChangelogTopics(topicGroups, tasksForTopicGroup); final Map<UUID, ClientState> clientStates = new HashMap<>(); final boolean lagComputationSuccessful = populateClientStatesMap(clientStates, clientMetadataMap, taskForPartition, changelogsByStatefulTask); final Set<TaskId> allTasks = partitionsForTask.keySet(); final Set<TaskId> statefulTasks = changelogsByStatefulTask.keySet(); log.debug("Assigning tasks to clients with number of replicas {}", allTasks, clientStates, numStandbyReplicas()); final TaskAssignor taskAssignor = createTaskAssignor(lagComputationSuccessful); final boolean probingRebalanceNeeded = taskAssignor.assign(clientStates, allTasks, statefulTasks, assignmentConfigs); log.info("Assigned tasks to clients as {}{}.", Utils.NL, clientStates.entrySet().stream().map(Map.Entry::toString).collect(Collectors.joining(Utils.NL))); return probingRebalanceNeeded; }