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/07/26 23:02:58 UTC
[kafka] branch trunk updated: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks (#14030)
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 ed44bcd71b3 KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks (#14030)
ed44bcd71b3 is described below
commit ed44bcd71b3b9926c474033882eaa6c1cf35cfa4
Author: Hao Li <11...@users.noreply.github.com>
AuthorDate: Wed Jul 26 16:02:52 2023 -0700
KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks (#14030)
Part of KIP-925.
Reviewers: Matthias J. Sax <ma...@confluent.io>
---
checkstyle/suppressions.xml | 2 +-
.../processor/internals/assignment/Graph.java | 41 +-
.../assignment/RackAwareTaskAssignor.java | 232 +++++++-
.../internals/assignment/AssignmentTestUtils.java | 40 +-
.../processor/internals/assignment/GraphTest.java | 57 +-
.../assignment/RackAwareTaskAssignorTest.java | 653 +++++++++++++++++++--
6 files changed, 902 insertions(+), 123 deletions(-)
diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 547a8e4b24a..0c40bcf446b 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -201,7 +201,7 @@
files="Murmur3.java"/>
<suppress checks="(NPathComplexity|CyclomaticComplexity)"
- files="KStreamSlidingWindowAggregate.java"/>
+ files="(KStreamSlidingWindowAggregate|RackAwareTaskAssignor).java"/>
<!-- suppress FinalLocalVariable outside of the streams package. -->
<suppress checks="FinalLocalVariable"
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 baf6dec5343..8bb31f44def 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
@@ -42,8 +42,12 @@ public class Graph<V extends Comparable<V>> {
this(destination, capacity, cost, residualFlow, flow, true);
}
- public Edge(final V destination, final int capacity, final int cost, final int residualFlow, final int flow,
- final boolean forwardEdge) {
+ public Edge(final V destination,
+ final int capacity,
+ final int cost,
+ final int residualFlow,
+ final int flow,
+ final boolean forwardEdge) {
Objects.requireNonNull(destination);
if (capacity < 0) {
throw new IllegalArgumentException("Edge capacity cannot be negative");
@@ -72,8 +76,11 @@ public class Graph<V extends Comparable<V>> {
final Graph<?>.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 class Graph<V extends Comparable<V>> {
@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 class Graph<V extends Comparable<V>> {
addEdge(u, new Edge(v, capacity, cost, capacity - flow, flow));
}
- public Set<V> nodes() {
+ public SortedSet<V> nodes() {
return nodes;
}
- public Map<V, Edge> edges(final V node) {
- return adjList.get(node);
+ public SortedMap<V, Edge> edges(final V node) {
+ final SortedMap<V, Edge> edge = adjList.get(node);
+ return edge == null ? new TreeMap<>() : edge;
}
public boolean isResidualGraph() {
@@ -126,12 +141,12 @@ public class Graph<V extends Comparable<V>> {
sinkNode = node;
}
- public int totalCost() {
- int totalCost = 0;
+ public long totalCost() {
+ long totalCost = 0;
for (final Map.Entry<V, SortedMap<V, Edge>> nodeEdges : adjList.entrySet()) {
final SortedMap<V, Edge> edges = nodeEdges.getValue();
- for (final Entry<V, Edge> 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 0452620f972..0b52cad482d 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 @@ import org.slf4j.LoggerFactory;
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<TaskId, Set<TopicPartition>> partitionsForTask;
- private final Map<UUID, Map<String, Optional<String>>> processRacks;
+ private final Map<UUID, Map<String, Optional<String>>> racksForProcess;
private final AssignmentConfigs assignmentConfigs;
private final Map<TopicPartition, Set<String>> racksForPartition;
private final InternalTopicManager internalTopicManager;
@@ -49,18 +56,18 @@ public class RackAwareTaskAssignor {
public RackAwareTaskAssignor(final Cluster fullMetadata,
final Map<TaskId, Set<TopicPartition>> partitionsForTask,
final Map<Subtopology, Set<TaskId>> tasksForTopicGroup,
- final Map<UUID, Map<String, Optional<String>>> processRacks,
+ final Map<UUID, Map<String, Optional<String>>> 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 class RackAwareTaskAssignor {
}
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 class RackAwareTaskAssignor {
* 1. RackId exist for all clients
* 2. Different consumerId for same process should have same rackId
*/
- for (final Map.Entry<UUID, Map<String, Optional<String>>> entry : processRacks.entrySet()) {
+ for (final Map.Entry<UUID, Map<String, Optional<String>>> entry : racksForProcess.entrySet()) {
final UUID processId = entry.getKey();
KeyValue<String, String> previousRackInfo = null;
for (final Map.Entry<String, Optional<String>> rackEntry : entry.getValue().entrySet()) {
@@ -185,4 +188,213 @@ public class RackAwareTaskAssignor {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID processId, final boolean inCurrentAssignment, final int trafficCost, final int nonOverlapCost) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(processId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + processId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> 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<TopicPartition> 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<String> 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<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final int trafficCost, final int nonOverlapCost) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Graph<Integer> 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<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final int trafficCost,
+ final int nonOverlapCost) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ 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 = 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<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final int trafficCost,
+ final int nonOverlapCost) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> 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<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final Map<TaskId, UUID> taskClientMap) {
+ int tasksAssigned = 0;
+ 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 = 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<UUID, Integer> assignedTaskNumber = new HashMap<>();
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> 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<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/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 8993372dd5a..5f53e028a7a 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 java.util.List;
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 cf940f5a842..7b5f2fb76ee 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.MatcherAssert.assertThat;
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 class GraphTest {
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<Integer> 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<Integer, Graph<Integer>.Edge> edges = graph.edges(0);
assertEquals(2, edges.size());
@@ -74,19 +73,19 @@ public class GraphTest {
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 class GraphTest {
final Set<Integer> 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<Integer, Graph<Integer>.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 class GraphTest {
public void testMinCostDetectNodeNotInNegativeCycle() {
final Graph<Integer> 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 class GraphTest {
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 class GraphTest {
graph1.solveMinCostFlow();
assertEquals(2, graph1.totalCost());
- Map<Integer, Graph<Integer>.Edge> edges = graph1.edges(5);
+ Map<Integer, Graph<Integer>.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 class GraphTest {
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 7117d948122..554e3461d03 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.TaskId;
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<Object[]> 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<String, Object> configProps() {
final Map<String, Object> configurationMap = new HashMap<>();
configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
@@ -99,10 +164,10 @@ public class RackAwareTaskAssignorTest {
}
@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 class RackAwareTaskAssignorTest {
);
// 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 class RackAwareTaskAssignorTest {
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 class RackAwareTaskAssignorTest {
// 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<UUID, Map<String, Optional<String>>> 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 class RackAwareTaskAssignorTest {
);
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 class RackAwareTaskAssignorTest {
);
// 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0)
+ ));
+ final SortedSet<TaskId> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> 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<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize);
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getRandomCluster(nodeSize, tpSize),
+ taskTopicPartitionMap,
+ getTopologyGroupTaskMap(),
+ getRandomProcessRacks(clientSize, nodeSize),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize);
+ final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+
+ final Map<UUID, Integer> 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<UUID, ClientState> 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<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize);
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getRandomCluster(nodeSize, tpSize),
+ taskTopicPartitionMap,
+ getTopologyGroupTaskMap(),
+ getRandomProcessRacks(clientSize, nodeSize),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize);
+ final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ for (final Entry<UUID, ClientState> 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<TaskId, UUID> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_2, clientState0),
+ mkEntry(UUID_5, clientState1)
+ ));
+ final SortedSet<TaskId> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_2, clientState0),
+ mkEntry(UUID_5, clientState1)
+ ));
+ final SortedSet<TaskId> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_2, clientState0),
+ mkEntry(UUID_5, clientState1)
+ ));
+ final SortedSet<TaskId> 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<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_2, clientState0),
+ mkEntry(UUID_5, clientState1)
+ ));
+ final SortedSet<TaskId> 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<Node> 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<PartitionInfo> 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<UUID, Map<String, Optional<String>>> getRandomProcessRacks(final int clientSize, final int nodeSize) {
+ final List<String> racks = new ArrayList<>(nodeSize);
+ for (int i = 0; i < nodeSize; i++) {
+ racks.add("rack" + i);
+ }
+ Collections.shuffle(racks);
+ final Map<UUID, Map<String, Optional<String>>> 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<TaskId, Set<TopicPartition>> getTaskTopicPartitionMap(final int tpSize) {
+ final SortedMap<TaskId, Set<TopicPartition>> 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<UUID, ClientState> getRandomClientState(final int clientSize, final int tpSize) {
+ final SortedMap<UUID, ClientState> clientStates = new TreeMap<>();
+ final List<TaskId> 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<Entry<UUID, ClientState>> 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<UUID, Map<String, Optional<String>>> 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<UUID, Map<String, Optional<String>>> getProcessRacksForProcess0() {
return getProcessRacksForProcess0(false);
}
@@ -268,24 +783,34 @@ public class RackAwareTaskAssignorTest {
private Map<UUID, Map<String, Optional<String>>> getProcessRacksForProcess0(final boolean missingRack) {
final Map<UUID, Map<String, Optional<String>>> processRacks = new HashMap<>();
final Optional<String> 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<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForTask1() {
- return getTaskTopicPartitionMapForTask1(false);
+ private Map<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForTask0() {
+ return getTaskTopicPartitionMapForTask0(false);
}
- private Map<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForTask1(final boolean extraTopic) {
+ private Map<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForTask0(final boolean extraTopic) {
final Set<TopicPartition> 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<TaskId, Set<TopicPartition>> 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<Subtopology, Set<TaskId>> getTopologyGroupTaskMap() {
- return Collections.singletonMap(subtopology1, Collections.singleton(new TaskId(1, 1)));
+ return Collections.singletonMap(SUBTOPOLOGY_0, Collections.singleton(new TaskId(1, 1)));
}
}