diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/assignment/KafkaStreamsAssignment.java b/streams/src/main/java/org/apache/kafka/streams/processor/assignment/KafkaStreamsAssignment.java index a4947c36467ff..f5205c8422bdf 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/assignment/KafkaStreamsAssignment.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/assignment/KafkaStreamsAssignment.java @@ -16,11 +16,13 @@ */ package org.apache.kafka.streams.processor.assignment; +import static java.util.Collections.unmodifiableMap; + import java.time.Instant; -import java.util.HashSet; import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import java.util.stream.Collectors; import org.apache.kafka.streams.processor.TaskId; @@ -31,7 +33,7 @@ public class KafkaStreamsAssignment { private final ProcessId processId; - private final Map assignment; + private final Map tasks; private final Optional followupRebalanceDeadline; /** @@ -45,7 +47,8 @@ public class KafkaStreamsAssignment { * @return a new KafkaStreamsAssignment object with the given processId and assignment */ public static KafkaStreamsAssignment of(final ProcessId processId, final Set assignment) { - return new KafkaStreamsAssignment(processId, assignment, Optional.empty()); + final Map tasks = assignment.stream().collect(Collectors.toMap(AssignedTask::id, Function.identity())); + return new KafkaStreamsAssignment(processId, tasks, Optional.empty()); } /** @@ -62,14 +65,14 @@ public static KafkaStreamsAssignment of(final ProcessId processId, final Set assignment, + final Map tasks, final Optional followupRebalanceDeadline) { this.processId = processId; - this.assignment = assignment.stream().collect(Collectors.toMap(AssignedTask::id, t -> t)); + this.tasks = tasks; this.followupRebalanceDeadline = followupRebalanceDeadline; } @@ -83,24 +86,18 @@ public ProcessId processId() { /** * - * @return a set of assigned tasks that are part of this {@code KafkaStreamsAssignment} + * @return a read-only set of assigned tasks that are part of this {@code KafkaStreamsAssignment} */ - public Set assignment() { - // TODO change assignment to return a map so we aren't forced to copy this into a Set - return new HashSet<>(assignment.values()); - } - - // TODO: merge this with #assignment by having it return a Map - public Set assignedTaskIds() { - return assignment.keySet(); + public Map tasks() { + return unmodifiableMap(tasks); } public void assignTask(final AssignedTask newTask) { - assignment.put(newTask.id(), newTask); + tasks.put(newTask.id(), newTask); } public void removeTask(final AssignedTask removedTask) { - assignment.remove(removedTask.id()); + tasks.remove(removedTask.id()); } /** @@ -140,5 +137,25 @@ public TaskId id() { public Type type() { return taskType; } + + @Override + public int hashCode() { + final int prime = 31; + int result = prime + this.id.hashCode(); + result = prime * result + this.type().hashCode(); + return result; + } + + @Override + public boolean equals(final Object obj) { + if (this == obj) + return true; + if (obj == null) + return false; + if (getClass() != obj.getClass()) + return false; + final AssignedTask other = (AssignedTask) obj; + return this.id.equals(other.id()) && this.taskType == other.taskType; + } } } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/assignment/TaskAssignmentUtils.java b/streams/src/main/java/org/apache/kafka/streams/processor/assignment/TaskAssignmentUtils.java index d7179226a9b97..d2e30d538050b 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/assignment/TaskAssignmentUtils.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/assignment/TaskAssignmentUtils.java @@ -16,6 +16,8 @@ */ package org.apache.kafka.streams.processor.assignment; +import static org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.STANDBY_OPTIMIZER_MAX_ITERATION; + import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; @@ -26,9 +28,15 @@ import java.util.Set; import java.util.SortedSet; import java.util.UUID; +import java.util.function.BiConsumer; +import java.util.function.BiFunction; +import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.kafka.streams.KeyValue; import org.apache.kafka.streams.processor.TaskId; import org.apache.kafka.streams.processor.assignment.KafkaStreamsAssignment.AssignedTask; +import org.apache.kafka.streams.processor.internals.assignment.ConstrainedPrioritySet; import org.apache.kafka.streams.processor.internals.assignment.Graph; import org.apache.kafka.streams.processor.internals.assignment.MinTrafficGraphConstructor; import org.apache.kafka.streams.processor.internals.assignment.RackAwareGraphConstructor; @@ -72,6 +80,27 @@ public static Map identityAssignment(final Ap return assignments; } + /** + * Assign standby tasks to KafkaStreams clients according to the default logic. + *

+ * If rack-aware client tags are configured, the rack-aware standby task assignor will be used + * + * @param applicationState the metadata and other info describing the current application state + * @param kafkaStreamsAssignments the current assignment of tasks to KafkaStreams clients + * + * @return a new map containing the mappings from KafkaStreamsAssignments updated with the default standby assignment + */ + public static Map defaultStandbyTaskAssignment(final ApplicationState applicationState, + final Map kafkaStreamsAssignments) { + if (!applicationState.assignmentConfigs().rackAwareAssignmentTags().isEmpty()) { + return tagBasedStandbyTaskAssignment(applicationState, kafkaStreamsAssignments); + } else if (canPerformRackAwareOptimization(applicationState, AssignedTask.Type.STANDBY)) { + return tagBasedStandbyTaskAssignment(applicationState, kafkaStreamsAssignments); + } else { + return loadBasedStandbyTaskAssignment(applicationState, kafkaStreamsAssignments); + } + } + /** * Optimize active task assignment for rack awareness. This optimization is based on the * {@link StreamsConfig#RACK_AWARE_ASSIGNMENT_TRAFFIC_COST_CONFIG trafficCost} @@ -98,11 +127,9 @@ public static Map identityAssignment(final Ap * * @return a map with the KafkaStreamsAssignments updated to minimize cross-rack traffic for active tasks */ - public static Map optimizeRackAwareActiveTasks( - final ApplicationState applicationState, - final Map kafkaStreamsAssignments, - final SortedSet tasks - ) { + public static Map optimizeRackAwareActiveTasks(final ApplicationState applicationState, + final Map kafkaStreamsAssignments, + final SortedSet tasks) { if (tasks.isEmpty()) { return kafkaStreamsAssignments; } @@ -131,7 +158,7 @@ public static Map optimizeRackAwareActiveTask clientRacks.put(uuid, kafkaStreamsStates.get(entry.getKey()).rackId()); } - final long initialCost = computeInitialCost( + final long initialCost = computeTotalAssignmentCost( topicPartitionsByTaskId, taskIds, clientIds, @@ -172,7 +199,7 @@ public static Map optimizeRackAwareActiveTask assignmentGraph.clientByTask, (assignment, taskId) -> assignment.assignTask(new AssignedTask(taskId, AssignedTask.Type.ACTIVE)), (assignment, taskId) -> assignment.removeTask(new AssignedTask(taskId, AssignedTask.Type.ACTIVE)), - (assignment, taskId) -> assignment.assignedTaskIds().contains(taskId) + (assignment, taskId) -> assignment.tasks().containsKey(taskId) ); return kafkaStreamsAssignments; @@ -231,7 +258,7 @@ public static Map optimizeRackAwareStandbyTas clientRacks.put(uuid, kafkaStreamsStates.get(entry.getKey()).rackId()); } - final long initialCost = computeInitialCost( + final long initialCost = computeTotalAssignmentCost( topicPartitionsByTaskId, taskIds, clientIds, @@ -244,18 +271,112 @@ public static Map optimizeRackAwareStandbyTas ); LOG.info("Assignment before standby task optimization has cost {}", initialCost); - throw new UnsupportedOperationException("Not yet Implemented."); + final MoveStandbyTaskPredicate moveablePredicate = getStandbyTaskMovePredicate(applicationState); + final BiFunction> getMovableTasks = (source, destination) -> { + return source.tasks().values().stream() + .filter(task -> task.type() == AssignedTask.Type.STANDBY) + .filter(task -> !destination.tasks().containsKey(task.id())) + .filter(task -> { + final KafkaStreamsState sourceState = kafkaStreamsStates.get(source.processId()); + final KafkaStreamsState destinationState = kafkaStreamsStates.get(source.processId()); + return moveablePredicate.canMoveStandbyTask(sourceState, destinationState, task.id(), kafkaStreamsAssignments); + }) + .map(AssignedTask::id) + .sorted() + .collect(Collectors.toList()); + }; + + final long startTime = System.currentTimeMillis(); + boolean taskMoved = true; + int round = 0; + final RackAwareGraphConstructor graphConstructor = RackAwareGraphConstructorFactory.create( + applicationState.assignmentConfigs().rackAwareAssignmentStrategy(), taskIds); + while (taskMoved && round < STANDBY_OPTIMIZER_MAX_ITERATION) { + taskMoved = false; + round++; + for (int i = 0; i < kafkaStreamsAssignments.size(); i++) { + final UUID clientId1 = clientIds.get(i); + final KafkaStreamsAssignment clientState1 = kafkaStreamsAssignments.get(new ProcessId(clientId1)); + for (int j = i + 1; j < kafkaStreamsAssignments.size(); j++) { + final UUID clientId2 = clientIds.get(i); + final KafkaStreamsAssignment clientState2 = kafkaStreamsAssignments.get(new ProcessId(clientId2)); + + final String rack1 = clientRacks.get(clientState1.processId().id()).get(); + final String rack2 = clientRacks.get(clientState2.processId().id()).get(); + // Cross rack traffic can not be reduced if racks are the same + if (rack1.equals(rack2)) { + continue; + } + + final List movable1 = getMovableTasks.apply(clientState1, clientState2); + final List movable2 = getMovableTasks.apply(clientState2, clientState1); + + // There's no needed to optimize if one is empty because the optimization + // can only swap tasks to keep the client's load balanced + if (movable1.isEmpty() || movable2.isEmpty()) { + continue; + } + + final List taskIdList = Stream.concat(movable1.stream(), movable2.stream()) + .sorted() + .collect(Collectors.toList()); + final List clients = Stream.of(clientId1, clientId2).sorted().collect(Collectors.toList()); + + final AssignmentGraph assignmentGraph = buildTaskGraph( + assignmentsByUuid, + clientRacks, + taskIdList, + clients, + topicPartitionsByTaskId, + crossRackTrafficCost, + nonOverlapCost, + false, + false, + graphConstructor + ); + assignmentGraph.graph.solveMinCostFlow(); + + taskMoved |= graphConstructor.assignTaskFromMinCostFlow( + assignmentGraph.graph, + clientIds, + taskIds, + assignmentsByUuid, + assignmentGraph.taskCountByClient, + assignmentGraph.clientByTask, + (assignment, taskId) -> assignment.assignTask(new AssignedTask(taskId, AssignedTask.Type.STANDBY)), + (assignment, taskId) -> assignment.removeTask(new AssignedTask(taskId, AssignedTask.Type.STANDBY)), + (assignment, taskId) -> assignment.tasks().containsKey(taskId) && assignment.tasks().get(taskId).type() == AssignedTask.Type.STANDBY + ); + } + } + } + final long finalCost = computeTotalAssignmentCost( + topicPartitionsByTaskId, + taskIds, + clientIds, + assignmentsByUuid, + clientRacks, + crossRackTrafficCost, + nonOverlapCost, + true, + true + ); + + final long duration = System.currentTimeMillis() - startTime; + LOG.info("Assignment after {} rounds and {} milliseconds for standby task optimization is {}\n with cost {}", + round, duration, kafkaStreamsAssignments, finalCost); + return kafkaStreamsAssignments; } - private static long computeInitialCost(final Map> topicPartitionsByTaskId, - final List taskIds, - final List clientIds, - final Map assignmentsByUuid, - final Map> clientRacks, - final int crossRackTrafficCost, - final int nonOverlapCost, - final boolean hasReplica, - final boolean isStandby) { + private static long computeTotalAssignmentCost(final Map> topicPartitionsByTaskId, + final List taskIds, + final List clientIds, + final Map assignmentsByUuid, + final Map> clientRacks, + final int crossRackTrafficCost, + final int nonOverlapCost, + final boolean hasReplica, + final boolean isStandby) { if (taskIds.isEmpty()) { return 0; } @@ -296,7 +417,7 @@ private static AssignmentGraph buildTaskGraph(final Map assignment.assignedTaskIds().contains(taskId), + (assignment, taskId) -> assignment.tasks().containsKey(taskId), (taskId, processId, inCurrentAssignment, unused0, unused1, unused2) -> { final String clientRack = clientRacks.get(processId).get(); final int assignmentChangeCost = !inCurrentAssignment ? nonOverlapCost : 0; @@ -328,6 +449,14 @@ public AssignmentGraph(final Graph graph, } } + @FunctionalInterface + public interface MoveStandbyTaskPredicate { + boolean canMoveStandbyTask(final KafkaStreamsState source, + final KafkaStreamsState destination, + final TaskId taskId, + final Map kafkaStreamsAssignment); + } + /** * * @return the traffic cost of assigning this {@param task} to the client {@param streamsState}. @@ -354,7 +483,6 @@ private static boolean canPerformRackAwareOptimization(final ApplicationState ap final AssignedTask.Type taskType) { final String rackAwareAssignmentStrategy = applicationState.assignmentConfigs().rackAwareAssignmentStrategy(); if (StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE.equals(rackAwareAssignmentStrategy)) { - LOG.warn("KafkaStreams rack aware task assignment optimization was disabled in the StreamsConfig."); return false; } return hasValidRackInformation(applicationState, taskType); @@ -407,4 +535,317 @@ private static boolean hasValidRackInformation(final TaskInfo task, } return true; } + + private static Map tagBasedStandbyTaskAssignment(final ApplicationState applicationState, + final Map kafkaStreamsAssignments) { + final int numStandbyReplicas = applicationState.assignmentConfigs().numStandbyReplicas(); + final Map streamStates = applicationState.kafkaStreamsStates(false); + + final Set rackAwareAssignmentTags = new HashSet<>(applicationState.assignmentConfigs().rackAwareAssignmentTags()); + final TagStatistics tagStatistics = new TagStatistics(applicationState); + + final ConstrainedPrioritySet standbyTaskClientsByTaskLoad = standbyTaskPriorityListByLoad(streamStates, kafkaStreamsAssignments); + + final Set statefulTaskIds = applicationState.allTasks().values().stream() + .filter(TaskInfo::isStateful) + .map(TaskInfo::id) + .collect(Collectors.toSet()); + final Map tasksToRemainingStandbys = statefulTaskIds.stream() + .collect(Collectors.toMap(Function.identity(), t -> numStandbyReplicas)); + final Map clientsByUuid = kafkaStreamsAssignments.entrySet().stream().collect(Collectors.toMap( + entry -> entry.getKey().id(), + Map.Entry::getValue + )); + + final Map pendingStandbyTasksToClientId = new HashMap<>(); + for (final TaskId statefulTaskId : statefulTaskIds) { + for (final KafkaStreamsAssignment assignment : clientsByUuid.values()) { + if (assignment.tasks().containsKey(statefulTaskId)) { + assignStandbyTasksToClientsWithDifferentTags( + numStandbyReplicas, + standbyTaskClientsByTaskLoad, + statefulTaskId, + assignment.processId(), + rackAwareAssignmentTags, + streamStates, + kafkaStreamsAssignments, + tasksToRemainingStandbys, + tagStatistics.tagKeyToValues, + tagStatistics.tagEntryToClients, + pendingStandbyTasksToClientId + ); + } + } + } + + if (!tasksToRemainingStandbys.isEmpty()) { + assignPendingStandbyTasksToLeastLoadedClients(clientsByUuid, + numStandbyReplicas, + standbyTaskClientsByTaskLoad, + tasksToRemainingStandbys); + } + + return kafkaStreamsAssignments; + } + + private static Map loadBasedStandbyTaskAssignment(final ApplicationState applicationState, + final Map kafkaStreamsAssignments) { + final int numStandbyReplicas = applicationState.assignmentConfigs().numStandbyReplicas(); + final Map streamStates = applicationState.kafkaStreamsStates(false); + + final Set statefulTaskIds = applicationState.allTasks().values().stream() + .filter(TaskInfo::isStateful) + .map(TaskInfo::id) + .collect(Collectors.toSet()); + final Map tasksToRemainingStandbys = statefulTaskIds.stream() + .collect(Collectors.toMap(Function.identity(), t -> numStandbyReplicas)); + final Map clients = kafkaStreamsAssignments.entrySet().stream().collect(Collectors.toMap( + entry -> entry.getKey().id(), + Map.Entry::getValue + )); + + final ConstrainedPrioritySet standbyTaskClientsByTaskLoad = standbyTaskPriorityListByLoad(streamStates, kafkaStreamsAssignments); + standbyTaskClientsByTaskLoad.offerAll(streamStates.keySet().stream().map(ProcessId::id).collect(Collectors.toSet())); + for (final TaskId task : statefulTaskIds) { + assignStandbyTasksForActiveTask( + numStandbyReplicas, + clients, + tasksToRemainingStandbys, + standbyTaskClientsByTaskLoad, + task + ); + } + return kafkaStreamsAssignments; + } + + private static void assignStandbyTasksForActiveTask(final int numStandbyReplicas, + final Map clients, + final Map tasksToRemainingStandbys, + final ConstrainedPrioritySet standbyTaskClientsByTaskLoad, + final TaskId activeTaskId) { + int numRemainingStandbys = tasksToRemainingStandbys.get(activeTaskId); + while (numRemainingStandbys > 0) { + final UUID client = standbyTaskClientsByTaskLoad.poll(activeTaskId); + if (client == null) { + break; + } + clients.get(client).assignTask(new AssignedTask(activeTaskId, AssignedTask.Type.STANDBY)); + numRemainingStandbys--; + standbyTaskClientsByTaskLoad.offer(client); + } + + tasksToRemainingStandbys.put(activeTaskId, numRemainingStandbys); + if (numRemainingStandbys > 0) { + LOG.warn("Unable to assign {} of {} standby tasks for task [{}]. " + + "There is not enough available capacity. You should " + + "increase the number of application instances " + + "to maintain the requested number of standby replicas.", + numRemainingStandbys, numStandbyReplicas, activeTaskId); + } + } + + private static void assignStandbyTasksToClientsWithDifferentTags(final int numberOfStandbyClients, + final ConstrainedPrioritySet standbyTaskClientsByTaskLoad, + final TaskId activeTaskId, + final ProcessId activeTaskClient, + final Set rackAwareAssignmentTags, + final Map clientStates, + final Map kafkaStreamsAssignments, + final Map tasksToRemainingStandbys, + final Map> tagKeyToValues, + final Map, Set> tagEntryToClients, + final Map pendingStandbyTasksToClientId) { + standbyTaskClientsByTaskLoad.offerAll(clientStates.keySet().stream() + .map(ProcessId::id).collect(Collectors.toSet())); + + // We set countOfUsedClients as 1 because client where active task is located has to be considered as used. + int countOfUsedClients = 1; + int numRemainingStandbys = tasksToRemainingStandbys.get(activeTaskId); + + final Map, Set> tagEntryToUsedClients = new HashMap<>(); + + ProcessId lastUsedClient = activeTaskClient; + do { + updateClientsOnAlreadyUsedTagEntries( + clientStates.get(lastUsedClient), + countOfUsedClients, + rackAwareAssignmentTags, + tagEntryToClients, + tagKeyToValues, + tagEntryToUsedClients + ); + + final UUID clientOnUnusedTagDimensions = standbyTaskClientsByTaskLoad.poll( + activeTaskId, uuid -> !isClientUsedOnAnyOfTheTagEntries(new ProcessId(uuid), tagEntryToUsedClients) + ); + + if (clientOnUnusedTagDimensions == null) { + break; + } + + final KafkaStreamsState clientStateOnUsedTagDimensions = clientStates.get(new ProcessId(clientOnUnusedTagDimensions)); + countOfUsedClients++; + numRemainingStandbys--; + + LOG.debug("Assigning {} out of {} standby tasks for an active task [{}] with client tags {}. " + + "Standby task client tags are {}.", + numberOfStandbyClients - numRemainingStandbys, numberOfStandbyClients, activeTaskId, + clientStates.get(activeTaskClient).clientTags(), + clientStateOnUsedTagDimensions.clientTags()); + + kafkaStreamsAssignments.get(clientStateOnUsedTagDimensions.processId()).assignTask( + new AssignedTask(activeTaskId, AssignedTask.Type.STANDBY) + ); + lastUsedClient = new ProcessId(clientOnUnusedTagDimensions); + } while (numRemainingStandbys > 0); + + if (numRemainingStandbys > 0) { + pendingStandbyTasksToClientId.put(activeTaskId, activeTaskClient); + tasksToRemainingStandbys.put(activeTaskId, numRemainingStandbys); + LOG.warn("Rack aware standby task assignment was not able to assign {} of {} standby tasks for the " + + "active task [{}] with the rack aware assignment tags {}. " + + "This may happen when there aren't enough application instances on different tag " + + "dimensions compared to an active and corresponding standby task. " + + "Consider launching application instances on different tag dimensions than [{}]. " + + "Standby task assignment will fall back to assigning standby tasks to the least loaded clients.", + numRemainingStandbys, numberOfStandbyClients, + activeTaskId, rackAwareAssignmentTags, + clientStates.get(activeTaskClient).clientTags()); + + } else { + tasksToRemainingStandbys.remove(activeTaskId); + } + } + + private static boolean isClientUsedOnAnyOfTheTagEntries(final ProcessId client, + final Map, Set> tagEntryToUsedClients) { + return tagEntryToUsedClients.values().stream().anyMatch(usedClients -> usedClients.contains(client)); + } + + private static void updateClientsOnAlreadyUsedTagEntries(final KafkaStreamsState usedClient, + final int countOfUsedClients, + final Set rackAwareAssignmentTags, + final Map, Set> tagEntryToClients, + final Map> tagKeyToValues, + final Map, Set> tagEntryToUsedClients) { + final Map usedClientTags = usedClient.clientTags(); + + for (final Map.Entry usedClientTagEntry : usedClientTags.entrySet()) { + final String tagKey = usedClientTagEntry.getKey(); + + if (!rackAwareAssignmentTags.contains(tagKey)) { + LOG.warn("Client tag with key [{}] will be ignored when computing rack aware standby " + + "task assignment because it is not part of the configured rack awareness [{}].", + tagKey, rackAwareAssignmentTags); + continue; + } + + final Set allTagValues = tagKeyToValues.get(tagKey); + + if (allTagValues.size() <= countOfUsedClients) { + allTagValues.forEach(tagValue -> tagEntryToUsedClients.remove(new KeyValue<>(tagKey, tagValue))); + } else { + final String tagValue = usedClientTagEntry.getValue(); + final KeyValue tagEntry = new KeyValue<>(tagKey, tagValue); + final Set clientsOnUsedTagValue = tagEntryToClients.get(tagEntry); + tagEntryToUsedClients.put(tagEntry, clientsOnUsedTagValue); + } + } + } + + private static MoveStandbyTaskPredicate getStandbyTaskMovePredicate(final ApplicationState applicationState) { + final boolean hasRackAwareAssignmentTags = !applicationState.assignmentConfigs().rackAwareAssignmentTags().isEmpty(); + if (hasRackAwareAssignmentTags) { + final BiConsumer>> addTags = (cs, tagSet) -> { + final Map tags = cs.clientTags(); + if (tags != null) { + tagSet.addAll(tags.entrySet().stream() + .map(entry -> KeyValue.pair(entry.getKey(), entry.getValue())) + .collect(Collectors.toList()) + ); + } + }; + + final Map clients = applicationState.kafkaStreamsStates(false); + + return (source, destination, sourceTask, kafkaStreamsAssignments) -> { + final Set> tagsWithSource = new HashSet<>(); + final Set> tagsWithDestination = new HashSet<>(); + for (final KafkaStreamsAssignment assignment: kafkaStreamsAssignments.values()) { + final boolean hasAssignedTask = assignment.tasks().containsKey(sourceTask); + final boolean isSourceProcess = assignment.processId().equals(source.processId()); + final boolean isDestinationProcess = assignment.processId().equals(destination.processId()); + if (hasAssignedTask && !isSourceProcess && !isDestinationProcess) { + final KafkaStreamsState clientState = clients.get(assignment.processId()); + addTags.accept(clientState, tagsWithSource); + addTags.accept(clientState, tagsWithDestination); + } + } + addTags.accept(source, tagsWithSource); + addTags.accept(destination, tagsWithDestination); + return tagsWithDestination.size() >= tagsWithSource.size(); + }; + } else { + return (a, b, c, d) -> true; + } + } + + private static ConstrainedPrioritySet standbyTaskPriorityListByLoad(final Map streamStates, + final Map kafkaStreamsAssignments) { + return new ConstrainedPrioritySet( + (processId, taskId) -> kafkaStreamsAssignments.get(new ProcessId(processId)).tasks().containsKey(taskId), + processId -> { + final double capacity = streamStates.get(new ProcessId(processId)).numProcessingThreads(); + final double numTasks = kafkaStreamsAssignments.get(new ProcessId(processId)).tasks().size(); + return numTasks / capacity; + } + ); + } + + private static void assignPendingStandbyTasksToLeastLoadedClients(final Map clients, + final int numStandbyReplicas, + final ConstrainedPrioritySet standbyTaskClientsByTaskLoad, + final Map pendingStandbyTaskToNumberRemainingStandbys) { + // We need to re offer all the clients to find the least loaded ones + standbyTaskClientsByTaskLoad.offerAll(clients.keySet()); + + for (final Map.Entry pendingStandbyTaskAssignmentEntry : pendingStandbyTaskToNumberRemainingStandbys.entrySet()) { + final TaskId activeTaskId = pendingStandbyTaskAssignmentEntry.getKey(); + + assignStandbyTasksForActiveTask( + numStandbyReplicas, + clients, + pendingStandbyTaskToNumberRemainingStandbys, + standbyTaskClientsByTaskLoad, + activeTaskId + ); + } + } + + private static class TagStatistics { + private final Map> tagKeyToValues; + private final Map, Set> tagEntryToClients; + + private TagStatistics(final Map> tagKeyToValues, + final Map, Set> tagEntryToClients) { + this.tagKeyToValues = tagKeyToValues; + this.tagEntryToClients = tagEntryToClients; + } + + public TagStatistics(final ApplicationState applicationState) { + final Map clientStates = applicationState.kafkaStreamsStates(false); + + final Map> tagKeyToValues = new HashMap<>(); + final Map, Set> tagEntryToClients = new HashMap<>(); + for (final KafkaStreamsState state : clientStates.values()) { + state.clientTags().forEach((tagKey, tagValue) -> { + tagKeyToValues.computeIfAbsent(tagKey, ignored -> new HashSet<>()).add(tagValue); + tagEntryToClients.computeIfAbsent(new KeyValue<>(tagKey, tagValue), ignored -> new HashSet<>()).add(state.processId()); + }); + } + + this.tagKeyToValues = tagKeyToValues; + this.tagEntryToClients = tagEntryToClients; + } + } } \ No newline at end of file diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/assignment/assignors/StickyTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/assignment/assignors/StickyTaskAssignor.java index 74d272ef1ecf5..2174b6d823b6f 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/assignment/assignors/StickyTaskAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/assignment/assignors/StickyTaskAssignor.java @@ -278,7 +278,7 @@ public void processOptimizedAssignments(final Map entry : optimizedAssignments.entrySet()) { final ProcessId processId = entry.getKey(); - final Set assignedTasks = optimizedAssignments.get(processId).assignment(); + final Set assignedTasks = new HashSet<>(optimizedAssignments.get(processId).tasks().values()); newAssignments.put(processId, assignedTasks); for (final AssignedTask task : assignedTasks) { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 3fd54d70ab892..658ba1540da48 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -18,6 +18,7 @@ import java.time.Instant; import java.util.Optional; +import java.util.function.Function; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.ListOffsetsResult; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; @@ -555,18 +556,21 @@ private ApplicationState buildApplicationState(final TopologyMetadata topologyMe RackUtils.annotateTopicPartitionsWithRackInfo(cluster, internalTopicManager, allTopicPartitions); - final Set logicalTasks = logicalTaskIds.stream().map(taskId -> { - final Set stateStoreNames = topologyMetadata - .stateStoreNameToSourceTopicsForTopology(taskId.topologyName()) - .keySet(); - final Set topicPartitions = topicPartitionsForTask.get(taskId); - return new DefaultTaskInfo( - taskId, - !stateStoreNames.isEmpty(), - stateStoreNames, - topicPartitions - ); - }).collect(Collectors.toSet()); + final Map logicalTasks = logicalTaskIds.stream().collect(Collectors.toMap( + Function.identity(), + taskId -> { + final Set stateStoreNames = topologyMetadata + .stateStoreNameToSourceTopicsForTopology(taskId.topologyName()) + .keySet(); + final Set topicPartitions = topicPartitionsForTask.get(taskId); + return new DefaultTaskInfo( + taskId, + !stateStoreNames.isEmpty(), + stateStoreNames, + topicPartitions + ); + } + )); return new DefaultApplicationState( assignmentConfigs.toPublicAssignmentConfigs(), @@ -728,12 +732,12 @@ private void checkAllPartitions(final Set allSourceTopics, * populate the stateful tasks that have been assigned to the clients */ private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetadata, - final Set allSourceTopics, - final Map topicGroups, - final Map clientMetadataMap, - final Map> partitionsForTask, - final Map>> racksForProcessConsumer, - final Set statefulTasks) { + final Set allSourceTopics, + final Map topicGroups, + final Map clientMetadataMap, + final Map> partitionsForTask, + final Map>> racksForProcessConsumer, + final Set statefulTasks) { if (!statefulTasks.isEmpty()) { throw new TaskAssignmentException("The stateful tasks should not be populated before assigning tasks to clients"); } @@ -775,7 +779,7 @@ private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetada final Optional userTaskAssignor = userTaskAssignorSupplier.get(); - UserTaskAssignmentListener userTaskAssignmentListener = (GroupAssignment assignment, GroupSubscription subscription) -> { }; + final UserTaskAssignmentListener userTaskAssignmentListener; if (userTaskAssignor.isPresent()) { final ApplicationState applicationState = buildApplicationState( taskManager.topologyMetadata(), @@ -785,12 +789,11 @@ private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetada ); final org.apache.kafka.streams.processor.assignment.TaskAssignor assignor = userTaskAssignor.get(); final TaskAssignment taskAssignment = assignor.assign(applicationState); - processStreamsPartitionAssignment(clientMetadataMap, taskAssignment); final AssignmentError assignmentError = validateTaskAssignment(applicationState, taskAssignment); - userTaskAssignmentListener = (GroupAssignment assignment, GroupSubscription subscription) -> { - assignor.onAssignmentComputed(assignment, subscription, assignmentError); - }; + processStreamsPartitionAssignment(clientMetadataMap, taskAssignment); + userTaskAssignmentListener = (assignment, subscription) -> assignor.onAssignmentComputed(assignment, subscription, assignmentError); } else { + userTaskAssignmentListener = (assignment, subscription) -> { }; final TaskAssignor taskAssignor = createTaskAssignor(lagComputationSuccessful); final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor( fullMetadata, @@ -1564,7 +1567,7 @@ private AssignmentError validateTaskAssignment(final ApplicationState applicatio final Map standbyTasksInOutput = new HashMap<>(); for (final KafkaStreamsAssignment assignment : assignments) { final Set tasksForAssignment = new HashSet<>(); - for (final KafkaStreamsAssignment.AssignedTask task : assignment.assignment()) { + for (final KafkaStreamsAssignment.AssignedTask task : assignment.tasks().values()) { if (activeTasksInOutput.containsKey(task.id()) && task.type() == KafkaStreamsAssignment.AssignedTask.Type.ACTIVE) { log.error("Assignment is invalid: active task {} was assigned to multiple KafkaStreams clients: {} and {}", task.id(), assignment.processId().id(), activeTasksInOutput.get(task.id()).id()); @@ -1614,7 +1617,7 @@ private AssignmentError validateTaskAssignment(final ApplicationState applicatio final Set taskIdsInInput = applicationState.allTasks().keySet(); for (final KafkaStreamsAssignment assignment : assignments) { - for (final KafkaStreamsAssignment.AssignedTask task : assignment.assignment()) { + for (final KafkaStreamsAssignment.AssignedTask task : assignment.tasks().values()) { if (!taskIdsInInput.contains(task.id())) { log.error("Assignment is invalid: task {} assigned to KafkaStreams client {} was unknown", task.id(), assignment.processId().id()); diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index d24e9c19167c5..704c1d50885c4 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -85,8 +85,8 @@ public ClientState(final UUID processId, final Map clientTags) { } ClientState(final UUID processId, final int capacity, final Map clientTags) { - previousStandbyTasks.taskIds(new TreeSet<>()); - previousActiveTasks.taskIds(new TreeSet<>()); + previousStandbyTasks.setTaskIds(new TreeSet<>()); + previousActiveTasks.setTaskIds(new TreeSet<>()); taskOffsetSums = new TreeMap<>(); taskLagTotals = new TreeMap<>(); this.capacity = capacity; @@ -110,8 +110,8 @@ public ClientState(final Set previousActiveTasks, final Map clientTags, final int capacity, final UUID processId) { - this.previousStandbyTasks.taskIds(unmodifiableSet(new TreeSet<>(previousStandbyTasks))); - this.previousActiveTasks.taskIds(unmodifiableSet(new TreeSet<>(previousActiveTasks))); + this.previousStandbyTasks.setTaskIds(unmodifiableSet(new TreeSet<>(previousStandbyTasks))); + this.previousActiveTasks.setTaskIds(unmodifiableSet(new TreeSet<>(previousActiveTasks))); taskOffsetSums = emptyMap(); this.taskLagTotals = unmodifiableMap(taskLagTotals); this.capacity = capacity; @@ -489,14 +489,14 @@ public SortedMap> taskIdsByPreviousConsumer() { } public void setAssignedTasks(final KafkaStreamsAssignment assignment) { - final Set activeTasks = assignment.assignment().stream() + final Set activeTasks = assignment.tasks().values().stream() .filter(task -> task.type() == ACTIVE).map(KafkaStreamsAssignment.AssignedTask::id) .collect(Collectors.toSet()); - final Set standbyTasks = assignment.assignment().stream() + final Set standbyTasks = assignment.tasks().values().stream() .filter(task -> task.type() == STANDBY).map(KafkaStreamsAssignment.AssignedTask::id) .collect(Collectors.toSet()); - assignedActiveTasks.taskIds(activeTasks); - assignedStandbyTasks.taskIds(standbyTasks); + assignedActiveTasks.setTaskIds(activeTasks); + assignedStandbyTasks.setTaskIds(standbyTasks); } public String currentAssignment() { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTask.java index 92769699ccccd..1f098dc9aea81 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTask.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTask.java @@ -31,7 +31,7 @@ class ClientStateTask { this.consumerToTaskIds = consumerToTaskIds; } - void taskIds(final Set clientToTaskIds) { + void setTaskIds(final Set clientToTaskIds) { taskIds = clientToTaskIds; } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ConstrainedPrioritySet.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ConstrainedPrioritySet.java index 1de9dfc3c67b9..3705ee39af37d 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ConstrainedPrioritySet.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ConstrainedPrioritySet.java @@ -30,14 +30,14 @@ /** * Wraps a priority queue of clients and returns the next valid candidate(s) based on the current task assignment */ -class ConstrainedPrioritySet { +public class ConstrainedPrioritySet { private final PriorityQueue clientsByTaskLoad; private final BiFunction constraint; private final Set uniqueClients = new HashSet<>(); - ConstrainedPrioritySet(final BiFunction constraint, - final Function weight) { + public ConstrainedPrioritySet(final BiFunction constraint, + final Function weight) { this.constraint = constraint; clientsByTaskLoad = new PriorityQueue<>(Comparator.comparing(weight).thenComparing(clientId -> clientId)); } @@ -45,7 +45,7 @@ class ConstrainedPrioritySet { /** * @return the next least loaded client that satisfies the given criteria, or null if none do */ - UUID poll(final TaskId task, final Function extraConstraint) { + public UUID poll(final TaskId task, final Function extraConstraint) { final Set invalidPolledClients = new HashSet<>(); while (!clientsByTaskLoad.isEmpty()) { final UUID candidateClient = pollNextClient(); @@ -66,17 +66,17 @@ UUID poll(final TaskId task, final Function extraConstraint) { /** * @return the next least loaded client that satisfies the given criteria, or null if none do */ - UUID poll(final TaskId task) { + public UUID poll(final TaskId task) { return poll(task, client -> true); } - void offerAll(final Collection clients) { + public void offerAll(final Collection clients) { for (final UUID client : clients) { offer(client); } } - void offer(final UUID client) { + public void offer(final UUID client) { if (uniqueClients.contains(client)) { clientsByTaskLoad.remove(client); } else { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/DefaultApplicationState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/DefaultApplicationState.java index f92898ccf349e..b3b3084bc1751 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/DefaultApplicationState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/DefaultApplicationState.java @@ -21,9 +21,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Optional; -import java.util.Set; import java.util.UUID; -import java.util.stream.Collectors; import org.apache.kafka.streams.processor.assignment.TaskInfo; import org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.ClientMetadata; import org.apache.kafka.streams.processor.TaskId; @@ -42,10 +40,10 @@ public class DefaultApplicationState implements ApplicationState { private final Map> cachedKafkaStreamStates; public DefaultApplicationState(final AssignmentConfigs assignmentConfigs, - final Set tasks, + final Map tasks, final Map clientStates) { this.assignmentConfigs = assignmentConfigs; - this.tasks = unmodifiableMap(tasks.stream().collect(Collectors.toMap(TaskInfo::id, task -> task))); + this.tasks = unmodifiableMap(tasks); this.clientStates = clientStates; this.cachedKafkaStreamStates = new HashMap<>(); } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java index 1d671603e4cd7..4b430fbb2c1dd 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java @@ -77,7 +77,7 @@ int getCost(final TaskId taskId, // This is number is picked based on testing. Usually the optimization for standby assignment // stops after 3 rounds - private static final int STANDBY_OPTIMIZER_MAX_ITERATION = 4; + public static final int STANDBY_OPTIMIZER_MAX_ITERATION = 4; private final Cluster fullMetadata; private final Map> partitionsForTask;