Source:
<code>private void 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; <START> if (lagComputationSuccessful && highAvailabilityEnabled) { taskAssignor = new HighAvailabilityTaskAssignor<>(clientStates, allTasks, statefulTasks, assignmentConfigs); } else { taskAssignor = new StickyTaskAssignor<>(clientStates, allTasks, statefulTasks, assignmentConfigs); if (highAvailabilityEnabled) { ((StickyTaskAssignor) taskAssignor).preservePreviousTaskAssignment(); } } taskAssignor.assign(); log.info("Assigned tasks to clients as {}{}.", Utils.NL, clientStates.entrySet().stream().map(Map.Entry::toString).collect(Collectors.joining(Utils.NL))); <END> }</code><technical_language>suggestion if (highAvailabilityEnabled) { taskAssignor = new HighAvailabilityTaskAssignor<>(clientStates, allTasks, statefulTasks, assignmentConfigs); } else { taskAssignor = new StickyTaskAssignor<>(clientStates, allTasks, statefulTasks, assignmentConfigs); if (!lagComputationSuccessful) { ((StickyTaskAssignor) taskAssignor).preservePreviousTaskAssignment(); } }</technical_language>

Target:
private void 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<UUID> taskAssignor; if (highAvailabilityEnabled) { if (lagComputationSuccessful) { taskAssignor = new HighAvailabilityTaskAssignor<>( clientStates, allTasks, statefulTasks, assignmentConfigs); } else { taskAssignor = new StickyTaskAssignor<>(clientStates, allTasks, statefulTasks, assignmentConfigs, true); } } else { taskAssignor = new StickyTaskAssignor<>(clientStates, allTasks, statefulTasks, assignmentConfigs, false); } taskAssignor.assign(); log.info("Assigned tasks to clients as {}{}.", Utils.NL, clientStates.entrySet().stream().map(Map.Entry::toString).collect(Collectors.joining(Utils.NL))); }