You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by mj...@apache.org on 2023/12/06 06:00:11 UTC
(kafka) branch trunk updated: KAFKA-15022: introduce interface to control graph constructor (#14714)
This is an automated email from the ASF dual-hosted git repository.
mjsax pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new f6560ab1cd2 KAFKA-15022: introduce interface to control graph constructor (#14714)
f6560ab1cd2 is described below
commit f6560ab1cd226bbb33c6380d33bd6caddbdb3e68
Author: Hao Li <11...@users.noreply.github.com>
AuthorDate: Tue Dec 5 22:00:04 2023 -0800
KAFKA-15022: introduce interface to control graph constructor (#14714)
Part of KIP-925.
Refactor graph construction and assignment in RackAwareAssignor to new interface.
Reviewers: Matthias J. Sax <ma...@confluent.io>
---
.../processor/internals/TopologyMetadata.java | 18 +-
.../BalanceSubtopologyGraphConstructor.java | 254 +++++++++++++++++
.../processor/internals/assignment/Graph.java | 95 +++++++
.../assignment/MinTrafficGraphConstructor.java | 145 ++++++++++
.../assignment/RackAwareGraphConstructor.java | 142 ++++++++++
.../RackAwareGraphConstructorFactory.java | 39 +++
.../assignment/RackAwareTaskAssignor.java | 211 ++++-----------
.../processor/internals/TopologyMetadataTest.java | 29 +-
.../internals/assignment/AssignmentTestUtils.java | 34 ++-
.../processor/internals/assignment/GraphTest.java | 100 +++++++
.../RackAwareGraphConstructorFactoryTest.java | 46 ++++
.../assignment/RackAwareGraphConstructorTest.java | 301 +++++++++++++++++++++
.../assignment/RackAwareTaskAssignorTest.java | 85 ++++--
.../assignment/TaskAssignorConvergenceTest.java | 2 +-
14 files changed, 1314 insertions(+), 187 deletions(-)
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/TopologyMetadata.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/TopologyMetadata.java
index 6164a0e6171..150a6be9c5a 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/TopologyMetadata.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/TopologyMetadata.java
@@ -646,7 +646,7 @@ public class TopologyMetadata {
}
}
- public static class Subtopology {
+ public static class Subtopology implements Comparable<Subtopology> {
final int nodeGroupId;
final String namedTopology;
@@ -672,5 +672,21 @@ public class TopologyMetadata {
public int hashCode() {
return Objects.hash(nodeGroupId, namedTopology);
}
+
+ @Override
+ public int compareTo(final Subtopology other) {
+ if (nodeGroupId != other.nodeGroupId) {
+ return Integer.compare(nodeGroupId, other.nodeGroupId);
+ }
+ if (namedTopology == null) {
+ return other.namedTopology == null ? 0 : -1;
+ }
+ if (other.namedTopology == null) {
+ return 1;
+ }
+
+ // Both not null
+ return namedTopology.compareTo(other.namedTopology);
+ }
}
}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/BalanceSubtopologyGraphConstructor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/BalanceSubtopologyGraphConstructor.java
new file mode 100644
index 00000000000..e913955c185
--- /dev/null
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/BalanceSubtopologyGraphConstructor.java
@@ -0,0 +1,254 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor.internals.assignment;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.SortedMap;
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import java.util.UUID;
+import java.util.function.BiConsumer;
+import java.util.function.BiPredicate;
+import java.util.stream.Collectors;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.CostFunction;
+
+public class BalanceSubtopologyGraphConstructor implements RackAwareGraphConstructor {
+
+ private final Map<Subtopology, Set<TaskId>> tasksForTopicGroup;
+
+ public BalanceSubtopologyGraphConstructor(final Map<Subtopology, Set<TaskId>> tasksForTopicGroup) {
+ this.tasksForTopicGroup = tasksForTopicGroup;
+ }
+
+ @Override
+ public int getSinkNodeID(
+ final List<TaskId> taskIdList,
+ final List<UUID> clientList,
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup
+ ) {
+ return clientList.size() + taskIdList.size() + clientList.size() * tasksForTopicGroup.size();
+ }
+
+
+ @Override
+ public int getClientNodeId(final int clientIndex, final List<TaskId> taskIdList, final List<UUID> clientList, final int topicGroupIndex) {
+ return taskIdList.size() + clientList.size() * topicGroupIndex + clientIndex;
+ }
+
+ @Override
+ public int getClientIndex(final int clientNodeId, final List<TaskId> taskIdList, final List<UUID> clientList, final int topicGroupIndex) {
+ return clientNodeId - taskIdList.size() - clientList.size() * topicGroupIndex;
+ }
+
+ private static int getSecondStageClientNodeId(final List<TaskId> taskIdList, final List<UUID> clientList, final Map<Subtopology, Set<TaskId>> tasksForTopicGroup, final int clientIndex) {
+ return taskIdList.size() + clientList.size() * tasksForTopicGroup.size() + clientIndex;
+ }
+
+ @Override
+ public Graph<Integer> constructTaskGraph(
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask,
+ final CostFunction costFunction,
+ final int trafficCost,
+ final int nonOverlapCost,
+ final boolean hasReplica,
+ final boolean isStandby
+ ) {
+ validateTasks(taskIdList);
+
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : taskIdList) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (hasAssignedTask.test(clientState.getValue(), taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ constructEdges(
+ graph,
+ taskIdList,
+ clientList,
+ clientStates,
+ taskClientMap,
+ originalAssignedTaskNumber,
+ hasAssignedTask,
+ costFunction,
+ trafficCost,
+ nonOverlapCost,
+ hasReplica,
+ isStandby
+ );
+
+ // Run max flow algorithm to get a solution first
+ final long maxFlow = graph.calculateMaxFlow();
+ if (maxFlow != taskIdList.size()) {
+ throw new IllegalStateException("max flow calculated: " + maxFlow + " doesn't match taskSize: " + taskIdList.size());
+ }
+
+ return graph;
+ }
+
+ @Override
+ public boolean assignTaskFromMinCostFlow(
+ final Graph<Integer> graph,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final Map<TaskId, UUID> taskClientMap,
+ final BiConsumer<ClientState, TaskId> assignTask,
+ final BiConsumer<ClientState, TaskId> unAssignTask,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask
+ ) {
+ final SortedMap<Subtopology, Set<TaskId>> sortedTasksForTopicGroup = new TreeMap<>(tasksForTopicGroup);
+ final Set<TaskId> taskIdSet = new HashSet<>(taskIdList);
+
+ int taskNodeId = 0;
+ int topicGroupIndex = 0;
+ int tasksAssigned = 0;
+ boolean taskMoved = false;
+ for (final Entry<Subtopology, Set<TaskId>> kv : sortedTasksForTopicGroup.entrySet()) {
+ final SortedSet<TaskId> taskIds = new TreeSet<>(kv.getValue());
+ for (final TaskId taskId : taskIds) {
+ if (!taskIdSet.contains(taskId)) {
+ continue;
+ }
+ final KeyValue<Boolean, Integer> movedAndAssigned = assignTaskToClient(graph, taskId, taskNodeId, topicGroupIndex,
+ clientStates, clientList, taskIdList, taskClientMap, assignTask, unAssignTask);
+ taskMoved |= movedAndAssigned.key;
+ tasksAssigned += movedAndAssigned.value;
+ taskNodeId++;
+ }
+ topicGroupIndex++;
+ }
+
+ validateAssignedTask(taskIdList, tasksAssigned, clientStates, originalAssignedTaskNumber, hasAssignedTask);
+
+ return taskMoved;
+ }
+
+ private void validateTasks(final List<TaskId> taskIdList) {
+ final Set<TaskId> tasksInSubtopology = tasksForTopicGroup.values().stream().flatMap(
+ Collection::stream).collect(Collectors.toSet());
+ for (final TaskId taskId : taskIdList) {
+ if (!tasksInSubtopology.contains(taskId)) {
+ throw new IllegalStateException("Task " + taskId + " not in tasksForTopicGroup");
+ }
+ }
+ }
+
+ private void constructEdges(
+ final Graph<Integer> graph,
+ final List<TaskId> taskIdList,
+ final List<UUID> clientList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask,
+ final CostFunction costFunction,
+ final int trafficCost,
+ final int nonOverlapCost,
+ final boolean hasReplica,
+ final boolean isStandby
+ ) {
+ final Set<TaskId> taskIdSet = new HashSet<>(taskIdList);
+ final SortedMap<Subtopology, Set<TaskId>> sortedTasksForTopicGroup = new TreeMap<>(tasksForTopicGroup);
+ final int sinkId = getSinkNodeID(taskIdList, clientList, tasksForTopicGroup);
+
+ int taskNodeId = 0;
+ int topicGroupIndex = 0;
+ for (final Entry<Subtopology, Set<TaskId>> kv : sortedTasksForTopicGroup.entrySet()) {
+ final SortedSet<TaskId> taskIds = new TreeSet<>(kv.getValue());
+ for (int clientIndex = 0; clientIndex < clientList.size(); clientIndex++) {
+ final UUID processId = clientList.get(clientIndex);
+ final int clientNodeId = getClientNodeId(clientIndex, taskIdList, clientList, topicGroupIndex);
+ int startingTaskNodeId = taskNodeId;
+ boolean hasTask = false;
+ for (final TaskId taskId : taskIds) {
+ // It's possible some taskId is not in the tasks we want to assign. For example, taskIdSet is only stateless tasks,
+ // but the tasks in subtopology map contains all tasks including stateful ones.
+ if (!taskIdSet.contains(taskId)) {
+ continue;
+ }
+ hasTask = true;
+ final boolean inCurrentAssignment = hasAssignedTask.test(clientStates.get(processId), taskId);
+ graph.addEdge(startingTaskNodeId, clientNodeId, 1, costFunction.getCost(taskId, processId, inCurrentAssignment, trafficCost, nonOverlapCost, isStandby), 0);
+ startingTaskNodeId++;
+ if (inCurrentAssignment) {
+ if (!hasReplica && taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + processId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, processId);
+ }
+ }
+
+ if (hasTask) {
+ final int secondStageClientNodeId = getSecondStageClientNodeId(taskIdList,
+ clientList, tasksForTopicGroup, clientIndex);
+ final int capacity =
+ originalAssignedTaskNumber.containsKey(processId) ? (int) Math.ceil(
+ originalAssignedTaskNumber.get(processId) * 1.0 / taskIdList.size()
+ * taskIds.size()) : 0;
+ graph.addEdge(clientNodeId, secondStageClientNodeId, capacity, 0, 0);
+ }
+ }
+
+ taskNodeId += (int) taskIds.stream().filter(taskIdSet::contains).count();
+ topicGroupIndex++;
+ }
+
+ // Add edges from source to all tasks
+ taskNodeId = 0;
+ for (final Entry<Subtopology, Set<TaskId>> kv : sortedTasksForTopicGroup.entrySet()) {
+ final SortedSet<TaskId> taskIds = new TreeSet<>(kv.getValue());
+ for (final TaskId taskId : taskIds) {
+ if (!taskIdSet.contains(taskId)) {
+ continue;
+ }
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 0);
+ taskNodeId++;
+ }
+ }
+
+ // Add sink
+ for (int clientIndex = 0; clientIndex < clientList.size(); clientIndex++) {
+ final UUID processId = clientList.get(clientIndex);
+ final int capacity = originalAssignedTaskNumber.getOrDefault(processId, 0);
+ final int secondStageClientNodeId = getSecondStageClientNodeId(taskIdList, clientList, tasksForTopicGroup, clientIndex);
+ graph.addEdge(secondStageClientNodeId, sinkId, capacity, 0, 0);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+ }
+}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/Graph.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/Graph.java
index 396f27cfcaf..6ab87706c0c 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/Graph.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/Graph.java
@@ -19,9 +19,11 @@ package org.apache.kafka.streams.processor.internals.assignment;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedList;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
+import java.util.Queue;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;
@@ -201,6 +203,8 @@ public class Graph<V extends Comparable<V>> {
residualGraph.addEdge(edge.destination, backwardEdge);
}
}
+ residualGraph.setSourceNode(sourceNode);
+ residualGraph.setSinkNode(sinkNode);
return residualGraph;
}
@@ -251,6 +255,97 @@ public class Graph<V extends Comparable<V>> {
}
}
+ public long flow() {
+ long flow = 0;
+ final SortedMap<V, Edge> edges = adjList.get(sourceNode);
+ if (edges != null) {
+ for (final Edge edge : edges.values()) {
+ flow += edge.flow;
+ }
+ }
+ return flow;
+ }
+
+ public long calculateMaxFlow() {
+ final Graph<V> residualGraph = residualGraph();
+ residualGraph.fordFulkson();
+
+ long maxFlow = 0;
+ for (final Entry<V, SortedMap<V, Edge>> nodeEdges : adjList.entrySet()) {
+ final V node = nodeEdges.getKey();
+ for (final Entry<V, Edge> nodeEdge : nodeEdges.getValue().entrySet()) {
+ final V destination = nodeEdge.getKey();
+ final Edge edge = nodeEdge.getValue();
+ final Edge residualEdge = residualGraph.adjList.get(node).get(destination);
+ edge.flow = residualEdge.flow;
+ edge.residualFlow = residualEdge.residualFlow;
+
+ if (node == sourceNode) {
+ maxFlow += edge.flow;
+ }
+ }
+ }
+
+ return maxFlow;
+ }
+
+ private void fordFulkson() {
+ if (!isResidualGraph) {
+ throw new IllegalStateException("Should be residual graph to cancel negative cycles");
+ }
+
+ Map<V, V> parents = new HashMap<>();
+ while (breadthFirstSearch(sourceNode, sinkNode, parents)) {
+ int possibleFlow = Integer.MAX_VALUE;
+ for (V node = sinkNode; node != sourceNode; node = parents.get(node)) {
+ final V parent = parents.get(node);
+ possibleFlow = Math.min(possibleFlow, adjList.get(parent).get(node).residualFlow);
+ }
+
+ for (V node = sinkNode; node != sourceNode; node = parents.get(node)) {
+ final V parent = parents.get(node);
+ final Edge parentEdge = adjList.get(parent).get(node);
+ final Edge counterEdge = parentEdge.counterEdge;
+ parentEdge.residualFlow -= possibleFlow;
+ if (parentEdge.forwardEdge) {
+ parentEdge.flow += possibleFlow;
+ }
+ counterEdge.residualFlow += possibleFlow;
+ if (counterEdge.forwardEdge && counterEdge.flow >= possibleFlow) {
+ counterEdge.flow -= possibleFlow;
+ }
+ }
+
+ parents = new HashMap<>();
+ }
+ }
+
+ private boolean breadthFirstSearch(final V source, final V target, final Map<V, V> parents) {
+ final Set<V> visited = new HashSet<>();
+ final Queue<V> queue = new LinkedList<>();
+ queue.add(source);
+ visited.add(source);
+ while (!queue.isEmpty()) {
+ final V node = queue.poll();
+ final SortedMap<V, Edge> nodeEdges = adjList.get(node);
+ for (final Entry<V, Edge> nodeEdge : nodeEdges.entrySet()) {
+ final V nextNode = nodeEdge.getKey();
+ if (visited.contains(nextNode) || nodeEdge.getValue().residualFlow <= 0) {
+ continue;
+ }
+ if (nodeEdge.getKey().equals(target)) {
+ parents.put(target, node);
+ return true;
+ }
+ queue.add(nodeEdge.getKey());
+ parents.put(nodeEdge.getKey(), node);
+ visited.add(nodeEdge.getKey());
+ }
+ }
+
+ return false;
+ }
+
private void populateInOutFlow(final Map<V, Long> inFlow, final Map<V, Long> outFlow) {
for (final Entry<V, SortedMap<V, Edge>> nodeEdges : adjList.entrySet()) {
final V node = nodeEdges.getKey();
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/MinTrafficGraphConstructor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/MinTrafficGraphConstructor.java
new file mode 100644
index 00000000000..ec877c66a42
--- /dev/null
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/MinTrafficGraphConstructor.java
@@ -0,0 +1,145 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor.internals.assignment;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.BiConsumer;
+import java.util.function.BiPredicate;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.CostFunction;
+
+public class MinTrafficGraphConstructor implements RackAwareGraphConstructor {
+
+ @Override
+ public int getSinkNodeID(
+ final List<TaskId> taskIdList,
+ final List<UUID> clientList,
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup
+ ) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ @Override
+ public int getClientNodeId(final int clientIndex, final List<TaskId> taskIdList, final List<UUID> clientList, final int topicGroupIndex) {
+ return clientIndex + taskIdList.size();
+ }
+
+ @Override
+ public int getClientIndex(final int clientNodeId, final List<TaskId> taskIdList, final List<UUID> clientList, final int topicGroupIndex) {
+ return clientNodeId - taskIdList.size();
+ }
+
+ @Override
+ public Graph<Integer> constructTaskGraph(
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask,
+ final CostFunction costFunction,
+ final int trafficCost,
+ final int nonOverlapCost,
+ final boolean hasReplica,
+ final boolean isStandby
+ ) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : taskIdList) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (hasAssignedTask.test(clientState.getValue(), taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = getClientNodeId(j, taskIdList, null, -1);
+ final UUID processId = clientList.get(j);
+
+ final int flow = hasAssignedTask.test(clientStates.get(processId), taskId) ? 1 : 0;
+ final int cost = costFunction.getCost(taskId, processId, flow == 1, trafficCost,
+ nonOverlapCost, isStandby);
+ if (flow == 1) {
+ if (!hasReplica && taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + processId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, processId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+
+ // Add edge from source to task
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ final int sinkId = getSinkNodeID(taskIdList, clientList, null);
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have higher traffic cost. This is to maintain the original assigned task count
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = getClientNodeId(i, taskIdList, null, -1);
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ @Override
+ public boolean assignTaskFromMinCostFlow(
+ final Graph<Integer> graph,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final Map<TaskId, UUID> taskClientMap,
+ final BiConsumer<ClientState, TaskId> assignTask,
+ final BiConsumer<ClientState, TaskId> unAssignTask,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask
+ ) {
+ int tasksAssigned = 0;
+ boolean taskMoved = false;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ final KeyValue<Boolean, Integer> movedAndAssigned = assignTaskToClient(graph, taskId, taskNodeId, -1,
+ clientStates, clientList, taskIdList, taskClientMap, assignTask, unAssignTask);
+ taskMoved |= movedAndAssigned.key;
+ tasksAssigned += movedAndAssigned.value;
+ }
+
+ validateAssignedTask(taskIdList, tasksAssigned, clientStates, originalAssignedTaskNumber, hasAssignedTask);
+ return taskMoved;
+ }
+}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructor.java
new file mode 100644
index 00000000000..4b46c147e17
--- /dev/null
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructor.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor.internals.assignment;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Objects;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.BiConsumer;
+import java.util.function.BiPredicate;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.CostFunction;
+
+/**
+ * Construct graph for rack aware task assignor
+ */
+public interface RackAwareGraphConstructor {
+ int SOURCE_ID = -1;
+
+ int getSinkNodeID(final List<TaskId> taskIdList, final List<UUID> clientList, final Map<Subtopology, Set<TaskId>> tasksForTopicGroup);
+
+ int getClientNodeId(final int clientIndex, final List<TaskId> taskIdList, final List<UUID> clientList, final int topicGroupIndex);
+
+ int getClientIndex(final int clientNodeId, final List<TaskId> taskIdList, final List<UUID> clientList, final int topicGroupIndex);
+
+ Graph<Integer> constructTaskGraph(
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask,
+ final CostFunction costFunction,
+ final int trafficCost,
+ final int nonOverlapCost,
+ final boolean hasReplica,
+ final boolean isStandby);
+
+ boolean assignTaskFromMinCostFlow(
+ final Graph<Integer> graph,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final Map<TaskId, UUID> taskClientMap,
+ final BiConsumer<ClientState, TaskId> assignTask,
+ final BiConsumer<ClientState, TaskId> unAssignTask,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask);
+
+ default KeyValue<Boolean, Integer> assignTaskToClient(
+ final Graph<Integer> graph,
+ final TaskId taskId,
+ final int taskNodeId,
+ final int topicGroupIndex,
+ final Map<UUID, ClientState> clientStates,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<TaskId, UUID> taskClientMap,
+ final BiConsumer<ClientState, TaskId> assignTask,
+ final BiConsumer<ClientState, TaskId> unAssignTask
+ ) {
+ int tasksAssigned = 0;
+ boolean taskMoved = false;
+ final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ if (edge.flow > 0) {
+ tasksAssigned++;
+ final int clientIndex = getClientIndex(edge.destination, taskIdList, clientList, topicGroupIndex);
+ final UUID processId = clientList.get(clientIndex);
+ final UUID originalProcessId = taskClientMap.get(taskId);
+
+ // Don't need to assign this task to other client
+ if (processId.equals(originalProcessId)) {
+ break;
+ }
+
+ unAssignTask.accept(clientStates.get(originalProcessId), taskId);
+ assignTask.accept(clientStates.get(processId), taskId);
+ taskMoved = true;
+ }
+ }
+ return KeyValue.pair(taskMoved, tasksAssigned);
+ }
+
+ default void validateAssignedTask(
+ final List<TaskId> taskIdList,
+ final int tasksAssigned,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final BiPredicate<ClientState, TaskId> hasAssignedTask
+ ) {
+ // Validate task assigned
+ if (tasksAssigned != taskIdList.size()) {
+ throw new IllegalStateException("Computed active task assignment number "
+ + tasksAssigned + " is different size " + taskIdList.size());
+ }
+
+ // Validate original assigned task number matches
+ final Map<UUID, Integer> assignedTaskNumber = new HashMap<>();
+ for (final TaskId taskId : taskIdList) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (hasAssignedTask.test(clientState.getValue(), taskId)) {
+ assignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ if (originalAssignedTaskNumber.size() != assignedTaskNumber.size()) {
+ throw new IllegalStateException("There are " + originalAssignedTaskNumber.size() + " clients have "
+ + " active tasks before assignment, but " + assignedTaskNumber.size() + " clients have"
+ + " active tasks after assignment");
+ }
+
+ for (final Entry<UUID, Integer> originalCapacity : originalAssignedTaskNumber.entrySet()) {
+ final int capacity = assignedTaskNumber.getOrDefault(originalCapacity.getKey(), 0);
+ if (!Objects.equals(originalCapacity.getValue(), capacity)) {
+ throw new IllegalStateException("There are " + originalCapacity.getValue() + " tasks assigned to"
+ + " client " + originalCapacity.getKey() + " before assignment, but " + capacity + " tasks "
+ + " are assigned to it after assignment");
+ }
+ }
+ }
+}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorFactory.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorFactory.java
new file mode 100644
index 00000000000..c23dde358b3
--- /dev/null
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorFactory.java
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor.internals.assignment;
+
+import java.util.Map;
+import java.util.Set;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
+
+public class RackAwareGraphConstructorFactory {
+
+ static RackAwareGraphConstructor create(final AssignmentConfigs assignmentConfigs, final Map<Subtopology, Set<TaskId>> tasksForTopicGroup) {
+ switch (assignmentConfigs.rackAwareAssignmentStrategy) {
+ case StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC:
+ return new MinTrafficGraphConstructor();
+ case StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY:
+ return new BalanceSubtopologyGraphConstructor(tasksForTopicGroup);
+ default:
+ throw new IllegalArgumentException("Rack aware assignment is disabled");
+ }
+ }
+
+}
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 18e3d78d4d7..b124d07424a 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
@@ -23,7 +23,6 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
@@ -31,7 +30,6 @@ import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.UUID;
-import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.BiPredicate;
import java.util.stream.Collectors;
@@ -61,12 +59,22 @@ public class RackAwareTaskAssignor {
final Map<UUID, ClientState> clientStateMap);
}
+ @FunctionalInterface
+ public interface CostFunction {
+ int getCost(final TaskId taskId,
+ final UUID processId,
+ final boolean inCurrentAssignment,
+ final int trafficCost,
+ final int nonOverlapCost,
+ final boolean isStandby);
+ }
+
// For stateless tasks, it's ok to move them around. So we have 0 non_overlap_cost
public static final int STATELESS_TRAFFIC_COST = 1;
public static final int STATELESS_NON_OVERLAP_COST = 0;
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
- private static final int SOURCE_ID = -1;
+
// 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;
@@ -74,6 +82,7 @@ public class RackAwareTaskAssignor {
private final Cluster fullMetadata;
private final Map<TaskId, Set<TopicPartition>> partitionsForTask;
private final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask;
+ private final Map<Subtopology, Set<TaskId>> tasksForTopicGroup;
private final AssignmentConfigs assignmentConfigs;
private final Map<TopicPartition, Set<String>> racksForPartition;
private final Map<UUID, String> racksForProcess;
@@ -93,6 +102,7 @@ public class RackAwareTaskAssignor {
this.fullMetadata = fullMetadata;
this.partitionsForTask = partitionsForTask;
this.changelogPartitionsForTask = changelogPartitionsForTask;
+ this.tasksForTopicGroup = tasksForTopicGroup;
this.internalTopicManager = internalTopicManager;
this.assignmentConfigs = assignmentConfigs;
this.racksForPartition = new HashMap<>();
@@ -283,18 +293,6 @@ public class RackAwareTaskAssignor {
return cost;
}
- private static int getSinkNodeID(final List<UUID> clientList, final List<TaskId> taskIdList) {
- return clientList.size() + taskIdList.size();
- }
-
- private static int getClientNodeId(final List<TaskId> taskIdList, final int clientIndex) {
- return clientIndex + taskIdList.size();
- }
-
- private static int getClientIndex(final List<TaskId> taskIdList, final int clientNodeId) {
- return clientNodeId - taskIdList.size();
- }
-
/**
* Compute the cost for the provided {@code activeTasks}. The passed in active tasks must be contained in {@code clientState}.
*/
@@ -327,8 +325,20 @@ public class RackAwareTaskAssignor {
}
final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
final List<TaskId> taskIdList = new ArrayList<>(tasks);
- final Graph<Integer> graph = constructTaskGraph(clientList, taskIdList,
- clientStates, new HashMap<>(), new HashMap<>(), hasAssignedTask, trafficCost, nonOverlapCost, hasReplica, isStandby);
+ final Graph<Integer> graph = new MinTrafficGraphConstructor()
+ .constructTaskGraph(
+ clientList,
+ taskIdList,
+ clientStates,
+ new HashMap<>(),
+ new HashMap<>(),
+ hasAssignedTask,
+ this::getCost,
+ trafficCost,
+ nonOverlapCost,
+ hasReplica,
+ isStandby
+ );
return graph.totalCost();
}
@@ -363,13 +373,25 @@ public class RackAwareTaskAssignor {
final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
final Map<TaskId, UUID> taskClientMap = new HashMap<>();
final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
- final Graph<Integer> graph = constructTaskGraph(clientList, taskIdList,
- clientStates, taskClientMap, originalAssignedTaskNumber, ClientState::hasActiveTask, trafficCost, nonOverlapCost, false, false);
+ final RackAwareGraphConstructor graphConstructor = RackAwareGraphConstructorFactory.create(assignmentConfigs, tasksForTopicGroup);
+ final Graph<Integer> graph = graphConstructor.constructTaskGraph(
+ clientList,
+ taskIdList,
+ clientStates,
+ taskClientMap,
+ originalAssignedTaskNumber,
+ ClientState::hasActiveTask,
+ this::getCost,
+ trafficCost,
+ nonOverlapCost,
+ false,
+ false
+ );
graph.solveMinCostFlow();
final long cost = graph.totalCost();
- assignTaskFromMinCostFlow(graph, clientList, taskIdList, clientStates, originalAssignedTaskNumber,
+ graphConstructor.assignTaskFromMinCostFlow(graph, clientList, taskIdList, clientStates, originalAssignedTaskNumber,
taskClientMap, ClientState::assignActive, ClientState::unassignActive, ClientState::hasActiveTask);
final long duration = time.milliseconds() - startTime;
@@ -397,6 +419,7 @@ public class RackAwareTaskAssignor {
boolean taskMoved = true;
int round = 0;
+ final RackAwareGraphConstructor graphConstructor = new MinTrafficGraphConstructor();
while (taskMoved && round < STANDBY_OPTIMIZER_MAX_ITERATION) {
taskMoved = false;
round++;
@@ -432,12 +455,22 @@ public class RackAwareTaskAssignor {
Collectors.toList());
final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
- final Graph<Integer> graph = constructTaskGraph(clients, taskIdList,
- clientStates, taskClientMap, originalAssignedTaskNumber,
- ClientState::hasStandbyTask, trafficCost, nonOverlapCost, true, true);
+ final Graph<Integer> graph = graphConstructor.constructTaskGraph(
+ clients,
+ taskIdList,
+ clientStates,
+ taskClientMap,
+ originalAssignedTaskNumber,
+ ClientState::hasStandbyTask,
+ this::getCost,
+ trafficCost,
+ nonOverlapCost,
+ true,
+ true
+ );
graph.solveMinCostFlow();
- taskMoved |= assignTaskFromMinCostFlow(graph, clients, taskIdList, clientStates,
+ taskMoved |= graphConstructor.assignTaskFromMinCostFlow(graph, clients, taskIdList, clientStates,
originalAssignedTaskNumber,
taskClientMap, ClientState::assignStandby, ClientState::unassignStandby,
ClientState::hasStandbyTask);
@@ -450,134 +483,4 @@ public class RackAwareTaskAssignor {
log.info("Assignment after {} rounds and {} milliseconds for standby task optimization is {}\n with cost {}", round, duration, clientStates, cost);
return cost;
}
-
- private Graph<Integer> constructTaskGraph(final List<UUID> clientList,
- final List<TaskId> taskIdList,
- final Map<UUID, ClientState> clientStates,
- final Map<TaskId, UUID> taskClientMap,
- final Map<UUID, Integer> originalAssignedTaskNumber,
- final BiPredicate<ClientState, TaskId> hasAssignedTask,
- final int trafficCost,
- final int nonOverlapCost,
- final boolean hasReplica,
- final boolean isStandby) {
- final Graph<Integer> graph = new Graph<>();
-
- for (final TaskId taskId : taskIdList) {
- for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
- if (hasAssignedTask.test(clientState.getValue(), taskId)) {
- originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
- }
- }
- }
-
- // Make task and client Node id in graph deterministic
- for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
- final TaskId taskId = taskIdList.get(taskNodeId);
- for (int j = 0; j < clientList.size(); j++) {
- final int clientNodeId = getClientNodeId(taskIdList, j);
- final UUID processId = clientList.get(j);
-
- final int flow = hasAssignedTask.test(clientStates.get(processId), taskId) ? 1 : 0;
- final int cost = getCost(taskId, processId, flow == 1, trafficCost, nonOverlapCost, isStandby);
- if (flow == 1) {
- if (!hasReplica && taskClientMap.containsKey(taskId)) {
- throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
- + processId + ", " + taskClientMap.get(taskId));
- }
- taskClientMap.put(taskId, processId);
- }
-
- graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
- }
- if (!taskClientMap.containsKey(taskId)) {
- throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
- }
-
- // Add edge from source to task
- graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
- }
-
- final int sinkId = getSinkNodeID(clientList, taskIdList);
- // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
- // even though it may have higher traffic cost. This is to maintain the original assigned task count
- for (int i = 0; i < clientList.size(); i++) {
- final int clientNodeId = getClientNodeId(taskIdList, i);
- final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
- // Flow equals to capacity for edges to sink
- graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
- }
-
- graph.setSourceNode(SOURCE_ID);
- graph.setSinkNode(sinkId);
-
- return graph;
- }
-
- private boolean assignTaskFromMinCostFlow(final Graph<Integer> graph,
- final List<UUID> clientList,
- final List<TaskId> taskIdList,
- final Map<UUID, ClientState> clientStates,
- final Map<UUID, Integer> originalAssignedTaskNumber,
- final Map<TaskId, UUID> taskClientMap,
- final BiConsumer<ClientState, TaskId> assignTask,
- final BiConsumer<ClientState, TaskId> unAssignTask,
- final BiPredicate<ClientState, TaskId> hasAssignedTask) {
- int tasksAssigned = 0;
- boolean taskMoved = false;
- for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
- final TaskId taskId = taskIdList.get(taskNodeId);
- final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
- for (final Graph<Integer>.Edge edge : edges.values()) {
- if (edge.flow > 0) {
- tasksAssigned++;
- final int clientIndex = getClientIndex(taskIdList, edge.destination);
- final UUID processId = clientList.get(clientIndex);
- final UUID originalProcessId = taskClientMap.get(taskId);
-
- // Don't need to assign this task to other client
- if (processId.equals(originalProcessId)) {
- break;
- }
-
- unAssignTask.accept(clientStates.get(originalProcessId), taskId);
- assignTask.accept(clientStates.get(processId), taskId);
- taskMoved = true;
- }
- }
- }
-
- // Validate task assigned
- if (tasksAssigned != taskIdList.size()) {
- throw new IllegalStateException("Computed active task assignment number "
- + tasksAssigned + " is different size " + taskIdList.size());
- }
-
- // Validate original assigned task number matches
- final Map<UUID, Integer> assignedTaskNumber = new HashMap<>();
- for (final TaskId taskId : taskIdList) {
- for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
- if (hasAssignedTask.test(clientState.getValue(), taskId)) {
- assignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
- }
- }
- }
-
- if (originalAssignedTaskNumber.size() != assignedTaskNumber.size()) {
- throw new IllegalStateException("There are " + originalAssignedTaskNumber.size() + " clients have "
- + " active tasks before assignment, but " + assignedTaskNumber.size() + " clients have"
- + " active tasks after assignment");
- }
-
- for (final Entry<UUID, Integer> originalCapacity : originalAssignedTaskNumber.entrySet()) {
- final int capacity = assignedTaskNumber.getOrDefault(originalCapacity.getKey(), 0);
- if (!Objects.equals(originalCapacity.getValue(), capacity)) {
- throw new IllegalStateException("There are " + originalCapacity.getValue() + " tasks assigned to"
- + " client " + originalCapacity.getKey() + " before assignment, but " + capacity + " tasks "
- + " are assigned to it after assignment");
- }
- }
-
- return taskMoved;
- }
}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TopologyMetadataTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TopologyMetadataTest.java
index 15cf2ffeeb7..8856f6ae2fd 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TopologyMetadataTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TopologyMetadataTest.java
@@ -17,12 +17,15 @@
package org.apache.kafka.streams.processor.internals;
import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
import org.apache.kafka.streams.processor.internals.testutil.DummyStreamsConfig;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.junit.MockitoJUnitRunner;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
@@ -43,11 +46,35 @@ public class TopologyMetadataTest {
Assert.assertFalse(topologyMetadata.isPaused(TOPOLOGY2));
topologyMetadata.pauseTopology(TOPOLOGY1);
- Assert.assertTrue(topologyMetadata.isPaused(TOPOLOGY1));
+ assertTrue(topologyMetadata.isPaused(TOPOLOGY1));
Assert.assertFalse(topologyMetadata.isPaused(TOPOLOGY2));
topologyMetadata.resumeTopology(TOPOLOGY1);
Assert.assertFalse(topologyMetadata.isPaused(TOPOLOGY1));
Assert.assertFalse(topologyMetadata.isPaused(TOPOLOGY2));
}
+
+ @Test
+ public void testSubtopologyCompare() {
+ Subtopology subtopology1 = new Subtopology(0, "1");
+ Subtopology subtopology2 = new Subtopology(1, "1");
+
+ assertTrue(subtopology1.compareTo(subtopology2) < 0);
+
+ subtopology1 = new Subtopology(0, null);
+ subtopology2 = new Subtopology(0, null);
+ assertEquals(0, subtopology1.compareTo(subtopology2));
+
+ subtopology1 = new Subtopology(0, null);
+ subtopology2 = new Subtopology(0, "1");
+ assertTrue(subtopology1.compareTo(subtopology2) < 0);
+
+ subtopology1 = new Subtopology(0, "1");
+ subtopology2 = new Subtopology(0, null);
+ assertTrue(subtopology1.compareTo(subtopology2) > 0);
+
+ subtopology1 = new Subtopology(0, "1");
+ subtopology2 = new Subtopology(0, "2");
+ assertTrue(subtopology1.compareTo(subtopology2) < 0);
+ }
}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java
index 76348fbcef5..626032f146e 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java
@@ -676,18 +676,27 @@ public final class AssignmentTestUtils {
return taskTopicPartitionMap;
}
- static Map<String, Object> configProps(final boolean enableRackAwareAssignor) {
- return configProps(enableRackAwareAssignor, 0);
+ static Map<Subtopology, Set<TaskId>> getTasksForTopicGroup(final int tpSize, final int partitionSize) {
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = new HashMap<>();
+ for (int i = 0; i < tpSize; i++) {
+ for (int j = 0; j < partitionSize; j++) {
+ final Subtopology subtopology = new Subtopology(i, null);
+ tasksForTopicGroup.computeIfAbsent(subtopology, k -> new HashSet<>()).add(new TaskId(i, j));
+ }
+ }
+ return tasksForTopicGroup;
+ }
+
+ static Map<String, Object> configProps(final String rackAwareConfig) {
+ return configProps(rackAwareConfig, 0);
}
- static Map<String, Object> configProps(final boolean enableRackAwareAssignor, final int replicaNum) {
+ static Map<String, Object> configProps(final String rackAwareConfig, final int replicaNum) {
final Map<String, Object> configurationMap = new HashMap<>();
configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
configurationMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, USER_END_POINT);
configurationMap.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, replicaNum);
- if (enableRackAwareAssignor) {
- configurationMap.put(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_CONFIG, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC);
- }
+ configurationMap.put(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_CONFIG, rackAwareConfig);
final ReferenceContainer referenceContainer = new ReferenceContainer();
configurationMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
@@ -696,7 +705,7 @@ public final class AssignmentTestUtils {
static InternalTopicManager mockInternalTopicManagerForRandomChangelog(final int nodeSize, final int tpSize, final int partitionSize) {
final MockTime time = new MockTime();
- final StreamsConfig streamsConfig = new StreamsConfig(configProps(true));
+ final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
final MockClientSupplier mockClientSupplier = new MockClientSupplier();
final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
time,
@@ -846,6 +855,15 @@ public final class AssignmentTestUtils {
);
}
+ static Map<Subtopology, Set<TaskId>> getTasksForTopicGroup() {
+ return mkMap(
+ mkEntry(new Subtopology(0, null), mkSet(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_0_4, TASK_0_5, TASK_0_6)),
+ mkEntry(new Subtopology(1, null), mkSet(TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3)),
+ mkEntry(new Subtopology(2, null), mkSet(TASK_2_0, TASK_2_1, TASK_2_2, TASK_2_3)),
+ mkEntry(new Subtopology(3, null), mkSet(TASK_3_0, TASK_3_1, TASK_3_2))
+ );
+ }
+
static Map<TaskId, Set<TopicPartition>> getTaskChangelogMapForAllTasks() {
return mkMap(
mkEntry(TASK_0_0, mkSet(CHANGELOG_TP_0_0)),
@@ -871,7 +889,7 @@ public final class AssignmentTestUtils {
static InternalTopicManager mockInternalTopicManagerForChangelog() {
final MockTime time = new MockTime();
- final StreamsConfig streamsConfig = new StreamsConfig(configProps(true));
+ final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
final MockClientSupplier mockClientSupplier = new MockClientSupplier();
final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
time,
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/GraphTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/GraphTest.java
index 629c572dec0..10a7f2ec69e 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/GraphTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/GraphTest.java
@@ -394,6 +394,106 @@ public class GraphTest {
}
}
+ @Test
+ public void testMaxFlowOnlySourceAndSink() {
+ final Graph<Integer> graph1 = new Graph<>();
+ graph1.addEdge(0, 1, 100, 0, 0);
+ graph1.setSourceNode(0);
+ graph1.setSinkNode(1);
+
+ final long maxFlow = graph1.calculateMaxFlow();
+ assertEquals(100, maxFlow);
+
+ final Map<Integer, Graph<Integer>.Edge> edges = graph1.edges(0);
+ assertEquals(100, edges.get(1).flow);
+ assertEquals(0, edges.get(1).residualFlow);
+
+ }
+
+ @Test
+ public void testMaxFlowBoundBySinkEdges() {
+ // Edges connected to sink have less capacity
+ final Graph<Integer> graph1 = new Graph<>();
+ graph1.addEdge(0, 1, 10, 0, 0);
+ graph1.addEdge(0, 2, 10, 0, 0);
+ graph1.addEdge(0, 3, 10, 0, 0);
+ graph1.addEdge(1, 4, 5, 0, 0);
+ graph1.addEdge(2, 5, 5, 0, 0);
+ graph1.addEdge(3, 6, 5, 0, 0);
+ graph1.addEdge(4, 7, 2, 0, 0);
+ graph1.addEdge(5, 7, 2, 0, 0);
+ graph1.addEdge(6, 7, 2, 0, 0);
+ graph1.setSourceNode(0);
+ graph1.setSinkNode(7);
+ final long maxFlow = graph1.calculateMaxFlow();
+ assertEquals(6, maxFlow);
+
+ Map<Integer, Graph<Integer>.Edge> edges = graph1.edges(0);
+ assertEquals(2, edges.get(1).flow);
+ assertEquals(2, edges.get(2).flow);
+ assertEquals(2, edges.get(3).flow);
+
+ edges = graph1.edges(1);
+ assertEquals(2, edges.get(4).flow);
+
+ edges = graph1.edges(2);
+ assertEquals(2, edges.get(5).flow);
+
+ edges = graph1.edges(3);
+ assertEquals(2, edges.get(6).flow);
+
+ edges = graph1.edges(4);
+ assertEquals(2, edges.get(7).flow);
+
+ edges = graph1.edges(5);
+ assertEquals(2, edges.get(7).flow);
+
+ edges = graph1.edges(6);
+ assertEquals(2, edges.get(7).flow);
+ }
+
+ @Test
+ public void testMaxFlowBoundBySourceEdges() {
+ // Edges connected to source have less capacity
+ final Graph<Integer> graph1 = new Graph<>();
+ graph1.addEdge(0, 1, 1, 0, 0);
+ graph1.addEdge(0, 2, 2, 0, 0);
+ graph1.addEdge(0, 3, 3, 0, 0);
+ graph1.addEdge(1, 4, 5, 0, 0);
+ graph1.addEdge(2, 5, 5, 0, 0);
+ graph1.addEdge(3, 6, 5, 0, 0);
+ graph1.addEdge(4, 7, 10, 0, 0);
+ graph1.addEdge(5, 7, 10, 0, 0);
+ graph1.addEdge(6, 7, 10, 0, 0);
+ graph1.setSourceNode(0);
+ graph1.setSinkNode(7);
+ final long maxFlow = graph1.calculateMaxFlow();
+ assertEquals(6, maxFlow);
+
+ Map<Integer, Graph<Integer>.Edge> edges = graph1.edges(0);
+ assertEquals(1, edges.get(1).flow);
+ assertEquals(2, edges.get(2).flow);
+ assertEquals(3, edges.get(3).flow);
+
+ edges = graph1.edges(1);
+ assertEquals(1, edges.get(4).flow);
+
+ edges = graph1.edges(2);
+ assertEquals(2, edges.get(5).flow);
+
+ edges = graph1.edges(3);
+ assertEquals(3, edges.get(6).flow);
+
+ edges = graph1.edges(4);
+ assertEquals(1, edges.get(7).flow);
+
+ edges = graph1.edges(5);
+ assertEquals(2, edges.get(7).flow);
+
+ edges = graph1.edges(6);
+ assertEquals(3, edges.get(7).flow);
+ }
+
private static Graph<Integer>.Edge getEdge(final int destination, final int capacity, final int cost, final int residualFlow, final int flow) {
return getEdge(destination, capacity, cost, residualFlow, flow, true);
}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorFactoryTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorFactoryTest.java
new file mode 100644
index 00000000000..a4c17ea3e3c
--- /dev/null
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorFactoryTest.java
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor.internals.assignment;
+
+import static org.apache.kafka.common.utils.Utils.mkMap;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.configProps;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.instanceOf;
+
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
+import org.junit.jupiter.api.Test;
+
+public class RackAwareGraphConstructorFactoryTest {
+
+ @Test
+ public void shouldReturnMinCostConstructor() {
+ final AssignmentConfigs config = new AssignorConfiguration(
+ new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC)).originals()).assignmentConfigs();
+ final RackAwareGraphConstructor constructor = RackAwareGraphConstructorFactory.create(config, mkMap());
+ assertThat(constructor, instanceOf(MinTrafficGraphConstructor.class));
+ }
+
+ @Test
+ public void shouldReturnBalanceSubtopologyConstructor() {
+ final AssignmentConfigs config = new AssignorConfiguration(
+ new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY)).originals()).assignmentConfigs();
+ final RackAwareGraphConstructor constructor = RackAwareGraphConstructorFactory.create(config, mkMap());
+ assertThat(constructor, instanceOf(BalanceSubtopologyGraphConstructor.class));
+ }
+}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorTest.java
new file mode 100644
index 00000000000..607dff93825
--- /dev/null
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareGraphConstructorTest.java
@@ -0,0 +1,301 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.processor.internals.assignment;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptySet;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertBalancedTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertValidAssignment;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomClientState;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMap;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTasksForTopicGroup;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.SortedMap;
+import java.util.SortedSet;
+import java.util.UUID;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+
+@RunWith(Parameterized.class)
+public class RackAwareGraphConstructorTest {
+ private static final String MIN_COST = "min_cost";
+ private static final String BALANCE_SUBTOPOLOGY = "balance_sub_topology";
+
+ private static final int TP_SIZE = 40;
+ private static final int PARTITION_SIZE = 3;
+ private static final int TOPIC_GROUP_SIZE = TP_SIZE;
+ private static final int CLIENT_SIZE = 20;
+
+ private Graph<Integer> graph;
+ private final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
+ TP_SIZE, PARTITION_SIZE, false);
+ private final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ private final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ private final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(CLIENT_SIZE,
+ TP_SIZE, PARTITION_SIZE, 1, false, taskIds);
+ private final List<UUID> clientList = new ArrayList<>(clientStateMap.keySet());
+ private final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ private final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ private final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = getTasksForTopicGroup(TP_SIZE,
+ PARTITION_SIZE);
+ private RackAwareGraphConstructor constructor;
+
+ @Parameter
+ public String constructorType;
+
+ @Parameterized.Parameters(name = "constructorType={0}")
+ public static Collection<Object[]> getParamStoreType() {
+ return asList(new Object[][] {
+ {MIN_COST},
+ {BALANCE_SUBTOPOLOGY}
+ });
+ }
+
+ @Before
+ public void setUp() {
+ randomAssignTasksToClient(taskIdList, clientStateMap);
+
+ if (constructorType.equals(MIN_COST)) {
+ constructor = new MinTrafficGraphConstructor();
+ } else if (constructorType.equals(BALANCE_SUBTOPOLOGY)) {
+ constructor = new BalanceSubtopologyGraphConstructor(tasksForTopicGroup);
+ }
+ graph = constructor.constructTaskGraph(
+ clientList, taskIdList, clientStateMap, taskClientMap, originalAssignedTaskNumber, ClientState::hasAssignedTask, this::getCost, 10, 1, false, false);
+ }
+
+ private int getCost(final TaskId taskId, final UUID processId, final boolean inCurrentAssignment, final int trafficCost, final int nonOverlapCost, final boolean isStandby) {
+ return 1;
+ }
+
+ @Test
+ public void testSubtopologyShouldContainAllTasks() {
+ if (constructorType.equals(MIN_COST)) {
+ return;
+ }
+ taskIdList.add(new TaskId(41, 0)); // Extra task not in subtopology map
+ assertThrows(IllegalStateException.class, () -> graph = constructor.constructTaskGraph(
+ clientList, taskIdList, clientStateMap, taskClientMap, originalAssignedTaskNumber,
+ ClientState::hasAssignedTask, this::getCost, 10, 1, false, false));
+ }
+
+ @Test
+ public void testMinCostGraphConstructor() {
+ if (constructorType.equals(BALANCE_SUBTOPOLOGY)) {
+ return;
+ }
+ // Flow should equal to task size which means each task is assigned now
+ assertEquals(taskIdList.size(), graph.flow());
+
+ // Total node size should match. Extra 2 is for source and sink nodes
+ assertEquals(taskIdList.size() + clientList.size() + 2, graph.nodes().size());
+
+ // Check edges from source to task nodes
+ SortedMap<Integer, Graph<Integer>.Edge> edges = graph.edges(RackAwareGraphConstructor.SOURCE_ID);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ assertEquals(1, edge.flow);
+ assertEquals(1, edge.capacity);
+ assertEquals(0, edge.residualFlow);
+ assertEquals(0, edge.cost);
+ Assertions.assertTrue(edge.forwardEdge);
+ }
+
+ // Check edges from task nodes to clients nodes,
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ edges = graph.edges(taskNodeId);
+ assertEquals(clientList.size(), edges.size());
+ int assignedClient = 0;
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ final int flow = edge.flow;
+ if (flow == 1) {
+ assignedClient++;
+ }
+ assertEquals(1, edge.capacity);
+ assertEquals(flow == 1 ? 0 : 1, edge.residualFlow);
+ assertEquals(1, edge.cost);
+ Assertions.assertTrue(edge.forwardEdge);
+ }
+ // Task should be assigned to exact 1 client
+ assertEquals(1, assignedClient);
+ taskNodeId++;
+ }
+
+ // Check edges from second stage client node to sink
+ final int sinkId = clientList.size() + taskIdList.size();
+ int totalFlow = 0;
+ for (int i = 0; i < clientList.size(); i++) {
+ final UUID clientId = clientList.get(i);
+ final int originalAssignedCount = originalAssignedTaskNumber.get(clientId);
+
+ final int clientNodeId = i + taskIdList.size();
+ edges = graph.edges(clientNodeId);
+ assertEquals(1, edges.size());
+ for (final Entry<Integer, Graph<Integer>.Edge> nodeEdge : edges.entrySet()) {
+ final Integer nodeId = nodeEdge.getKey();
+ assertEquals(sinkId, nodeId);
+ totalFlow += nodeEdge.getValue().flow;
+ assertEquals(originalAssignedCount, nodeEdge.getValue().capacity);
+ Assertions.assertTrue(nodeEdge.getValue().forwardEdge);
+ }
+ }
+ assertEquals(taskIdList.size(), totalFlow);
+ }
+
+ @Test
+ public void testBalanceSubtopologyGraphConstructor() {
+ if (constructorType.equals(MIN_COST)) {
+ return;
+ }
+ // Flow should equal to task size which means each task is assigned now
+ assertEquals(taskIdList.size(), graph.flow());
+
+ // Total node size should match. Extra 2 is for source and sink nodes
+ assertEquals(taskIdList.size() + TOPIC_GROUP_SIZE * clientList.size() + clientList.size() + 2, graph.nodes().size());
+
+ // Check edges from source to task nodes
+ SortedMap<Integer, Graph<Integer>.Edge> edges = graph.edges(RackAwareGraphConstructor.SOURCE_ID);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ assertEquals(1, edge.flow);
+ assertEquals(1, edge.capacity);
+ assertEquals(0, edge.residualFlow);
+ assertEquals(0, edge.cost);
+ Assertions.assertTrue(edge.forwardEdge);
+ }
+
+ // Check edges from task nodes to first stage clients node,
+ int taskNodeId = 0;
+ for (final Set<TaskId> unsortedTaskIds : tasksForTopicGroup.values()) {
+ for (int i = 0; i < unsortedTaskIds.size(); i++) {
+ edges = graph.edges(taskNodeId);
+ assertEquals(clientList.size(), edges.size());
+ int assignedClient = 0;
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ final int flow = edge.flow;
+ if (flow == 1) {
+ assignedClient++;
+ }
+ assertEquals(1, edge.capacity);
+ assertEquals(flow == 1 ? 0 : 1, edge.residualFlow);
+ assertEquals(1, edge.cost);
+ Assertions.assertTrue(edge.forwardEdge);
+ }
+ // Task should be assigned to exact 1 client
+ assertEquals(1, assignedClient);
+ taskNodeId++;
+ }
+ }
+
+ // Check edges from first stage client node to second stage client node
+ int topicGroupIndex = 0;
+ for (final Set<TaskId> tasks : tasksForTopicGroup.values()) {
+ final int taskCount = tasks.size();
+ for (int j = 0; j < clientList.size(); j++) {
+ final UUID clientId = clientList.get(j);
+ final int originalAssignedCount = originalAssignedTaskNumber.get(clientId);
+ final int expectedCapacity = (int) Math.ceil(originalAssignedCount * 1.0 / taskIdList.size() * taskCount);
+ final int clientNodeId = topicGroupIndex * clientList.size() + taskIdList.size() + j;
+ edges = graph.edges(clientNodeId);
+ assertEquals(1, edges.size());
+ for (final Entry<Integer, Graph<Integer>.Edge> nodeEdge : edges.entrySet()) {
+ final Integer nodeId = nodeEdge.getKey();
+ assertEquals(clientList.size() * tasksForTopicGroup.size() + taskIdList.size() + j, nodeId);
+ final Graph<Integer>.Edge edge = nodeEdge.getValue();
+ assertEquals(expectedCapacity, edge.capacity);
+ assertEquals(0, edge.cost);
+ Assertions.assertTrue(edge.forwardEdge);
+ }
+ }
+ topicGroupIndex++;
+ }
+
+ // Check edges from second stage client node to sink
+ final int sinkId = clientList.size() + tasksForTopicGroup.size() * clientList.size() + taskIdList.size();
+ int totalFlow = 0;
+ for (int i = 0; i < clientList.size(); i++) {
+ final UUID clientId = clientList.get(i);
+ final int originalAssignedCount = originalAssignedTaskNumber.get(clientId);
+
+ final int clientNodeId =
+ i + tasksForTopicGroup.size() * clientList.size() + taskIdList.size();
+ edges = graph.edges(clientNodeId);
+ assertEquals(1, edges.size());
+ for (final Entry<Integer, Graph<Integer>.Edge> nodeEdge : edges.entrySet()) {
+ final Integer nodeId = nodeEdge.getKey();
+ assertEquals(sinkId, nodeId);
+ totalFlow += nodeEdge.getValue().flow;
+ assertEquals(originalAssignedCount, nodeEdge.getValue().capacity);
+ Assertions.assertTrue(nodeEdge.getValue().forwardEdge);
+ }
+ }
+ assertEquals(taskIdList.size(), totalFlow);
+ }
+
+ @Test
+ public void testAssignTaskFromMinCostFlow() {
+ graph.solveMinCostFlow();
+ constructor.assignTaskFromMinCostFlow(
+ graph,
+ clientList,
+ taskIdList,
+ clientStateMap,
+ originalAssignedTaskNumber,
+ taskClientMap,
+ ClientState::assignActive,
+ ClientState::unassignActive,
+ ClientState::hasAssignedTask
+ );
+ assertValidAssignment(0, taskIds, emptySet(), clientStateMap, new StringBuilder());
+ if (constructorType.equals(BALANCE_SUBTOPOLOGY)) {
+ assertBalancedTasks(clientStateMap);
+ }
+ }
+
+ private void randomAssignTasksToClient(final List<TaskId> taskIdList, final SortedMap<UUID, ClientState> clientStateMap) {
+ int totalAssigned = 0;
+ for (final ClientState clientState : clientStateMap.values()) {
+ clientState.assignActive(taskIdList.get(totalAssigned++));
+ clientState.assignActive(taskIdList.get(totalAssigned++));
+ }
+ while (totalAssigned < taskIdList.size()) {
+ for (final ClientState clientState : clientStateMap.values()) {
+ if (AssignmentTestUtils.getRandom().nextInt(3) == 0) {
+ clientState.assignActive(taskIdList.get(totalAssigned));
+ totalAssigned++;
+ if (totalAssigned >= taskIdList.size()) {
+ break;
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java
index 58dd4313dbd..2c7d17cb5a9 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java
@@ -51,6 +51,8 @@ import static org.apache.kafka.streams.processor.internals.assignment.Assignment
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_5;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_6;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_7;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertBalancedTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertValidAssignment;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.clientTaskCount;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.configProps;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getClusterForAllTopics;
@@ -61,6 +63,7 @@ import static org.apache.kafka.streams.processor.internals.assignment.Assignment
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskChangelogMapForAllTasks;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMap;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMapForAllTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTasksForTopicGroup;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTopologyGroupTaskMap;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForChangelog;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForRandomChangelog;
@@ -105,6 +108,7 @@ import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
import org.apache.kafka.test.MockClientSupplier;
import org.apache.kafka.test.MockInternalTopicManager;
@@ -120,7 +124,7 @@ import org.mockito.Mockito;
public class RackAwareTaskAssignorTest {
private final MockTime time = new MockTime();
- private final StreamsConfig streamsConfig = new StreamsConfig(configProps(true));
+ private final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
private final MockClientSupplier mockClientSupplier = new MockClientSupplier();
private int trafficCost;
@@ -129,11 +133,16 @@ public class RackAwareTaskAssignorTest {
@Parameter
public boolean stateful;
- @Parameterized.Parameters(name = "stateful={0}")
+ @Parameter(1)
+ public String assignmentStrategy;
+
+ @Parameterized.Parameters(name = "stateful={0}, Strategy={1}")
public static Collection<Object[]> getParamStoreType() {
return asList(new Object[][] {
- {true},
- {false}
+ {true, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY},
+ {false, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC},
+ {true, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC},
+ {false, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY},
});
}
@@ -157,17 +166,17 @@ public class RackAwareTaskAssignorTest {
private AssignmentConfigs getRackAwareEnabledConfig() {
return new AssignorConfiguration(
- new StreamsConfig(configProps(true)).originals()).assignmentConfigs();
+ new StreamsConfig(configProps(assignmentStrategy)).originals()).assignmentConfigs();
}
private AssignmentConfigs getRackAwareEnabledConfigWithStandby(final int replicaNum) {
return new AssignorConfiguration(
- new StreamsConfig(configProps(true, replicaNum)).originals()).assignmentConfigs();
+ new StreamsConfig(configProps(assignmentStrategy, replicaNum)).originals()).assignmentConfigs();
}
private AssignmentConfigs getRackAwareDisabledConfig() {
return new AssignorConfiguration(
- new StreamsConfig(configProps(false)).originals()).assignmentConfigs();
+ new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE)).originals()).assignmentConfigs();
}
@Test
@@ -499,11 +508,15 @@ public class RackAwareTaskAssignorTest {
@Test
public void shouldOptimizeActiveTasks() {
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
+ mkEntry(new Subtopology(0, null), mkSet(TASK_0_0, TASK_0_1)),
+ mkEntry(new Subtopology(1, null), mkSet(TASK_1_0, TASK_1_1))
+ );
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForAllTopics(),
getTaskTopicPartitionMapForAllTasks(),
mkMap(),
- getTopologyGroupTaskMap(),
+ tasksForTopicGroup,
getProcessRacksForAllProcess(),
mockInternalTopicManager,
getRackAwareEnabledConfig(),
@@ -555,7 +568,7 @@ public class RackAwareTaskAssignorTest {
getRandomCluster(nodeSize, tpSize, partitionSize),
taskTopicPartitionMap,
mkMap(),
- getTopologyGroupTaskMap(),
+ getTasksForTopicGroup(tpSize, partitionSize),
getRandomProcessRacks(clientSize, nodeSize),
mockInternalTopicManager,
getRackAwareEnabledConfig(),
@@ -575,6 +588,10 @@ public class RackAwareTaskAssignorTest {
for (final Entry<UUID, ClientState> entry : clientStateMap.entrySet()) {
assertEquals((int) clientTaskCount.get(entry.getKey()), entry.getValue().activeTasks().size());
}
+
+ if (assignmentStrategy.equals(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY)) {
+ assertBalancedTasks(clientStateMap);
+ }
}
@Test
@@ -583,12 +600,13 @@ public class RackAwareTaskAssignorTest {
final int tpSize = 40;
final int partitionSize = 3;
final int clientSize = 30;
+ final int topicGroupSize = 20;
final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, false);
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getRandomCluster(nodeSize, tpSize, partitionSize),
taskTopicPartitionMap,
mkMap(),
- getTopologyGroupTaskMap(),
+ getTasksForTopicGroup(tpSize, partitionSize),
getRandomProcessRacks(clientSize, nodeSize),
mockInternalTopicManager,
getRackAwareEnabledConfig(),
@@ -610,22 +628,32 @@ public class RackAwareTaskAssignorTest {
assertEquals(0, originalCost);
final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, 0, 1);
- assertEquals(0, cost);
- // Make sure assignment doesn't change
- for (final Entry<TaskId, UUID> entry : taskClientMap.entrySet()) {
- final ClientState clientState = clientStateMap.get(entry.getValue());
- assertTrue(clientState.hasAssignedTask(entry.getKey()));
+ if (assignmentStrategy.equals(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC)) {
+ assertEquals(0, cost);
+
+ // Make sure assignment doesn't change
+ for (final Entry<TaskId, UUID> entry : taskClientMap.entrySet()) {
+ final ClientState clientState = clientStateMap.get(entry.getValue());
+ assertTrue(clientState.hasAssignedTask(entry.getKey()));
+ }
+ } else {
+ assertValidAssignment(0, taskIds, emptySet(), clientStateMap, new StringBuilder());
+ assertBalancedTasks(clientStateMap);
}
}
@Test
public void shouldOptimizeActiveTasksWithMoreClients() {
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
+ mkEntry(new Subtopology(0, null), mkSet(TASK_0_0)),
+ mkEntry(new Subtopology(1, null), mkSet(TASK_1_0))
+ );
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForAllTopics(),
getTaskTopicPartitionMapForAllTasks(),
mkMap(),
- getTopologyGroupTaskMap(),
+ tasksForTopicGroup,
getProcessRacksForAllProcess(),
mockInternalTopicManager,
getRackAwareEnabledConfig(),
@@ -667,11 +695,15 @@ public class RackAwareTaskAssignorTest {
@Test
public void shouldOptimizeActiveTasksWithMoreClientsWithMoreThanOneTask() {
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
+ mkEntry(new Subtopology(0, null), mkSet(TASK_0_0, TASK_0_1)),
+ mkEntry(new Subtopology(1, null), mkSet(TASK_1_0))
+ );
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForAllTopics(),
getTaskTopicPartitionMapForAllTasks(),
mkMap(),
- getTopologyGroupTaskMap(),
+ tasksForTopicGroup,
getProcessRacksForAllProcess(),
mockInternalTopicManager,
getRackAwareEnabledConfig(),
@@ -713,11 +745,15 @@ public class RackAwareTaskAssignorTest {
@Test
public void shouldBalanceAssignmentWithMoreCost() {
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
+ mkEntry(new Subtopology(0, null), mkSet(TASK_0_0, TASK_0_1)),
+ mkEntry(new Subtopology(1, null), mkSet(TASK_1_1))
+ );
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForAllTopics(),
getTaskTopicPartitionMapForAllTasks(),
mkMap(),
- getTopologyGroupTaskMap(),
+ tasksForTopicGroup,
getProcessRacksForAllProcess(),
mockInternalTopicManager,
getRackAwareEnabledConfig(),
@@ -748,10 +784,15 @@ public class RackAwareTaskAssignorTest {
final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
assertEquals(expectedCost, cost);
- // Even though assigning all tasks to UUID_2 will result in min cost, but it's not balanced
- // assignment. That's why TASK_0_1 is still assigned to UUID_5
- assertEquals(mkSet(TASK_0_0, TASK_1_1), clientState1.activeTasks());
- assertEquals(mkSet(TASK_0_1), clientState2.activeTasks());
+ if (stateful || assignmentStrategy.equals(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC)) {
+ // Even though assigning all tasks to UUID_2 will result in min cost, but it's not balanced
+ // assignment. That's why TASK_0_1 is still assigned to UUID_5
+ assertEquals(mkSet(TASK_0_0, TASK_1_1), clientState1.activeTasks());
+ assertEquals(mkSet(TASK_0_1), clientState2.activeTasks());
+ } else {
+ assertEquals(mkSet(TASK_0_0, TASK_0_1), clientState1.activeTasks());
+ assertEquals(mkSet(TASK_1_1), clientState2.activeTasks());
+ }
}
@Test
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java
index 19f865f8cf7..3db5dc3c803 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java
@@ -164,7 +164,7 @@ public class TaskAssignorConvergenceTest {
}
final MockTime time = new MockTime();
- final StreamsConfig streamsConfig = new StreamsConfig(configProps(true));
+ final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
final MockClientSupplier mockClientSupplier = new MockClientSupplier();
final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
time,