diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index a2b07913356db..89985ae2c7349 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -201,7 +201,7 @@ files="Murmur3.java"/> + files="(KStreamSlidingWindowAggregate|RackAwareTaskAssignor).java"/> .Edge otherEdge = (Graph.Edge) other; - return destination.equals(otherEdge.destination) && capacity == otherEdge.capacity - && cost == otherEdge.cost && residualFlow == otherEdge.residualFlow && flow == otherEdge.flow + return destination.equals(otherEdge.destination) + && capacity == otherEdge.capacity + && cost == otherEdge.cost + && residualFlow == otherEdge.residualFlow + && flow == otherEdge.flow && forwardEdge == otherEdge.forwardEdge; } @@ -84,8 +91,15 @@ public int hashCode() { @Override public String toString() { - return "{destination= " + destination + ", capacity=" + capacity + ", cost=" + cost - + ", residualFlow=" + residualFlow + ", flow=" + flow + ", forwardEdge=" + forwardEdge; + return "Edge {" + + "destination= " + destination + + ", capacity=" + capacity + + ", cost=" + cost + + ", residualFlow=" + residualFlow + + ", flow=" + flow + + ", forwardEdge=" + forwardEdge + + "}"; + } } @@ -106,12 +120,13 @@ public void addEdge(final V u, final V v, final int capacity, final int cost, fi addEdge(u, new Edge(v, capacity, cost, capacity - flow, flow)); } - public Set nodes() { + public SortedSet nodes() { return nodes; } - public Map edges(final V node) { - return adjList.get(node); + public SortedMap edges(final V node) { + final SortedMap edge = adjList.get(node); + return edge == null ? new TreeMap<>() : edge; } public boolean isResidualGraph() { @@ -126,12 +141,12 @@ public void setSinkNode(final V node) { sinkNode = node; } - public int totalCost() { - int totalCost = 0; + public long totalCost() { + long totalCost = 0; for (final Map.Entry> nodeEdges : adjList.entrySet()) { final SortedMap edges = nodeEdges.getValue(); - for (final Entry nodeEdge : edges.entrySet()) { - totalCost += nodeEdge.getValue().cost * nodeEdge.getValue().flow; + for (final Edge nodeEdge : edges.values()) { + totalCost += (long) nodeEdge.cost * nodeEdge.flow; } } return totalCost; 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 0452620f972b1..0b52cad482d64 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 @@ -16,12 +16,17 @@ */ package org.apache.kafka.streams.processor.internals.assignment; +import java.util.ArrayList; 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; +import java.util.SortedMap; +import java.util.SortedSet; import java.util.UUID; import org.apache.kafka.common.Cluster; import org.apache.kafka.common.Node; @@ -39,9 +44,11 @@ public class RackAwareTaskAssignor { private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class); + private static final int SOURCE_ID = -1; + private final Cluster fullMetadata; private final Map> partitionsForTask; - private final Map>> processRacks; + private final Map>> racksForProcess; private final AssignmentConfigs assignmentConfigs; private final Map> racksForPartition; private final InternalTopicManager internalTopicManager; @@ -49,18 +56,18 @@ public class RackAwareTaskAssignor { public RackAwareTaskAssignor(final Cluster fullMetadata, final Map> partitionsForTask, final Map> tasksForTopicGroup, - final Map>> processRacks, + final Map>> racksForProcess, final InternalTopicManager internalTopicManager, final AssignmentConfigs assignmentConfigs) { this.fullMetadata = fullMetadata; this.partitionsForTask = partitionsForTask; - this.processRacks = processRacks; + this.racksForProcess = racksForProcess; this.internalTopicManager = internalTopicManager; this.assignmentConfigs = assignmentConfigs; this.racksForPartition = new HashMap<>(); } - public synchronized boolean canEnableRackAwareAssignorForActiveTasks() { + public synchronized boolean canEnableRackAwareAssignor() { /* TODO: enable this after we add the config if (StreamsConfig.RACK_AWARE_ASSSIGNMENT_STRATEGY_NONE.equals(assignmentConfigs.rackAwareAssignmentStrategy)) { @@ -74,11 +81,7 @@ public synchronized boolean canEnableRackAwareAssignorForActiveTasks() { } return validateTopicPartitionRack(); - } - - public boolean canEnableRackAwareAssignorForStandbyTasks() { - // TODO - return false; + // TODO: add changelog topic, standby task validation } // Visible for testing. This method also checks if all TopicPartitions exist in cluster @@ -159,7 +162,7 @@ public boolean validateClientRack() { * 1. RackId exist for all clients * 2. Different consumerId for same process should have same rackId */ - for (final Map.Entry>> entry : processRacks.entrySet()) { + for (final Map.Entry>> entry : racksForProcess.entrySet()) { final UUID processId = entry.getKey(); KeyValue previousRackInfo = null; for (final Map.Entry> rackEntry : entry.getValue().entrySet()) { @@ -185,4 +188,213 @@ public boolean validateClientRack() { } return true; } + + private int getCost(final TaskId taskId, final UUID processId, final boolean inCurrentAssignment, final int trafficCost, final int nonOverlapCost) { + final Map> clientRacks = racksForProcess.get(processId); + if (clientRacks == null) { + throw new IllegalStateException("Client " + processId + " doesn't exist in processRacks"); + } + final Optional> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst(); + if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) { + throw new IllegalStateException("Client " + processId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first"); + } + + final String clientRack = clientRackOpt.get().get(); + final Set topicPartitions = partitionsForTask.get(taskId); + if (topicPartitions == null || topicPartitions.isEmpty()) { + throw new IllegalStateException("Task " + taskId + " has no TopicPartitions"); + } + + int cost = 0; + for (final TopicPartition tp : topicPartitions) { + final Set tpRacks = racksForPartition.get(tp); + if (tpRacks == null || tpRacks.isEmpty()) { + throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first"); + } + if (!tpRacks.contains(clientRack)) { + cost += trafficCost; + } + } + + if (!inCurrentAssignment) { + cost += nonOverlapCost; + } + + return cost; + } + + private static int getSinkID(final List clientList, final List taskIdList) { + return clientList.size() + taskIdList.size(); + } + + // For testing. canEnableRackAwareAssignor must be called first + long activeTasksCost(final SortedMap clientStates, final SortedSet activeTasks, final int trafficCost, final int nonOverlapCost) { + final List clientList = new ArrayList<>(clientStates.keySet()); + final List taskIdList = new ArrayList<>(activeTasks); + final Graph graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList, + clientStates, new HashMap<>(), new HashMap<>(), trafficCost, nonOverlapCost); + return graph.totalCost(); + } + + /** + * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first. + * {@code trafficCost} and {@code nonOverlapCost} balance cross rack traffic optimization and task movement. + * If we set {@code trafficCost} to a larger number, we are more likely to compute an assignment with less + * cross rack traffic. However, tasks may be shuffled a lot across clients. If we set {@code nonOverlapCost} + * to a larger number, we are more likely to compute an assignment with similar to input assignment. However, + * cross rack traffic can be higher. In extreme case, if we set {@code nonOverlapCost} to 0 and @{code trafficCost} + * to a positive value, the computed assignment will be minimum for cross rack traffic. If we set {@code trafficCost} to 0, + * and {@code nonOverlapCost} to a positive value, the computed assignment should be the same as input + * @param clientStates Client states + * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates + * @param trafficCost Cost of cross rack traffic for each TopicPartition + * @param nonOverlapCost Cost of assign a task to a different client + * @return Total cost after optimization + */ + public long optimizeActiveTasks(final SortedMap clientStates, + final SortedSet activeTasks, + final int trafficCost, + final int nonOverlapCost) { + if (activeTasks.isEmpty()) { + return 0; + } + + final List clientList = new ArrayList<>(clientStates.keySet()); + final List taskIdList = new ArrayList<>(activeTasks); + final Map taskClientMap = new HashMap<>(); + final Map originalAssignedTaskNumber = new HashMap<>(); + final Graph graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList, + clientStates, taskClientMap, originalAssignedTaskNumber, trafficCost, nonOverlapCost); + + graph.solveMinCostFlow(); + final long cost = graph.totalCost(); + + assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList, + clientStates, originalAssignedTaskNumber, taskClientMap); + + return cost; + } + + private Graph constructActiveTaskGraph(final SortedSet activeTasks, + final List clientList, + final List taskIdList, + final Map clientStates, + final Map taskClientMap, + final Map originalAssignedTaskNumber, + final int trafficCost, + final int nonOverlapCost) { + final Graph graph = new Graph<>(); + + for (final TaskId taskId : activeTasks) { + for (final Entry clientState : clientStates.entrySet()) { + if (clientState.getValue().hasAssignedTask(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 = taskIdList.size() + j; + final UUID processId = clientList.get(j); + + final int flow = clientStates.get(processId).hasAssignedTask(taskId) ? 1 : 0; + final int cost = getCost(taskId, processId, flow == 1, trafficCost, nonOverlapCost); + if (flow == 1) { + if (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"); + } + } + + final int sinkId = getSinkID(clientList, taskIdList); + for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) { + graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1); + } + + // 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 = taskIdList.size() + 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 void assignActiveTaskFromMinCostFlow(final Graph graph, + final SortedSet activeTasks, + final List clientList, + final List taskIdList, + final Map clientStates, + final Map originalAssignedTaskNumber, + final Map taskClientMap) { + int tasksAssigned = 0; + for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) { + final TaskId taskId = taskIdList.get(taskNodeId); + final Map.Edge> edges = graph.edges(taskNodeId); + for (final Graph.Edge edge : edges.values()) { + if (edge.flow > 0) { + tasksAssigned++; + final int clientIndex = edge.destination - taskIdList.size(); + 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; + } + + clientStates.get(originalProcessId).unassignActive(taskId); + clientStates.get(processId).assignActive(taskId); + } + } + } + + // Validate task assigned + if (tasksAssigned != activeTasks.size()) { + throw new IllegalStateException("Computed active task assignment number " + + tasksAssigned + " is different size " + activeTasks.size()); + } + + // Validate original assigned task number matches + final Map assignedTaskNumber = new HashMap<>(); + for (final TaskId taskId : activeTasks) { + for (final Entry clientState : clientStates.entrySet()) { + if (clientState.getValue().hasAssignedTask(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 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/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 8993372dd5aad..5f53e028a7ab5 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 @@ -22,6 +22,8 @@ import org.apache.kafka.clients.admin.AdminClient; import org.apache.kafka.clients.admin.ListOffsetsResult; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.internals.KafkaFutureImpl; import org.apache.kafka.streams.processor.TaskId; @@ -69,12 +71,38 @@ public final class AssignmentTestUtils { public static final UUID UUID_8 = uuidForInt(8); public static final UUID UUID_9 = uuidForInt(9); - public static final TopicPartition TP_0_0 = new TopicPartition("topic0", 0); - public static final TopicPartition TP_0_1 = new TopicPartition("topic0", 1); - public static final TopicPartition TP_0_2 = new TopicPartition("topic0", 2); - public static final TopicPartition TP_1_0 = new TopicPartition("topic1", 0); - public static final TopicPartition TP_1_1 = new TopicPartition("topic1", 1); - public static final TopicPartition TP_1_2 = new TopicPartition("topic1", 2); + public static final String RACK_0 = "rock0"; + public static final String RACK_1 = "rock1"; + public static final String RACK_2 = "rock2"; + public static final String RACK_3 = "rock3"; + public static final String RACK_4 = "rock4"; + + public static final Node NODE_0 = new Node(0, "node0", 1, RACK_0); + public static final Node NODE_1 = new Node(1, "node1", 1, RACK_1); + public static final Node NODE_2 = new Node(2, "node2", 1, RACK_2); + public static final Node NODE_3 = new Node(2, "node2", 1, RACK_3); + public static final Node NO_RACK_NODE = new Node(3, "node3", 1); + + public static final Node[] REPLICA_0 = new Node[] {NODE_0, NODE_1}; + public static final Node[] REPLICA_1 = new Node[] {NODE_1, NODE_2}; + public static final Node[] REPLICA_2 = new Node[] {NODE_0, NODE_2}; + public static final Node[] REPLICA_3 = new Node[] {NODE_1, NODE_3}; + + public static final String TP_0_NAME = "topic0"; + public static final String TP_1_NAME = "topic1"; + + public static final TopicPartition TP_0_0 = new TopicPartition(TP_0_NAME, 0); + public static final TopicPartition TP_0_1 = new TopicPartition(TP_0_NAME, 1); + public static final TopicPartition TP_0_2 = new TopicPartition(TP_0_NAME, 2); + public static final TopicPartition TP_1_0 = new TopicPartition(TP_1_NAME, 0); + public static final TopicPartition TP_1_1 = new TopicPartition(TP_1_NAME, 1); + public static final TopicPartition TP_1_2 = new TopicPartition(TP_1_NAME, 2); + + public static final PartitionInfo PI_0_0 = new PartitionInfo(TP_0_NAME, 0, NODE_0, REPLICA_0, REPLICA_0); + public static final PartitionInfo PI_0_1 = new PartitionInfo(TP_0_NAME, 1, NODE_1, REPLICA_1, REPLICA_1); + public static final PartitionInfo PI_1_0 = new PartitionInfo(TP_1_NAME, 0, NODE_2, REPLICA_2, REPLICA_2); + public static final PartitionInfo PI_1_1 = new PartitionInfo(TP_1_NAME, 1, NODE_3, REPLICA_3, REPLICA_3); + public static final PartitionInfo PI_1_2 = new PartitionInfo(TP_1_NAME, 2, NODE_0, REPLICA_0, REPLICA_0); public static final TaskId TASK_0_0 = new TaskId(0, 0); public static final TaskId TASK_0_1 = new TaskId(0, 1); 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 cf940f5a84257..7b5f2fb76ee99 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 @@ -20,7 +20,6 @@ import static org.hamcrest.Matchers.contains; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -48,19 +47,19 @@ public void setUp() { graph.addEdge(0, 3, 1, 1, 0); graph.addEdge(2, 1, 1, 1, 0); graph.addEdge(2, 3, 1, 2, 1); - graph.addEdge(4, 0, 1, 0, 1); - graph.addEdge(4, 2, 1, 0, 1); - graph.addEdge(1, 5, 1, 0, 1); - graph.addEdge(3, 5, 1, 0, 1); - graph.setSourceNode(4); - graph.setSinkNode(5); + graph.addEdge(-1, 0, 1, 0, 1); + graph.addEdge(-1, 2, 1, 0, 1); + graph.addEdge(1, 99, 1, 0, 1); + graph.addEdge(3, 99, 1, 0, 1); + graph.setSourceNode(-1); + graph.setSinkNode(99); } @Test public void testBasic() { final Set nodes = graph.nodes(); assertEquals(6, nodes.size()); - assertThat(nodes, contains(0, 1, 2, 3, 4, 5)); + assertThat(nodes, contains(-1, 0, 1, 2, 3, 99)); Map.Edge> edges = graph.edges(0); assertEquals(2, edges.size()); @@ -74,19 +73,19 @@ public void testBasic() { edges = graph.edges(1); assertEquals(1, edges.size()); - assertEquals(getEdge(5, 1, 0, 0, 1), edges.get(5)); + assertEquals(getEdge(99, 1, 0, 0, 1), edges.get(99)); edges = graph.edges(3); assertEquals(1, edges.size()); - assertEquals(getEdge(5, 1, 0, 0, 1), edges.get(5)); + assertEquals(getEdge(99, 1, 0, 0, 1), edges.get(99)); - edges = graph.edges(4); + edges = graph.edges(-1); assertEquals(2, edges.size()); assertEquals(getEdge(0, 1, 0, 0, 1), edges.get(0)); assertEquals(getEdge(2, 1, 0, 0, 1), edges.get(2)); - edges = graph.edges(5); - assertNull(edges); + edges = graph.edges(99); + assertTrue(edges.isEmpty()); assertFalse(graph.isResidualGraph()); } @@ -99,31 +98,31 @@ public void testResidualGraph() { final Set nodes = residualGraph.nodes(); assertEquals(6, nodes.size()); - assertThat(nodes, contains(0, 1, 2, 3, 4, 5)); + assertThat(nodes, contains(-1, 0, 1, 2, 3, 99)); Map.Edge> edges = residualGraph.edges(0); assertEquals(3, edges.size()); assertEquals(getEdge(1, 1, 3, 0, 1), edges.get(1)); assertEquals(getEdge(3, 1, 1, 1, 0), edges.get(3)); - assertEquals(getEdge(4, 1, 0, 1, 0, false), edges.get(4)); + assertEquals(getEdge(-1, 1, 0, 1, 0, false), edges.get(-1)); edges = residualGraph.edges(2); assertEquals(3, edges.size()); assertEquals(getEdge(1, 1, 1, 1, 0), edges.get(1)); assertEquals(getEdge(3, 1, 2, 0, 1), edges.get(3)); - assertEquals(getEdge(4, 1, 0, 1, 0, false), edges.get(4)); + assertEquals(getEdge(-1, 1, 0, 1, 0, false), edges.get(-1)); edges = residualGraph.edges(1); assertEquals(3, edges.size()); assertEquals(getEdge(0, 1, -3, 1, 0, false), edges.get(0)); assertEquals(getEdge(2, 1, -1, 0, 0, false), edges.get(2)); - assertEquals(getEdge(5, 1, 0, 0, 1), edges.get(5)); + assertEquals(getEdge(99, 1, 0, 0, 1), edges.get(99)); edges = residualGraph.edges(3); assertEquals(3, edges.size()); assertEquals(getEdge(0, 1, -1, 0, 0, false), edges.get(0)); assertEquals(getEdge(2, 1, -2, 1, 0, false), edges.get(2)); - assertEquals(getEdge(5, 1, 0, 0, 1), edges.get(5)); + assertEquals(getEdge(99, 1, 0, 0, 1), edges.get(99)); assertTrue(residualGraph.isResidualGraph()); } @@ -288,8 +287,8 @@ public void testMinCostFlow() { public void testMinCostDetectNodeNotInNegativeCycle() { final Graph graph1 = new Graph<>(); - graph1.addEdge(5, 0, 1, 0, 1); - graph1.addEdge(5, 1, 1, 0, 1); + graph1.addEdge(-1, 0, 1, 0, 1); + graph1.addEdge(-1, 1, 1, 0, 1); graph1.addEdge(0, 2, 1, 1, 0); graph1.addEdge(0, 3, 1, 1, 0); @@ -299,12 +298,12 @@ public void testMinCostDetectNodeNotInNegativeCycle() { graph1.addEdge(1, 3, 1, 10, 1); graph1.addEdge(1, 4, 1, 1, 0); - graph1.addEdge(2, 6, 0, 0, 0); - graph1.addEdge(3, 6, 1, 0, 1); - graph1.addEdge(4, 6, 1, 0, 1); + graph1.addEdge(2, 99, 0, 0, 0); + graph1.addEdge(3, 99, 1, 0, 1); + graph1.addEdge(4, 99, 1, 0, 1); - graph1.setSourceNode(5); - graph1.setSinkNode(6); + graph1.setSourceNode(-1); + graph1.setSinkNode(99); assertEquals(20, graph1.totalCost()); @@ -313,7 +312,7 @@ public void testMinCostDetectNodeNotInNegativeCycle() { graph1.solveMinCostFlow(); assertEquals(2, graph1.totalCost()); - Map.Edge> edges = graph1.edges(5); + Map.Edge> edges = graph1.edges(-1); assertEquals(getEdge(0, 1, 0, 0, 1), edges.get(0)); assertEquals(getEdge(1, 1, 0, 0, 1), edges.get(1)); @@ -328,13 +327,13 @@ public void testMinCostDetectNodeNotInNegativeCycle() { assertEquals(getEdge(4, 1, 1, 0, 1), edges.get(4)); edges = graph1.edges(2); - assertEquals(getEdge(6, 0, 0, 0, 0), edges.get(6)); + assertEquals(getEdge(99, 0, 0, 0, 0), edges.get(99)); edges = graph1.edges(3); - assertEquals(getEdge(6, 1, 0, 0, 1), edges.get(6)); + assertEquals(getEdge(99, 1, 0, 0, 1), edges.get(99)); edges = graph1.edges(4); - assertEquals(getEdge(6, 1, 0, 0, 1), edges.get(6)); + assertEquals(getEdge(99, 1, 0, 0, 1), edges.get(99)); } @Test 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 7117d9481222a..554e3461d03d0 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 @@ -16,20 +16,75 @@ */ package org.apache.kafka.streams.processor.internals.assignment; +import static java.util.Arrays.asList; +import static java.util.Collections.emptyMap; +import static java.util.Collections.emptySet; +import static org.apache.kafka.common.utils.Utils.mkEntry; +import static org.apache.kafka.common.utils.Utils.mkMap; +import static org.apache.kafka.common.utils.Utils.mkSet; +import static org.apache.kafka.common.utils.Utils.mkSortedSet; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_CLIENT_TAGS; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_2; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_3; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NO_RACK_NODE; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PI_0_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PI_0_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PI_1_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PI_1_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PI_1_2; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_2; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_3; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_4; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.REPLICA_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.SUBTOPOLOGY_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_2; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_0_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_0_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_0_NAME; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_1_0; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_1_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_1_2; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_1; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_2; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_3; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_4; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_5; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.uuidForInt; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.lessThanOrEqualTo; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.spy; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.Iterator; +import java.util.List; import java.util.Map; +import java.util.Map.Entry; import java.util.Optional; import java.util.Set; +import java.util.SortedMap; +import java.util.SortedSet; +import java.util.TreeMap; import java.util.UUID; +import java.util.stream.Collectors; import org.apache.kafka.common.Cluster; import org.apache.kafka.common.Node; import org.apache.kafka.common.PartitionInfo; @@ -43,37 +98,36 @@ import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology; import org.apache.kafka.test.MockClientSupplier; import org.apache.kafka.test.MockInternalTopicManager; +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 RackAwareTaskAssignorTest { - private final static String TOPIC0 = "topic0"; - private final static String TOPIC1 = "topic1"; private static final String USER_END_POINT = "localhost:8080"; private static final String APPLICATION_ID = "stream-partition-assignor-test"; - private final Node node0 = new Node(0, "node0", 1, "rack1"); - private final Node node1 = new Node(1, "node1", 1, "rack2"); - private final Node node2 = new Node(2, "node2", 1, "rack3"); - private final Node noRackNode = new Node(3, "node3", 1); - private final Node[] replicas1 = new Node[] {node0, node1}; - private final Node[] replicas2 = new Node[] {node1, node2}; - - private final PartitionInfo partitionInfo00 = new PartitionInfo(TOPIC0, 0, node0, replicas1, replicas1); - private final PartitionInfo partitionInfo01 = new PartitionInfo(TOPIC0, 1, node0, replicas2, replicas2); - - private final TopicPartition partitionWithoutInfo00 = new TopicPartition(TOPIC0, 0); - private final TopicPartition partitionWithoutInfo01 = new TopicPartition(TOPIC0, 1); - private final TopicPartition partitionWithoutInfo10 = new TopicPartition(TOPIC1, 0); - - private final UUID process0UUID = UUID.randomUUID(); - private final UUID process1UUID = UUID.randomUUID(); - - private final Subtopology subtopology1 = new Subtopology(1, "topology1"); - private final MockTime time = new MockTime(); private final StreamsConfig streamsConfig = new StreamsConfig(configProps()); private final MockClientSupplier mockClientSupplier = new MockClientSupplier(); + private int trafficCost; + private int nonOverlapCost; + + @Parameter + public boolean stateful; + + @Parameterized.Parameters(name = "stateful={0}") + public static Collection getParamStoreType() { + return asList(new Object[][] { + {true}, + {false} + }); + } + private final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager( time, streamsConfig, @@ -81,6 +135,17 @@ public class RackAwareTaskAssignorTest { false ); + @Before + public void setUp() { + if (stateful) { + trafficCost = 10; + nonOverlapCost = 1; + } else { + trafficCost = 1; + nonOverlapCost = 0; + } + } + private Map configProps() { final Map configurationMap = new HashMap<>(); configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); @@ -99,10 +164,10 @@ private Map configProps() { } @Test - public void disableActiveSinceMissingClusterInfo() { + public void shouldDisableActiveWhenMissingClusterInfo() { final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterForTopic0(), - getTaskTopicPartitionMapForTask1(true), + getTaskTopicPartitionMapForTask0(true), getTopologyGroupTaskMap(), getProcessRacksForProcess0(), mockInternalTopicManager, @@ -110,16 +175,16 @@ public void disableActiveSinceMissingClusterInfo() { ); // False since partitionWithoutInfo10 is missing in cluster metadata - assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertFalse(assignor.canEnableRackAwareAssignor()); assertFalse(assignor.populateTopicsToDiscribe(new HashSet<>())); assertTrue(assignor.validateClientRack()); } @Test - public void disableActiveSinceRackMissingInNode() { + public void shouldDisableActiveWhenRackMissingInNode() { final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterWithPartitionMissingRack(), - getTaskTopicPartitionMapForTask1(), + getTaskTopicPartitionMapForTask0(), getTopologyGroupTaskMap(), getProcessRacksForProcess0(), mockInternalTopicManager, @@ -129,14 +194,14 @@ public void disableActiveSinceRackMissingInNode() { assertTrue(assignor.validateClientRack()); assertFalse(assignor.populateTopicsToDiscribe(new HashSet<>())); // False since nodeMissingRack has one node which doesn't have rack - assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertFalse(assignor.canEnableRackAwareAssignor()); } @Test - public void disableActiveSinceRackMissingInClient() { + public void shouldDisableActiveWhenRackMissingInClient() { final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterForTopic0(), - getTaskTopicPartitionMapForTask1(), + getTaskTopicPartitionMapForTask0(), getTopologyGroupTaskMap(), getProcessRacksForProcess0(true), mockInternalTopicManager, @@ -145,21 +210,21 @@ public void disableActiveSinceRackMissingInClient() { // False since process1 doesn't have rackId assertFalse(assignor.validateClientRack()); - assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertFalse(assignor.canEnableRackAwareAssignor()); } @Test - public void disableActiveSinceRackDiffersInSameProcess() { + public void shouldDisableActiveWhenRackDiffersInSameProcess() { final Map>> processRacks = new HashMap<>(); // Different consumers in same process have different rack ID. This shouldn't happen. // If happens, there's a bug somewhere - processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer1", Optional.of("rack1")); - processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer2", Optional.of("rack2")); + processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer1", Optional.of("rack1")); + processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer2", Optional.of("rack2")); final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterForTopic0(), - getTaskTopicPartitionMapForTask1(), + getTaskTopicPartitionMapForTask0(), getTopologyGroupTaskMap(), processRacks, mockInternalTopicManager, @@ -167,14 +232,14 @@ public void disableActiveSinceRackDiffersInSameProcess() { ); assertFalse(assignor.validateClientRack()); - assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertFalse(assignor.canEnableRackAwareAssignor()); } @Test - public void enableRackAwareAssignorForActiveWithoutDescribingTopics() { + public void shouldEnableRackAwareAssignorForActiveWithoutDescribingTopics() { final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterForTopic0(), - getTaskTopicPartitionMapForTask1(), + getTaskTopicPartitionMapForTask0(), getTopologyGroupTaskMap(), getProcessRacksForProcess0(), mockInternalTopicManager, @@ -182,85 +247,535 @@ public void enableRackAwareAssignorForActiveWithoutDescribingTopics() { ); // partitionWithoutInfo00 has rackInfo in cluster metadata - assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertTrue(assignor.canEnableRackAwareAssignor()); } @Test - public void enableRackAwareAssignorForActiveWithDescribingTopics() { + public void shouldEnableRackAwareAssignorForActiveWithDescribingTopics() { final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager); doReturn( Collections.singletonMap( - TOPIC0, + TP_0_NAME, Collections.singletonList( - new TopicPartitionInfo(0, node0, Arrays.asList(replicas1), Collections.emptyList()) + new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList()) ) ) - ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0)); + ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME)); final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterWithNoNode(), - getTaskTopicPartitionMapForTask1(), + getTaskTopicPartitionMapForTask0(), getTopologyGroupTaskMap(), getProcessRacksForProcess0(), spyTopicManager, new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() ); - assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertTrue(assignor.canEnableRackAwareAssignor()); } @Test - public void disableRackAwareAssignorForActiveWithDescribingTopicsFailure() { + public void shouldDisableRackAwareAssignorForActiveWithDescribingTopicsFailure() { final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager); - doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0)); + doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton( + TP_0_NAME)); final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( getClusterWithNoNode(), - getTaskTopicPartitionMapForTask1(), + getTaskTopicPartitionMapForTask0(), getTopologyGroupTaskMap(), getProcessRacksForProcess0(), spyTopicManager, new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() ); - assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks()); + assertFalse(assignor.canEnableRackAwareAssignor()); assertTrue(assignor.populateTopicsToDiscribe(new HashSet<>())); } + @Test + public void shouldOptimizeEmptyActiveTasks() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1)); + + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_1, clientState0) + )); + final SortedSet taskIds = mkSortedSet(); + + assertTrue(assignor.canEnableRackAwareAssignor()); + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(0, originalCost); + + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(0, cost); + + assertEquals(mkSet(TASK_0_1, TASK_1_1), clientState0.activeTasks()); + } + + @Test + public void shouldOptimizeActiveTasks() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1)); + clientState1.assignActive(TASK_1_0); + clientState2.assignActive(TASK_0_0); + + // task_0_0 has same rack as UUID_1 + // task_0_1 has same rack as UUID_2 and UUID_3 + // task_1_0 has same rack as UUID_1 and UUID_3 + // task_1_1 has same rack as UUID_2 + // Optimal assignment is UUID_1: {0_0, 1_0}, UUID_2: {1_1}, UUID_3: {0_1} which result in no cross rack traffic + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_1, clientState0), + mkEntry(UUID_2, clientState1), + mkEntry(UUID_3, clientState2) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1); + + assertTrue(assignor.canEnableRackAwareAssignor()); + int expected = stateful ? 40 : 4; + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(expected, originalCost); + + expected = stateful ? 4 : 0; + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(expected, cost); + + assertEquals(mkSet(TASK_0_0, TASK_1_0), clientState0.activeTasks()); + assertEquals(mkSet(TASK_1_1), clientState1.activeTasks()); + assertEquals(mkSet(TASK_0_1), clientState2.activeTasks()); + } + + @Test + public void shouldOptimizeRandom() { + final int nodeSize = 30; + final int tpSize = 40; + final int clientSize = 30; + final SortedMap> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize); + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getRandomCluster(nodeSize, tpSize), + taskTopicPartitionMap, + getTopologyGroupTaskMap(), + getRandomProcessRacks(clientSize, nodeSize), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final SortedMap clientStateMap = getRandomClientState(clientSize, tpSize); + final SortedSet taskIds = (SortedSet) taskTopicPartitionMap.keySet(); + + final Map clientTaskCount = clientStateMap.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().activeTasks().size())); + + assertTrue(assignor.canEnableRackAwareAssignor()); + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, trafficCost, nonOverlapCost); + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertThat(cost, lessThanOrEqualTo(originalCost)); + + for (final Entry entry : clientStateMap.entrySet()) { + assertEquals((int) clientTaskCount.get(entry.getKey()), entry.getValue().activeTasks().size()); + } + } + + @Test + public void shouldMaintainOriginalAssignment() { + final int nodeSize = 20; + final int tpSize = 40; + final int clientSize = 30; + final SortedMap> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize); + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getRandomCluster(nodeSize, tpSize), + taskTopicPartitionMap, + getTopologyGroupTaskMap(), + getRandomProcessRacks(clientSize, nodeSize), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final SortedMap clientStateMap = getRandomClientState(clientSize, tpSize); + final SortedSet taskIds = (SortedSet) taskTopicPartitionMap.keySet(); + + final Map taskClientMap = new HashMap<>(); + for (final Entry entry : clientStateMap.entrySet()) { + entry.getValue().activeTasks().forEach(t -> taskClientMap.put(t, entry.getKey())); + } + assertEquals(taskIds.size(), taskClientMap.size()); + + // Because trafficCost is 0, original assignment should be maintained + assertTrue(assignor.canEnableRackAwareAssignor()); + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, 0, 1); + assertEquals(0, originalCost); + + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, 0, 1); + assertEquals(0, cost); + + // Make sure assignment doesn't change + for (final Entry entry : taskClientMap.entrySet()) { + final ClientState clientState = clientStateMap.get(entry.getValue()); + assertTrue(clientState.hasAssignedTask(entry.getKey())); + } + } + + @Test + public void shouldOptimizeActiveTasksWithMoreClients() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState1.assignActive(TASK_1_0); + clientState2.assignActive(TASK_0_0); + + // task_0_0 has same rack as UUID_1 and UUID_2 + // task_1_0 has same rack as UUID_1 and UUID_3 + // Optimal assignment is UUID_1: {}, UUID_2: {0_0}, UUID_3: {1_0} which result in no cross rack traffic + // and keeps UUID_1 empty since it was originally empty + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_1, clientState0), + mkEntry(UUID_2, clientState1), + mkEntry(UUID_3, clientState2) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_1_0); + + assertTrue(assignor.canEnableRackAwareAssignor()); + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, trafficCost, nonOverlapCost); + int expected = stateful ? 20 : 2; + assertEquals(expected, originalCost); + + expected = stateful ? 2 : 0; + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(expected, cost); + + // UUID_1 remains empty + assertEquals(mkSet(), clientState0.activeTasks()); + assertEquals(mkSet(TASK_0_0), clientState1.activeTasks()); + assertEquals(mkSet(TASK_1_0), clientState2.activeTasks()); + } + + @Test + public void shouldOptimizeActiveTasksWithMoreClientsWithMoreThanOneTask() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState1.assignActiveTasks(mkSet(TASK_0_1, TASK_1_0)); + clientState2.assignActive(TASK_0_0); + + // task_0_0 has same rack as UUID_1 and UUID_2 + // task_0_1 has same rack as UUID_2 and UUID_3 + // task_1_0 has same rack as UUID_1 and UUID_3 + // Optimal assignment is UUID_1: {}, UUID_2: {0_0, 0_1}, UUID_3: {1_0} which result in no cross rack traffic + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_1, clientState0), + mkEntry(UUID_2, clientState1), + mkEntry(UUID_3, clientState2) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0); + + assertTrue(assignor.canEnableRackAwareAssignor()); + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, trafficCost, nonOverlapCost); + int expected = stateful ? 20 : 2; + assertEquals(expected, originalCost); + + expected = stateful ? 2 : 0; + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(expected, cost); + + // Because original assignment is not balanced (3 tasks but client 0 has no task), we maintain it + assertEquals(mkSet(), clientState0.activeTasks()); + assertEquals(mkSet(TASK_0_0, TASK_0_1), clientState1.activeTasks()); + assertEquals(mkSet(TASK_1_0), clientState2.activeTasks()); + } + + @Test + public void shouldBalanceAssignmentWithMoreCost() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState0.assignActiveTasks(mkSet(TASK_0_0, TASK_1_1)); + clientState1.assignActive(TASK_0_1); + + // task_0_0 has same rack as UUID_2 + // task_0_1 has same rack as UUID_2 + // task_1_1 has same rack as UUID_2 + // UUID_5 is not in same rack as any task + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_2, clientState0), + mkEntry(UUID_5, clientState1) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_1); + + assertTrue(assignor.canEnableRackAwareAssignor()); + final int expectedCost = stateful ? 10 : 1; + final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, trafficCost, nonOverlapCost); + assertEquals(expectedCost, originalCost); + + final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, 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), clientState0.activeTasks()); + assertEquals(mkSet(TASK_0_1), clientState1.activeTasks()); + } + + @Test + public void shouldThrowIfMissingCallcanEnableRackAwareAssignor() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState0.assignActiveTasks(mkSet(TASK_0_0, TASK_1_1)); + clientState1.assignActive(TASK_0_1); + + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_2, clientState0), + mkEntry(UUID_5, clientState1) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_1); + final Exception exception = assertThrows(IllegalStateException.class, + () -> assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost)); + Assertions.assertEquals("TopicPartition topic0-0 has no rack information. Maybe forgot to call " + + "canEnableRackAwareAssignor first", exception.getMessage()); + } + + @Test + public void shouldThrowIfTaskInMultipleClients() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState0.assignActiveTasks(mkSet(TASK_0_0, TASK_1_1)); + clientState1.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1)); + + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_2, clientState0), + mkEntry(UUID_5, clientState1) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_1); + assertTrue(assignor.canEnableRackAwareAssignor()); + final Exception exception = assertThrows(IllegalArgumentException.class, + () -> assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost)); + Assertions.assertEquals( + "Task 1_1 assigned to multiple clients 00000000-0000-0000-0000-000000000005, " + + "00000000-0000-0000-0000-000000000002", exception.getMessage()); + } + + @Test + public void shouldThrowIfTaskMissingInClients() { + final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor( + getClusterForTopic0And1(), + getTaskTopicPartitionMapForAllTasks(), + getTopologyGroupTaskMap(), + getProcessRacksForAllProcess(), + mockInternalTopicManager, + new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs() + ); + + final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + + clientState0.assignActiveTasks(mkSet(TASK_0_0, TASK_1_1)); + clientState1.assignActive(TASK_0_1); + + final SortedMap clientStateMap = new TreeMap<>(mkMap( + mkEntry(UUID_2, clientState0), + mkEntry(UUID_5, clientState1) + )); + final SortedSet taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1); + assertTrue(assignor.canEnableRackAwareAssignor()); + final Exception exception = assertThrows(IllegalArgumentException.class, + () -> assignor.optimizeActiveTasks(clientStateMap, taskIds, trafficCost, nonOverlapCost)); + Assertions.assertEquals( + "Task 1_0 not assigned to any client", exception.getMessage()); + } + + private Cluster getRandomCluster(final int nodeSize, final int tpSize) { + final List nodeList = new ArrayList<>(nodeSize); + for (int i = 0; i < nodeSize; i++) { + nodeList.add(new Node(i, "node" + i, 1, "rack" + i)); + } + Collections.shuffle(nodeList); + final Set partitionInfoSet = new HashSet<>(); + for (int i = 0; i < tpSize; i++) { + final Node firstNode = nodeList.get(i % nodeSize); + final Node secondNode = nodeList.get((i + 1) % nodeSize); + final Node[] replica = new Node[] {firstNode, secondNode}; + partitionInfoSet.add(new PartitionInfo("topic" + i, 0, firstNode, replica, replica)); + } + + return new Cluster( + "cluster", + new HashSet<>(nodeList), + partitionInfoSet, + Collections.emptySet(), + Collections.emptySet() + ); + } + + private Map>> getRandomProcessRacks(final int clientSize, final int nodeSize) { + final List racks = new ArrayList<>(nodeSize); + for (int i = 0; i < nodeSize; i++) { + racks.add("rack" + i); + } + Collections.shuffle(racks); + final Map>> processRacks = new HashMap<>(); + for (int i = 0; i < clientSize; i++) { + final String rack = racks.get(i % nodeSize); + processRacks.put(uuidForInt(i), mkMap(mkEntry("1", Optional.of(rack)))); + } + return processRacks; + } + + private SortedMap> getTaskTopicPartitionMap(final int tpSize) { + final SortedMap> taskTopicPartitionMap = new TreeMap<>(); + for (int i = 0; i < tpSize; i++) { + taskTopicPartitionMap.put(new TaskId(i, 0), mkSet(new TopicPartition("topic" + i, 0))); + } + return taskTopicPartitionMap; + } + + private SortedMap getRandomClientState(final int clientSize, final int tpSize) { + final SortedMap clientStates = new TreeMap<>(); + final List taskIds = new ArrayList<>(tpSize); + for (int i = 0; i < tpSize; i++) { + taskIds.add(new TaskId(i, 0)); + } + Collections.shuffle(taskIds); + for (int i = 0; i < clientSize; i++) { + final ClientState clientState = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1); + clientStates.put(uuidForInt(i), clientState); + } + Iterator> iterator = clientStates.entrySet().iterator(); + for (final TaskId taskId : taskIds) { + if (iterator.hasNext()) { + iterator.next().getValue().assignActive(taskId); + } else { + iterator = clientStates.entrySet().iterator(); + iterator.next().getValue().assignActive(taskId); + } + } + return clientStates; + } + + private Cluster getClusterForTopic0And1() { + return new Cluster( + "cluster", + mkSet(NODE_0, NODE_1, NODE_2, NODE_3), + mkSet(PI_0_0, PI_0_1, PI_1_0, PI_1_1, PI_1_2), + Collections.emptySet(), + Collections.emptySet() + ); + } + private Cluster getClusterForTopic0() { return new Cluster( "cluster", - new HashSet<>(Arrays.asList(node0, node1, node2)), - new HashSet<>(Arrays.asList(partitionInfo00, partitionInfo01)), + mkSet(NODE_0, NODE_1, NODE_2), + mkSet(PI_0_0, PI_0_1), Collections.emptySet(), Collections.emptySet() ); } private Cluster getClusterWithPartitionMissingRack() { - final Node[] nodeMissingRack = new Node[]{node0, noRackNode}; - final PartitionInfo partitionInfoMissingNode = new PartitionInfo(TOPIC0, 0, node0, nodeMissingRack, nodeMissingRack); + final Node[] nodeMissingRack = new Node[]{NODE_0, NO_RACK_NODE}; + final PartitionInfo partitionInfoMissingNode = new PartitionInfo(TP_0_NAME, 0, NODE_0, nodeMissingRack, nodeMissingRack); return new Cluster( "cluster", - new HashSet<>(Arrays.asList(node0, node1, node2)), - new HashSet<>(Arrays.asList(partitionInfoMissingNode, partitionInfo01)), + mkSet(NODE_0, NODE_1, NODE_2), + mkSet(partitionInfoMissingNode, PI_0_1), Collections.emptySet(), Collections.emptySet() ); } private Cluster getClusterWithNoNode() { - final PartitionInfo noNodeInfo = new PartitionInfo(TOPIC0, 0, null, new Node[0], new Node[0]); + final PartitionInfo noNodeInfo = new PartitionInfo(TP_0_NAME, 0, null, new Node[0], new Node[0]); return new Cluster( "cluster", - new HashSet<>(Arrays.asList(node0, node1, node2, Node.noNode())), // mockClientSupplier.setCluster requires noNode + mkSet(NODE_0, NODE_1, NODE_2, Node.noNode()), // mockClientSupplier.setCluster requires noNode Collections.singleton(noNodeInfo), Collections.emptySet(), Collections.emptySet() ); } + private Map>> getProcessRacksForAllProcess() { + return mkMap( + mkEntry(UUID_1, mkMap(mkEntry("1", Optional.of(RACK_0)))), + mkEntry(UUID_2, mkMap(mkEntry("1", Optional.of(RACK_1)))), + mkEntry(UUID_3, mkMap(mkEntry("1", Optional.of(RACK_2)))), + mkEntry(UUID_4, mkMap(mkEntry("1", Optional.of(RACK_3)))), + mkEntry(UUID_5, mkMap(mkEntry("1", Optional.of(RACK_4)))) + ); + } + private Map>> getProcessRacksForProcess0() { return getProcessRacksForProcess0(false); } @@ -268,24 +783,34 @@ private Map>> getProcessRacksForProcess0() { private Map>> getProcessRacksForProcess0(final boolean missingRack) { final Map>> processRacks = new HashMap<>(); final Optional rack = missingRack ? Optional.empty() : Optional.of("rack1"); - processRacks.put(process0UUID, Collections.singletonMap("consumer1", rack)); + processRacks.put(UUID_1, Collections.singletonMap("consumer1", rack)); return processRacks; } - private Map> getTaskTopicPartitionMapForTask1() { - return getTaskTopicPartitionMapForTask1(false); + private Map> getTaskTopicPartitionMapForTask0() { + return getTaskTopicPartitionMapForTask0(false); } - private Map> getTaskTopicPartitionMapForTask1(final boolean extraTopic) { + private Map> getTaskTopicPartitionMapForTask0(final boolean extraTopic) { final Set topicPartitions = new HashSet<>(); - topicPartitions.add(partitionWithoutInfo00); + topicPartitions.add(TP_0_0); if (extraTopic) { - topicPartitions.add(partitionWithoutInfo10); + topicPartitions.add(TP_1_0); } - return Collections.singletonMap(new TaskId(1, 1), topicPartitions); + return Collections.singletonMap(TASK_0_0, topicPartitions); + } + + private Map> getTaskTopicPartitionMapForAllTasks() { + return mkMap( + mkEntry(TASK_0_0, mkSet(TP_0_0)), + mkEntry(TASK_0_1, mkSet(TP_0_1)), + mkEntry(TASK_1_0, mkSet(TP_1_0)), + mkEntry(TASK_1_1, mkSet(TP_1_1)), + mkEntry(TASK_1_2, mkSet(TP_1_2)) + ); } private Map> getTopologyGroupTaskMap() { - return Collections.singletonMap(subtopology1, Collections.singleton(new TaskId(1, 1))); + return Collections.singletonMap(SUBTOPOLOGY_0, Collections.singleton(new TaskId(1, 1))); } }