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/08/10 02:50:50 UTC
[kafka] branch trunk updated: KAFKA-15022: [9/N] use RackAwareTaskAssignor in StickyTaskAssignor (#14178)
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 96631c25d5f KAFKA-15022: [9/N] use RackAwareTaskAssignor in StickyTaskAssignor (#14178)
96631c25d5f is described below
commit 96631c25d5feb0d946dace56ddcd02515cc5e3a0
Author: Hao Li <11...@users.noreply.github.com>
AuthorDate: Wed Aug 9 19:50:43 2023 -0700
KAFKA-15022: [9/N] use RackAwareTaskAssignor in StickyTaskAssignor (#14178)
Part of KIP-925.
Use rack aware assignor in StickyTaskAssignor.
Reviewers: Matthias J. Sax <ma...@confluent.io>
---
.../internals/StreamsPartitionAssignor.java | 12 +-
.../internals/assignment/ClientState.java | 12 +
.../assignment/FallbackPriorTaskAssignor.java | 3 +-
.../assignment/HighAvailabilityTaskAssignor.java | 4 +-
.../assignment/RackAwareTaskAssignor.java | 20 +-
.../internals/assignment/StickyTaskAssignor.java | 56 ++-
.../integration/utils/EmbeddedKafkaCluster.java | 4 +
.../internals/assignment/AssignmentTestUtils.java | 257 +++++++++++---
.../internals/assignment/ClientStateTest.java | 14 +
.../HighAvailabilityTaskAssignorTest.java | 238 ++++++++-----
.../assignment/RackAwareTaskAssignorTest.java | 108 ++++--
.../assignment/StickyTaskAssignorTest.java | 395 ++++++++++++++++++++-
.../assignment/TaskAssignorConvergenceTest.java | 11 +-
13 files changed, 935 insertions(+), 199 deletions(-)
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
index 453097cc3f3..330d46a1780 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
@@ -640,8 +640,16 @@ public class StreamsPartitionAssignor implements ConsumerPartitionAssignor, Conf
final TaskAssignor taskAssignor = createTaskAssignor(lagComputationSuccessful);
- final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(fullMetadata, partitionsForTask,
- changelogTopics.changelogPartionsForTask(), tasksForTopicGroup, racksForProcessConsumer, internalTopicManager, assignmentConfigs);
+ final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(
+ fullMetadata,
+ partitionsForTask,
+ changelogTopics.changelogPartionsForTask(),
+ tasksForTopicGroup,
+ racksForProcessConsumer,
+ internalTopicManager,
+ assignmentConfigs,
+ time
+ );
final boolean probingRebalanceNeeded = taskAssignor.assign(clientStates,
allTasks,
statefulTasks,
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java
index 32599287bef..71fc1c3ef00 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java
@@ -112,6 +112,18 @@ public class ClientState {
this.processId = processId;
}
+ // For testing only
+ public ClientState(final ClientState clientState) {
+ this(
+ new HashSet<>(clientState.previousActiveTasks.taskIds()),
+ new HashSet<>(clientState.previousStandbyTasks.taskIds()),
+ clientState.taskLagTotals,
+ clientState.clientTags,
+ clientState.capacity,
+ clientState.processId
+ );
+ }
+
int capacity() {
return capacity;
}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/FallbackPriorTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/FallbackPriorTaskAssignor.java
index 48533910309..8e261eb296b 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/FallbackPriorTaskAssignor.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/FallbackPriorTaskAssignor.java
@@ -44,7 +44,8 @@ public class FallbackPriorTaskAssignor implements TaskAssignor {
final Set<TaskId> statefulTaskIds,
final RackAwareTaskAssignor rackAwareTaskAssignor,
final AssignmentConfigs configs) {
- delegate.assign(clients, allTaskIds, statefulTaskIds, rackAwareTaskAssignor, configs);
+ // Pass null for RackAwareTaskAssignor to disable it if we fallback
+ delegate.assign(clients, allTaskIds, statefulTaskIds, null, configs);
return true;
}
}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java
index a319fc15d7d..6dc8d1886ce 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignor.java
@@ -38,6 +38,8 @@ import java.util.function.BiPredicate;
import java.util.function.Function;
import static org.apache.kafka.common.utils.Utils.diff;
+import static org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.STATELESS_NON_OVERLAP_COST;
+import static org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.STATELESS_TRAFFIC_COST;
import static org.apache.kafka.streams.processor.internals.assignment.TaskMovement.assignActiveTaskMovements;
import static org.apache.kafka.streams.processor.internals.assignment.TaskMovement.assignStandbyTaskMovements;
@@ -45,8 +47,6 @@ public class HighAvailabilityTaskAssignor implements TaskAssignor {
private static final Logger log = LoggerFactory.getLogger(HighAvailabilityTaskAssignor.class);
private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
private static final int DEFAULT_STATEFUL_NON_OVERLAP_COST = 1;
- private static final int STATELESS_TRAFFIC_COST = 1;
- private static final int STATELESS_NON_OVERLAP_COST = 0;
@Override
public boolean assign(final Map<UUID, ClientState> clients,
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 c3921fc9b75..0c8f9492a3e 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
@@ -41,6 +41,7 @@ import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;
+import org.apache.kafka.common.utils.Time;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.processor.TaskId;
@@ -60,8 +61,11 @@ public class RackAwareTaskAssignor {
final Map<UUID, ClientState> clientStateMap);
}
- private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ // For stateless tasks, it's ok to move them around. So we have 0 non_overlap_cost
+ public static final int STATELESS_TRAFFIC_COST = 1;
+ public static final int STATELESS_NON_OVERLAP_COST = 0;
+ private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
private static final int SOURCE_ID = -1;
// This is number is picked based on testing. Usually the optimization for standby assignment
// stops after 3 rounds
@@ -75,6 +79,7 @@ public class RackAwareTaskAssignor {
private final Map<UUID, String> racksForProcess;
private final InternalTopicManager internalTopicManager;
private final boolean validClientRack;
+ private final Time time;
private Boolean canEnable = null;
public RackAwareTaskAssignor(final Cluster fullMetadata,
@@ -83,7 +88,8 @@ public class RackAwareTaskAssignor {
final Map<Subtopology, Set<TaskId>> tasksForTopicGroup,
final Map<UUID, Map<String, Optional<String>>> racksForProcessConsumer,
final InternalTopicManager internalTopicManager,
- final AssignmentConfigs assignmentConfigs) {
+ final AssignmentConfigs assignmentConfigs,
+ final Time time) {
this.fullMetadata = fullMetadata;
this.partitionsForTask = partitionsForTask;
this.changelogPartitionsForTask = changelogPartitionsForTask;
@@ -91,6 +97,7 @@ public class RackAwareTaskAssignor {
this.assignmentConfigs = assignmentConfigs;
this.racksForPartition = new HashMap<>();
this.racksForProcess = new HashMap<>();
+ this.time = Objects.requireNonNull(time, "Time was not specified");
validClientRack = validateClientRack(racksForProcessConsumer);
}
@@ -346,6 +353,7 @@ public class RackAwareTaskAssignor {
log.info("Assignment before active task optimization is {}\n with cost {}", clientStates,
activeTasksCost(activeTasks, clientStates, trafficCost, nonOverlapCost));
+ final long startTime = time.milliseconds();
final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
final Map<TaskId, UUID> taskClientMap = new HashMap<>();
@@ -359,7 +367,8 @@ public class RackAwareTaskAssignor {
assignTaskFromMinCostFlow(graph, clientList, taskIdList, clientStates, originalAssignedTaskNumber,
taskClientMap, ClientState::assignActive, ClientState::unassignActive, ClientState::hasActiveTask);
- log.info("Assignment after active task optimization is {}\n with cost {}", clientStates, cost);
+ final long duration = time.milliseconds() - startTime;
+ log.info("Assignment after {} milliseconds for active task optimization is {}\n with cost {}", duration, clientStates, cost);
return cost;
}
@@ -373,6 +382,7 @@ public class RackAwareTaskAssignor {
.sorted()
.collect(Collectors.toList());
+ final long startTime = time.milliseconds();
final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
final SortedSet<TaskId> standbyTasks = new TreeSet<>();
clientStates.values().forEach(clientState -> standbyTasks.addAll(clientState.standbyTasks()));
@@ -430,7 +440,9 @@ public class RackAwareTaskAssignor {
}
}
final long cost = standByTasksCost(standbyTasks, clientStates, trafficCost, nonOverlapCost);
- log.info("Assignment after {} rounds of standby task optimization is {}\n with cost {}", round, clientStates, cost);
+
+ final long duration = time.milliseconds() - startTime;
+ log.info("Assignment after {} rounds and {} milliseconds for standby task optimization is {}\n with cost {}", round, duration, clientStates, cost);
return cost;
}
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignor.java
index ec0eff41259..355ecd2e804 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignor.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignor.java
@@ -16,6 +16,13 @@
*/
package org.apache.kafka.streams.processor.internals.assignment;
+import static org.apache.kafka.common.utils.Utils.diff;
+import static org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.STATELESS_NON_OVERLAP_COST;
+import static org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor.STATELESS_TRAFFIC_COST;
+
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
import org.slf4j.Logger;
@@ -36,11 +43,18 @@ import java.util.UUID;
public class StickyTaskAssignor implements TaskAssignor {
private static final Logger log = LoggerFactory.getLogger(StickyTaskAssignor.class);
+
+ // For stateful tasks, by default we want to maintain stickiness. So we have higher non_overlap_cost
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 1;
+ private static final int DEFAULT_STATEFUL_NON_OVERLAP_COST = 10;
+
private Map<UUID, ClientState> clients;
private Set<TaskId> allTaskIds;
- private Set<TaskId> standbyTaskIds;
+ private Set<TaskId> statefulTaskIds;
private final Map<TaskId, UUID> previousActiveTaskAssignment = new HashMap<>();
private final Map<TaskId, Set<UUID>> previousStandbyTaskAssignment = new HashMap<>();
+ private RackAwareTaskAssignor rackAwareTaskAssignor; // nullable if passed from FallbackPriorTaskAssignor
+ private AssignmentConfigs configs;
private TaskPairs taskPairs;
private final boolean mustPreserveActiveTaskAssignment;
@@ -61,19 +75,55 @@ public class StickyTaskAssignor implements TaskAssignor {
final AssignmentConfigs configs) {
this.clients = clients;
this.allTaskIds = allTaskIds;
- this.standbyTaskIds = statefulTaskIds;
+ this.statefulTaskIds = statefulTaskIds;
+ this.rackAwareTaskAssignor = rackAwareTaskAssignor;
+ this.configs = configs;
final int maxPairs = allTaskIds.size() * (allTaskIds.size() - 1) / 2;
taskPairs = new TaskPairs(maxPairs);
mapPreviousTaskAssignment(clients);
assignActive();
+ optimizeActive();
assignStandby(configs.numStandbyReplicas);
+ optimizeStandby();
return false;
}
+ private void optimizeStandby() {
+ if (configs.numStandbyReplicas > 0 && rackAwareTaskAssignor != null && rackAwareTaskAssignor.canEnableRackAwareAssignor()) {
+ final int trafficCost = configs.rackAwareAssignmentTrafficCost == null ?
+ DEFAULT_STATEFUL_TRAFFIC_COST : configs.rackAwareAssignmentTrafficCost;
+ final int nonOverlapCost = configs.rackAwareAssignmentNonOverlapCost == null ?
+ DEFAULT_STATEFUL_NON_OVERLAP_COST : configs.rackAwareAssignmentNonOverlapCost;
+ final TreeMap<UUID, ClientState> clientStates = new TreeMap<>(clients);
+
+ rackAwareTaskAssignor.optimizeStandbyTasks(clientStates, trafficCost, nonOverlapCost, (s, d, t, c) -> true);
+ }
+ }
+
+ private void optimizeActive() {
+ if (rackAwareTaskAssignor != null && rackAwareTaskAssignor.canEnableRackAwareAssignor()) {
+ final int trafficCost = configs.rackAwareAssignmentTrafficCost == null ?
+ DEFAULT_STATEFUL_TRAFFIC_COST : configs.rackAwareAssignmentTrafficCost;
+ final int nonOverlapCost = configs.rackAwareAssignmentNonOverlapCost == null ?
+ DEFAULT_STATEFUL_NON_OVERLAP_COST : configs.rackAwareAssignmentNonOverlapCost;
+
+ final SortedSet<TaskId> statefulTasks = new TreeSet<>(statefulTaskIds);
+ final TreeMap<UUID, ClientState> clientStates = new TreeMap<>(clients);
+
+ rackAwareTaskAssignor.optimizeActiveTasks(statefulTasks, clientStates, trafficCost, nonOverlapCost);
+
+ final TreeSet<TaskId> statelessTasks = (TreeSet<TaskId>) diff(TreeSet::new, allTaskIds, statefulTasks);
+ if (!statelessTasks.isEmpty()) {
+ rackAwareTaskAssignor.optimizeActiveTasks(statelessTasks, clientStates,
+ STATELESS_TRAFFIC_COST, STATELESS_NON_OVERLAP_COST);
+ }
+ }
+ }
+
private void assignStandby(final int numStandbyReplicas) {
- for (final TaskId taskId : standbyTaskIds) {
+ for (final TaskId taskId : statefulTaskIds) {
for (int i = 0; i < numStandbyReplicas; i++) {
final Set<UUID> ids = findClientsWithoutAssignedTask(taskId);
if (ids.isEmpty()) {
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
index dc3f091fec5..304402449e2 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java
@@ -89,6 +89,10 @@ public class EmbeddedKafkaCluster {
final List<Properties> brokerConfigOverrides,
final long mockTimeMillisStart,
final long mockTimeNanoStart) {
+ if (!brokerConfigOverrides.isEmpty() && brokerConfigOverrides.size() != numBrokers) {
+ throw new IllegalArgumentException("Size of brokerConfigOverrides " + brokerConfigOverrides.size()
+ + " must match broker number " + numBrokers);
+ }
brokers = new KafkaEmbedded[numBrokers];
this.brokerConfig = brokerConfig;
time = new MockTime(mockTimeMillisStart, mockTimeNanoStart);
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 7f91d46bb0d..76348fbcef5 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
@@ -43,6 +43,7 @@ import org.apache.kafka.streams.processor.internals.InternalTopicManager;
import org.apache.kafka.streams.processor.internals.Task;
import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
import org.apache.kafka.test.MockClientSupplier;
import org.apache.kafka.test.MockInternalTopicManager;
import org.hamcrest.BaseMatcher;
@@ -77,10 +78,15 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anySet;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
public final class AssignmentTestUtils {
@@ -117,15 +123,20 @@ public final class AssignmentTestUtils {
public static final String TP_0_NAME = "topic0";
public static final String TP_1_NAME = "topic1";
public static final String TP_2_NAME = "topic2";
+ public static final String TP_3_NAME = "topic3";
public static final String CHANGELOG_TP_0_NAME = "store-0-changelog";
public static final String CHANGELOG_TP_1_NAME = "store-1-changelog";
public static final String CHANGELOG_TP_2_NAME = "store-2-changelog";
+ public static final String CHANGELOG_TP_3_NAME = "store-3-changelog";
public static final TopicPartition CHANGELOG_TP_0_0 = new TopicPartition(CHANGELOG_TP_0_NAME, 0);
public static final TopicPartition CHANGELOG_TP_0_1 = new TopicPartition(CHANGELOG_TP_0_NAME, 1);
public static final TopicPartition CHANGELOG_TP_0_2 = new TopicPartition(CHANGELOG_TP_0_NAME, 2);
public static final TopicPartition CHANGELOG_TP_0_3 = new TopicPartition(CHANGELOG_TP_0_NAME, 3);
+ public static final TopicPartition CHANGELOG_TP_0_4 = new TopicPartition(CHANGELOG_TP_0_NAME, 4);
+ public static final TopicPartition CHANGELOG_TP_0_5 = new TopicPartition(CHANGELOG_TP_0_NAME, 5);
+ public static final TopicPartition CHANGELOG_TP_0_6 = new TopicPartition(CHANGELOG_TP_0_NAME, 6);
public static final TopicPartition CHANGELOG_TP_1_0 = new TopicPartition(CHANGELOG_TP_1_NAME, 0);
public static final TopicPartition CHANGELOG_TP_1_1 = new TopicPartition(CHANGELOG_TP_1_NAME, 1);
public static final TopicPartition CHANGELOG_TP_1_2 = new TopicPartition(CHANGELOG_TP_1_NAME, 2);
@@ -133,11 +144,18 @@ public final class AssignmentTestUtils {
public static final TopicPartition CHANGELOG_TP_2_0 = new TopicPartition(CHANGELOG_TP_2_NAME, 0);
public static final TopicPartition CHANGELOG_TP_2_1 = new TopicPartition(CHANGELOG_TP_2_NAME, 1);
public static final TopicPartition CHANGELOG_TP_2_2 = new TopicPartition(CHANGELOG_TP_2_NAME, 2);
+ public static final TopicPartition CHANGELOG_TP_2_3 = new TopicPartition(CHANGELOG_TP_2_NAME, 3);
+ public static final TopicPartition CHANGELOG_TP_3_0 = new TopicPartition(CHANGELOG_TP_3_NAME, 0);
+ public static final TopicPartition CHANGELOG_TP_3_1 = new TopicPartition(CHANGELOG_TP_3_NAME, 1);
+ public static final TopicPartition CHANGELOG_TP_3_2 = new TopicPartition(CHANGELOG_TP_3_NAME, 2);
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_0_3 = new TopicPartition(TP_0_NAME, 3);
+ public static final TopicPartition TP_0_4 = new TopicPartition(TP_0_NAME, 4);
+ public static final TopicPartition TP_0_5 = new TopicPartition(TP_0_NAME, 5);
+ public static final TopicPartition TP_0_6 = new TopicPartition(TP_0_NAME, 6);
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);
@@ -145,11 +163,18 @@ public final class AssignmentTestUtils {
public static final TopicPartition TP_2_0 = new TopicPartition(TP_2_NAME, 0);
public static final TopicPartition TP_2_1 = new TopicPartition(TP_2_NAME, 1);
public static final TopicPartition TP_2_2 = new TopicPartition(TP_2_NAME, 2);
+ public static final TopicPartition TP_2_3 = new TopicPartition(TP_2_NAME, 3);
+ public static final TopicPartition TP_3_0 = new TopicPartition(TP_3_NAME, 0);
+ public static final TopicPartition TP_3_1 = new TopicPartition(TP_3_NAME, 1);
+ public static final TopicPartition TP_3_2 = new TopicPartition(TP_3_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_0_2 = new PartitionInfo(TP_0_NAME, 2, NODE_1, REPLICA_1, REPLICA_1);
public static final PartitionInfo PI_0_3 = new PartitionInfo(TP_0_NAME, 3, NODE_2, REPLICA_2, REPLICA_2);
+ public static final PartitionInfo PI_0_4 = new PartitionInfo(TP_0_NAME, 4, NODE_3, REPLICA_3, REPLICA_3);
+ public static final PartitionInfo PI_0_5 = new PartitionInfo(TP_0_NAME, 5, NODE_4, REPLICA_4, REPLICA_4);
+ public static final PartitionInfo PI_0_6 = new PartitionInfo(TP_0_NAME, 6, NODE_2, REPLICA_2, REPLICA_2);
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);
@@ -157,6 +182,10 @@ public final class AssignmentTestUtils {
public static final PartitionInfo PI_2_0 = new PartitionInfo(TP_2_NAME, 0, NODE_4, REPLICA_4, REPLICA_4);
public static final PartitionInfo PI_2_1 = new PartitionInfo(TP_2_NAME, 1, NODE_3, REPLICA_3, REPLICA_3);
public static final PartitionInfo PI_2_2 = new PartitionInfo(TP_2_NAME, 2, NODE_1, REPLICA_4, REPLICA_4);
+ public static final PartitionInfo PI_2_3 = new PartitionInfo(TP_2_NAME, 3, NODE_0, REPLICA_0, REPLICA_0);
+ public static final PartitionInfo PI_3_0 = new PartitionInfo(TP_3_NAME, 0, NODE_2, REPLICA_2, REPLICA_2);
+ public static final PartitionInfo PI_3_1 = new PartitionInfo(TP_3_NAME, 1, NODE_3, REPLICA_3, REPLICA_3);
+ public static final PartitionInfo PI_3_2 = new PartitionInfo(TP_3_NAME, 2, NODE_4, REPLICA_4, REPLICA_4);
public static final TaskId TASK_0_0 = new TaskId(0, 0);
public static final TaskId TASK_0_1 = new TaskId(0, 1);
@@ -173,6 +202,9 @@ public final class AssignmentTestUtils {
public static final TaskId TASK_2_1 = new TaskId(2, 1);
public static final TaskId TASK_2_2 = new TaskId(2, 2);
public static final TaskId TASK_2_3 = new TaskId(2, 3);
+ public static final TaskId TASK_3_0 = new TaskId(3, 0);
+ public static final TaskId TASK_3_1 = new TaskId(3, 1);
+ public static final TaskId TASK_3_2 = new TaskId(3, 2);
public static final TaskId NAMED_TASK_T0_0_0 = new TaskId(0, 0, "topology0");
public static final TaskId NAMED_TASK_T0_0_1 = new TaskId(0, 1, "topology0");
@@ -194,6 +226,7 @@ public final class AssignmentTestUtils {
private static final String USER_END_POINT = "localhost:8080";
private static final String APPLICATION_ID = "stream-partition-assignor-test";
+ private static Random random;
public static final String TOPIC_PREFIX = "topic";
public static final String CHANGELOG_TOPIC_PREFIX = "changelog-topic";
public static final String RACK_PREFIX = "rack";
@@ -583,22 +616,26 @@ public final class AssignmentTestUtils {
for (int i = 0; i < nodeSize; i++) {
nodeList.add(new Node(i, "node" + i, 1, RACK_PREFIX + i));
}
- Collections.shuffle(nodeList);
+ final Random rand = getRandom();
+ Collections.shuffle(nodeList, rand);
return nodeList;
}
- static Node[] getRandomReplica(final List<Node> nodeList, final int index) {
- final Node firstNode = nodeList.get(index % nodeList.size());
- final Node secondNode = nodeList.get((index + 1) % nodeList.size());
+ static Node[] getRandomReplica(final List<Node> nodeList, final int index, final int partition) {
+ final Node firstNode = nodeList.get((index * partition) % nodeList.size());
+ final Node secondNode = nodeList.get((index * partition + 1) % nodeList.size());
return new Node[] {firstNode, secondNode};
}
- static Cluster getRandomCluster(final int nodeSize, final int tpSize) {
+ static Cluster getRandomCluster(final int nodeSize, final int tpSize, final int partitionSize) {
final List<Node> nodeList = getRandomNodes(nodeSize);
final Set<PartitionInfo> partitionInfoSet = new HashSet<>();
for (int i = 0; i < tpSize; i++) {
- final Node[] replica = getRandomReplica(nodeList, i);
- partitionInfoSet.add(new PartitionInfo(TOPIC_PREFIX + i, 0, replica[0], replica, replica));
+ for (int j = 0; j < partitionSize; j++) {
+ final Node[] replica = getRandomReplica(nodeList, i, j);
+ partitionInfoSet.add(
+ new PartitionInfo(TOPIC_PREFIX + i, j, replica[0], replica, replica));
+ }
}
return new Cluster(
@@ -615,23 +652,26 @@ public final class AssignmentTestUtils {
for (int i = 0; i < nodeSize; i++) {
racks.add(RACK_PREFIX + i);
}
- Collections.shuffle(racks);
+ final Random rand = getRandom();
+ Collections.shuffle(racks, rand);
final Map<UUID, Map<String, Optional<String>>> processRacks = new HashMap<>();
- for (int i = 0; i < clientSize; i++) {
+ for (int i = 1; i <= clientSize; i++) {
final String rack = racks.get(i % nodeSize);
processRacks.put(uuidForInt(i), mkMap(mkEntry("1", Optional.of(rack))));
}
return processRacks;
}
- static SortedMap<TaskId, Set<TopicPartition>> getTaskTopicPartitionMap(final int tpSize, final boolean changelog) {
+ static SortedMap<TaskId, Set<TopicPartition>> getTaskTopicPartitionMap(final int tpSize, final int partitionSize, final boolean changelog) {
final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = new TreeMap<>();
final String topicName = changelog ? CHANGELOG_TOPIC_PREFIX : TOPIC_PREFIX;
for (int i = 0; i < tpSize; i++) {
- taskTopicPartitionMap.put(new TaskId(i, 0), mkSet(
- new TopicPartition(topicName + i, 0),
- new TopicPartition(topicName + (i + 1) % tpSize, 0)
- ));
+ for (int j = 0; j < partitionSize; j++) {
+ taskTopicPartitionMap.put(new TaskId(i, j), mkSet(
+ new TopicPartition(topicName + i, j),
+ new TopicPartition(topicName + ((i + 1) % tpSize), j)
+ ));
+ }
}
return taskTopicPartitionMap;
}
@@ -654,7 +694,7 @@ public final class AssignmentTestUtils {
return configurationMap;
}
- static InternalTopicManager mockInternalTopicManagerForRandomChangelog(final int nodeSize, final int tpSize) {
+ static InternalTopicManager mockInternalTopicManagerForRandomChangelog(final int nodeSize, final int tpSize, final int partitionSize) {
final MockTime time = new MockTime();
final StreamsConfig streamsConfig = new StreamsConfig(configProps(true));
final MockClientSupplier mockClientSupplier = new MockClientSupplier();
@@ -671,12 +711,14 @@ public final class AssignmentTestUtils {
for (int i = 0; i < tpSize; i++) {
final String topicName = CHANGELOG_TOPIC_PREFIX + i;
changelogNames.add(topicName);
+ for (int j = 0; j < partitionSize; j++) {
- final Node firstNode = nodeList.get(i % nodeSize);
- final Node secondNode = nodeList.get((i + 1) % nodeSize);
- final TopicPartitionInfo info = new TopicPartitionInfo(0, firstNode, Arrays.asList(firstNode, secondNode), Collections.emptyList());
+ final Node[] replica = getRandomReplica(nodeList, i, j);
+ final TopicPartitionInfo info = new TopicPartitionInfo(j, replica[0],
+ Arrays.asList(replica), Arrays.asList(replica));
- topicPartitionInfo.computeIfAbsent(topicName, tp -> new ArrayList<>()).add(info);
+ topicPartitionInfo.computeIfAbsent(topicName, tp -> new ArrayList<>()).add(info);
+ }
}
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
@@ -684,34 +726,65 @@ public final class AssignmentTestUtils {
return spyTopicManager;
}
- static SortedMap<UUID, ClientState> getRandomClientState(final int clientSize, final int tpSize, final int maxCapacity) {
- return getRandomClientState(clientSize, tpSize, maxCapacity, true);
+ static SortedMap<UUID, ClientState> getRandomClientState(final int clientSize, final int tpSize, final int partitionSize, final int maxCapacity, final Set<TaskId> statefulTasks) {
+ return getRandomClientState(clientSize, tpSize, partitionSize, maxCapacity, true, statefulTasks);
}
- static SortedMap<UUID, ClientState> getRandomClientState(final int clientSize, final int tpSize, final int maxCapacity, final boolean initialAssignment) {
+ static List<Set<TaskId>> getRandomSubset(final Set<TaskId> taskIds, final int listSize) {
+ final Random rand = getRandom();
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ Collections.shuffle(taskIdList, rand);
+ int start = 0;
+ final List<Set<TaskId>> subSets = new ArrayList<>(listSize);
+ for (int i = 0; i < listSize; i++) {
+ final int remaining = taskIdList.size() - start;
+ final Set<TaskId> subset = new HashSet<>();
+ if (remaining != 0) {
+ // In last round, get all tasks
+ final int subSetSize = (i == listSize - 1) ? remaining : rand.nextInt(remaining) + 1;
+ for (int j = 0; j < subSetSize; j++) {
+ subset.add(taskIdList.get(start + j));
+ }
+ start += subSetSize;
+ }
+ subSets.add(subset);
+ }
+ return subSets;
+ }
+
+ static SortedMap<UUID, ClientState> getRandomClientState(final int clientSize, final int tpSize, final int partitionSize, final int maxCapacity, final boolean initialAssignment, final Set<TaskId> statefulTasks) {
final SortedMap<UUID, ClientState> clientStates = new TreeMap<>();
- final Map<TaskId, Long> taskLags = new HashMap<>();
+ final Map<TaskId, Long> taskLags = statefulTasks.stream().collect(Collectors.toMap(taskId -> taskId, taskId -> 0L));
+ final Set<TaskId> taskIds = new HashSet<>();
for (int i = 0; i < tpSize; i++) {
- taskLags.put(new TaskId(i, 0), 0L);
+ for (int j = 0; j < partitionSize; j++) {
+ taskIds.add(new TaskId(i, j));
+ }
}
- final long seed = System.currentTimeMillis();
- System.out.println("seed: " + seed);
- final Random random = new Random(seed);
- for (int i = 0; i < clientSize; i++) {
- final int capacity = random.nextInt(maxCapacity) + 1;
- final ClientState clientState = new ClientState(emptySet(), emptySet(), taskLags, EMPTY_CLIENT_TAGS, capacity, uuidForInt(i));
- clientStates.put(uuidForInt(i), clientState);
+ final Set<TaskId> missingTaskIds = taskLags.keySet().stream().filter(id -> !taskIds.contains(id)).collect(
+ Collectors.toSet());
+ if (!missingTaskIds.isEmpty()) {
+ throw new IllegalArgumentException(missingTaskIds + " missing in all task ids " + taskIds);
+ }
+
+ final List<Set<TaskId>> previousActives = getRandomSubset(taskIds, clientSize);
+ final List<Set<TaskId>> previousStandbys = getRandomSubset(statefulTasks, clientSize);
+
+ final Random rand = getRandom();
+
+ for (int i = 1; i <= clientSize; i++) {
+ final int capacity = rand.nextInt(maxCapacity) + 1;
+ final UUID processId = uuidForInt(i);
+ final ClientState clientState = new ClientState(previousActives.get(i - 1), previousStandbys.get(i - 1), taskLags, EMPTY_CLIENT_TAGS, capacity, processId);
+ clientStates.put(processId, clientState);
}
if (initialAssignment) {
Iterator<Entry<UUID, ClientState>> iterator = clientStates.entrySet().iterator();
- final List<TaskId> taskIds = new ArrayList<>(tpSize);
- for (int i = 0; i < tpSize; i++) {
- taskIds.add(new TaskId(i, 0));
- }
- Collections.shuffle(taskIds);
- for (final TaskId taskId : taskIds) {
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ Collections.shuffle(taskIdList, rand);
+ for (final TaskId taskId : taskIdList) {
if (!iterator.hasNext()) {
iterator = clientStates.entrySet().iterator();
}
@@ -725,7 +798,26 @@ public final class AssignmentTestUtils {
return new Cluster(
"cluster",
mkSet(NODE_0, NODE_1, NODE_2, NODE_3, NODE_4),
- mkSet(PI_0_0, PI_0_1, PI_0_2, PI_0_3, PI_1_0, PI_1_1, PI_1_2, PI_1_3, PI_2_0, PI_2_1, PI_2_2),
+ mkSet(
+ PI_0_0,
+ PI_0_1,
+ PI_0_2,
+ PI_0_3,
+ PI_0_4,
+ PI_0_5,
+ PI_0_6,
+ PI_1_0,
+ PI_1_1,
+ PI_1_2,
+ PI_1_3,
+ PI_2_0,
+ PI_2_1,
+ PI_2_2,
+ PI_2_3,
+ PI_3_0,
+ PI_3_1,
+ PI_3_2
+ ),
Collections.emptySet(),
Collections.emptySet()
);
@@ -737,13 +829,20 @@ public final class AssignmentTestUtils {
mkEntry(TASK_0_1, mkSet(TP_0_1)),
mkEntry(TASK_0_2, mkSet(TP_0_2)),
mkEntry(TASK_0_3, mkSet(TP_0_3)),
+ mkEntry(TASK_0_4, mkSet(TP_0_4)),
+ mkEntry(TASK_0_5, mkSet(TP_0_5)),
+ mkEntry(TASK_0_6, mkSet(TP_0_6)),
mkEntry(TASK_1_0, mkSet(TP_1_0)),
mkEntry(TASK_1_1, mkSet(TP_1_1)),
mkEntry(TASK_1_2, mkSet(TP_1_2)),
mkEntry(TASK_1_3, mkSet(TP_1_3)),
mkEntry(TASK_2_0, mkSet(TP_2_0)),
mkEntry(TASK_2_1, mkSet(TP_2_1)),
- mkEntry(TASK_2_2, mkSet(TP_2_2))
+ mkEntry(TASK_2_2, mkSet(TP_2_2)),
+ mkEntry(TASK_2_3, mkSet(TP_2_3)),
+ mkEntry(TASK_3_0, mkSet(TP_3_0)),
+ mkEntry(TASK_3_1, mkSet(TP_3_1)),
+ mkEntry(TASK_3_2, mkSet(TP_3_2))
);
}
@@ -753,13 +852,20 @@ public final class AssignmentTestUtils {
mkEntry(TASK_0_1, mkSet(CHANGELOG_TP_0_1)),
mkEntry(TASK_0_2, mkSet(CHANGELOG_TP_0_2)),
mkEntry(TASK_0_3, mkSet(CHANGELOG_TP_0_3)),
+ mkEntry(TASK_0_4, mkSet(CHANGELOG_TP_0_4)),
+ mkEntry(TASK_0_5, mkSet(CHANGELOG_TP_0_5)),
+ mkEntry(TASK_0_6, mkSet(CHANGELOG_TP_0_6)),
mkEntry(TASK_1_0, mkSet(CHANGELOG_TP_1_0)),
mkEntry(TASK_1_1, mkSet(CHANGELOG_TP_1_1)),
mkEntry(TASK_1_2, mkSet(CHANGELOG_TP_1_2)),
mkEntry(TASK_1_3, mkSet(CHANGELOG_TP_1_3)),
mkEntry(TASK_2_0, mkSet(CHANGELOG_TP_2_0)),
mkEntry(TASK_2_1, mkSet(CHANGELOG_TP_2_1)),
- mkEntry(TASK_2_2, mkSet(CHANGELOG_TP_2_2))
+ mkEntry(TASK_2_2, mkSet(CHANGELOG_TP_2_2)),
+ mkEntry(TASK_2_3, mkSet(CHANGELOG_TP_2_3)),
+ mkEntry(TASK_3_0, mkSet(CHANGELOG_TP_3_0)),
+ mkEntry(TASK_3_1, mkSet(CHANGELOG_TP_3_1)),
+ mkEntry(TASK_3_2, mkSet(CHANGELOG_TP_3_2))
);
}
@@ -782,7 +888,10 @@ public final class AssignmentTestUtils {
new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_0), Collections.emptyList()),
new TopicPartitionInfo(1, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList()),
new TopicPartitionInfo(2, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList()),
- new TopicPartitionInfo(3, NODE_2, Arrays.asList(REPLICA_2), Collections.emptyList())
+ new TopicPartitionInfo(3, NODE_2, Arrays.asList(REPLICA_2), Collections.emptyList()),
+ new TopicPartitionInfo(4, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList()),
+ new TopicPartitionInfo(5, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList()),
+ new TopicPartitionInfo(6, NODE_0, Arrays.asList(REPLICA_0), Collections.emptyList())
)
),
mkEntry(
@@ -797,11 +906,19 @@ public final class AssignmentTestUtils {
CHANGELOG_TP_2_NAME, Arrays.asList(
new TopicPartitionInfo(0, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList()),
new TopicPartitionInfo(1, NODE_2, Arrays.asList(REPLICA_2), Collections.emptyList()),
- new TopicPartitionInfo(2, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList())
+ new TopicPartitionInfo(2, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList()),
+ new TopicPartitionInfo(3, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList())
+ )
+ ),
+ mkEntry(
+ CHANGELOG_TP_3_NAME, Arrays.asList(
+ new TopicPartitionInfo(0, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList()),
+ new TopicPartitionInfo(1, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList()),
+ new TopicPartitionInfo(2, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList())
)
)
)
- ).when(spyTopicManager).getTopicPartitionInfo(mkSet(CHANGELOG_TP_0_NAME, CHANGELOG_TP_1_NAME, CHANGELOG_TP_2_NAME));
+ ).when(spyTopicManager).getTopicPartitionInfo(anySet());
return spyTopicManager;
}
@@ -876,4 +993,58 @@ public final class AssignmentTestUtils {
mkEntry(UUID_7, mkMap(mkEntry("1", Optional.of(RACK_1))))
);
}
+
+ static RackAwareTaskAssignor getRackAwareTaskAssignor(final AssignmentConfigs configs) {
+ return spy(
+ new RackAwareTaskAssignor(
+ getClusterForAllTopics(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTaskChangelogMapForAllTasks(),
+ new HashMap<>(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManagerForChangelog(),
+ configs,
+ new MockTime()
+ )
+ );
+ }
+
+ static void verifyTaskPlacementWithRackAwareAssignor(final RackAwareTaskAssignor rackAwareTaskAssignor,
+ final Set<TaskId> allTaskIds,
+ final Map<UUID, ClientState> clientStates,
+ final boolean hasStandby,
+ final boolean enableRackAwareTaskAssignor) {
+ // Verifies active and standby are in different clients
+ verifyStandbySatisfyRackReplica(allTaskIds, rackAwareTaskAssignor.racksForProcess(), clientStates, null, true, null);
+
+ if (enableRackAwareTaskAssignor) {
+ verify(rackAwareTaskAssignor, times(2)).optimizeActiveTasks(any(), any(), anyInt(), anyInt());
+ verify(rackAwareTaskAssignor, hasStandby ? times(1) : never()).optimizeStandbyTasks(any(), anyInt(), anyInt(), any());
+ } else {
+ verify(rackAwareTaskAssignor, never()).optimizeActiveTasks(any(), any(), anyInt(), anyInt());
+ verify(rackAwareTaskAssignor, never()).optimizeStandbyTasks(any(), anyInt(), anyInt(), any());
+ }
+ }
+
+ static SortedMap<UUID, ClientState> copyClientStateMap(final Map<UUID, ClientState> originalMap) {
+ return new TreeMap<>(originalMap
+ .entrySet()
+ .stream()
+ .collect(
+ Collectors.toMap(
+ Entry::getKey,
+ entry -> new ClientState(entry.getValue())
+ )
+ )
+ );
+ }
+
+ static synchronized Random getRandom() {
+ if (random == null) {
+ final long seed = System.currentTimeMillis();
+ System.out.println("seed for getRandom: " + seed);
+ random = new Random(seed);
+ }
+ return random;
+ }
}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java
index 54df914cf53..341b98130b8 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java
@@ -51,6 +51,7 @@ import static org.apache.kafka.streams.processor.internals.assignment.Assignment
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_3;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasActiveTasks;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasStandbyTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.uuidForInt;
import static org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo.UNKNOWN_OFFSET_SUM;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.MatcherAssert.assertThat;
@@ -556,4 +557,17 @@ public class ClientStateTest {
assertEquals(UUID_3, new ClientState(UUID_3, 1, mkMap()).processId());
assertNull(new ClientState().processId());
}
+
+ @Test
+ public void shouldCopyState() {
+ final ClientState clientState = new ClientState(mkSet(new TaskId(0, 0)), mkSet(new TaskId(0, 1)), Collections.emptyMap(), EMPTY_CLIENT_TAGS, 1, uuidForInt(1));
+ final ClientState clientStateCopy = new ClientState(clientState);
+
+ assertEquals(clientStateCopy.processId(), clientState.processId());
+ assertEquals(clientStateCopy.capacity(), clientState.capacity());
+ assertEquals(clientStateCopy.prevActiveTasks(), clientStateCopy.prevActiveTasks());
+ assertEquals(clientStateCopy.prevStandbyTasks(), clientStateCopy.prevStandbyTasks());
+ assertThat(clientStateCopy.prevActiveTasks(), equalTo(clientState.prevActiveTasks()));
+ assertThat(clientStateCopy.prevStandbyTasks(), equalTo(clientState.prevStandbyTasks()));
+ }
}
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java
index 4acfe6a47fd..5828c0eabcb 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/HighAvailabilityTaskAssignorTest.java
@@ -17,12 +17,20 @@
package org.apache.kafka.streams.processor.internals.assignment;
import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
import java.util.SortedMap;
import java.util.SortedSet;
+import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.InternalTopicManager;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
+import org.hamcrest.Matchers;
import org.junit.Before;
import org.junit.Test;
@@ -65,22 +73,20 @@ import static org.apache.kafka.streams.processor.internals.assignment.Assignment
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertBalancedStatefulAssignment;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertBalancedTasks;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertValidAssignment;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.copyClientStateMap;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getClientStatesMap;
-import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getClusterForAllTopics;
-import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getProcessRacksForAllProcess;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRackAwareTaskAssignor;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomClientState;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomCluster;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomProcessRacks;
-import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskChangelogMapForAllTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomSubset;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMap;
-import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMapForAllTasks;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTopologyGroupTaskMap;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasActiveTasks;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasAssignedTasks;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasStandbyTasks;
-import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForChangelog;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForRandomChangelog;
-import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.verifyStandbySatisfyRackReplica;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.verifyTaskPlacementWithRackAwareAssignor;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.empty;
@@ -88,12 +94,7 @@ import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.fail;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
@RunWith(Parameterized.class)
public class HighAvailabilityTaskAssignorTest {
@@ -127,6 +128,8 @@ public class HighAvailabilityTaskAssignorTest {
);
}
+ private final Time time = new MockTime();
+
@Parameter
public boolean enableRackAwareTaskAssignor;
@@ -140,7 +143,7 @@ public class HighAvailabilityTaskAssignorTest {
}
@Parameterized.Parameters(name = "enableRackAwareTaskAssignor={0}")
- public static Collection<Object[]> getParamStoreType() {
+ public static Collection<Object[]> data() {
return asList(new Object[][] {
{true},
{false}
@@ -188,7 +191,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(unstable, is(true));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -229,7 +232,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(clientState3, hasAssignedTasks(6));
assertThat(unstable, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -270,7 +273,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedTasks(clientStates);
}
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -312,7 +315,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedTasks(clientStates);
}
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -349,7 +352,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTaskIds, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -394,7 +397,7 @@ public class HighAvailabilityTaskAssignorTest {
fail("Expected a skewed task assignment, but was: " + taskSkewReport);
}
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -432,7 +435,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTaskIds, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -474,7 +477,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedTasks(clientStates);
}
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -514,7 +517,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(notCaughtUpClientState2.standbyTaskCount(), greaterThanOrEqualTo(allTaskIds.size() / 3));
assertValidAssignment(0, allTaskIds.size() / 3 + 1, allTaskIds, emptySet(), clientStates, new StringBuilder());
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -567,7 +570,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTaskIds, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -605,7 +608,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(clientState1, hasActiveTasks(6));
assertThat(clientState2, hasActiveTasks(3));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTaskIds, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -632,7 +635,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -659,7 +662,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -689,7 +692,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -725,7 +728,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedStatefulAssignment(allTasks, clientStates, new StringBuilder());
assertBalancedTasks(clientStates);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -750,7 +753,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(client2.standbyTasks(), equalTo(mkSet(TASK_0_0)));
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -775,7 +778,7 @@ public class HighAvailabilityTaskAssignorTest {
assertHasNoStandbyTasks(client1, client2);
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -800,7 +803,7 @@ public class HighAvailabilityTaskAssignorTest {
assertHasNoActiveTasks(client2);
assertThat(probingRebalanceNeeded, is(true));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@@ -840,7 +843,7 @@ public class HighAvailabilityTaskAssignorTest {
assertHasNoActiveTasks(client2);
assertThat(probingRebalanceNeeded, is(true));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -878,7 +881,7 @@ public class HighAvailabilityTaskAssignorTest {
assertHasNoActiveTasks(client2);
assertThat(probingRebalanceNeeded, is(true));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -899,7 +902,7 @@ public class HighAvailabilityTaskAssignorTest {
assertHasNoStandbyTasks(client1);
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -919,7 +922,7 @@ public class HighAvailabilityTaskAssignorTest {
assertHasNoStandbyTasks(client1);
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -948,7 +951,7 @@ public class HighAvailabilityTaskAssignorTest {
);
assertThat(probingRebalanceNeeded, is(true));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -985,7 +988,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(probingRebalanceNeeded, is(true));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, true, enableRackAwareTaskAssignor);
}
@Test
@@ -1011,7 +1014,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(client3.activeTasks(), not(empty()));
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1035,7 +1038,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(client1.activeTasks(), equalTo(client1.prevActiveTasks()));
assertThat(client2.activeTasks(), equalTo(client2.prevActiveTasks()));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1055,7 +1058,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(probingRebalanceNeeded, is(false));
assertHasNoStandbyTasks(client1, client2);
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1074,7 +1077,7 @@ public class HighAvailabilityTaskAssignorTest {
assertThat(probingRebalanceNeeded, is(true));
assertThat(client2.standbyTaskCount(), equalTo(1));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1120,7 +1123,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedActiveAssignment(clientStates, new StringBuilder());
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1166,7 +1169,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedActiveAssignment(clientStates, new StringBuilder());
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1212,7 +1215,7 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedActiveAssignment(clientStates, new StringBuilder());
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
@@ -1258,51 +1261,148 @@ public class HighAvailabilityTaskAssignorTest {
assertBalancedActiveAssignment(clientStates, new StringBuilder());
assertThat(probingRebalanceNeeded, is(false));
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, allTasks, clientStates, false, enableRackAwareTaskAssignor);
}
@Test
public void shouldAssignRandomInput() {
final int nodeSize = 50;
final int tpSize = 60;
+ final int partitionSize = 3;
final int clientSize = 50;
final int replicaCount = 3;
final int maxCapacity = 3;
final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
- tpSize, false);
+ tpSize, partitionSize, false);
final AssignmentConfigs assignorConfiguration = getConfigWithStandbys(replicaCount);
final RackAwareTaskAssignor rackAwareTaskAssignor = spy(new RackAwareTaskAssignor(
- getRandomCluster(nodeSize, tpSize),
+ getRandomCluster(nodeSize, tpSize, partitionSize),
taskTopicPartitionMap,
- getTaskTopicPartitionMap(tpSize, true),
+ getTaskTopicPartitionMap(tpSize, partitionSize, true),
getTopologyGroupTaskMap(),
getRandomProcessRacks(clientSize, nodeSize),
- mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize),
- assignorConfiguration
+ mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize, partitionSize),
+ assignorConfiguration,
+ time
));
- final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
- tpSize, maxCapacity, false);
final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final List<Set<TaskId>> statefulAndStatelessTasks = getRandomSubset(taskIds, 2);
+ final Set<TaskId> statefulTasks = statefulAndStatelessTasks.get(0);
+ final Set<TaskId> statelessTasks = statefulAndStatelessTasks.get(1);
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
+ tpSize, maxCapacity, partitionSize, false, statefulTasks);
+
new HighAvailabilityTaskAssignor().assign(
clientStateMap,
taskIds,
- taskIds,
+ statefulTasks,
rackAwareTaskAssignor,
assignorConfiguration
);
assertValidAssignment(
replicaCount,
- taskIds,
- mkSet(),
+ statefulTasks,
+ statelessTasks,
clientStateMap,
new StringBuilder()
);
assertBalancedActiveAssignment(clientStateMap, new StringBuilder());
- verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, taskIds, clientStateMap, true);
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, taskIds, clientStateMap, true, enableRackAwareTaskAssignor);
+ }
+
+ @Test
+ public void shouldRemainOriginalAssignmentWithoutTrafficCost() {
+ // This test tests that if the traffic cost is 0, we should have same assignment with or without
+ // rack aware assignor enabled
+ final int nodeSize = 50;
+ final int tpSize = 60;
+ final int partitionSize = 3;
+ final int clientSize = 50;
+ final int replicaCount = 1;
+ final int maxCapacity = 3;
+ final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
+ tpSize, partitionSize, false);
+ final Cluster cluster = getRandomCluster(nodeSize, tpSize, partitionSize);
+ final Map<TaskId, Set<TopicPartition>> taskChangelogTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, true);
+ final Map<Subtopology, Set<TaskId>> subtopologySetMap = getTopologyGroupTaskMap();
+ final Map<UUID, Map<String, Optional<String>>> processRackMap = getRandomProcessRacks(clientSize, nodeSize);
+ final InternalTopicManager mockInternalTopicManager = mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize, partitionSize);
+
+ AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ replicaCount,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ 0,
+ 10,
+ rackAwareStrategy
+ );
+
+ RackAwareTaskAssignor rackAwareTaskAssignor = spy(new RackAwareTaskAssignor(
+ cluster,
+ taskTopicPartitionMap,
+ taskChangelogTopicPartitionMap,
+ subtopologySetMap,
+ processRackMap,
+ mockInternalTopicManager,
+ configs,
+ time
+ ));
+
+ final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final List<Set<TaskId>> statefulAndStatelessTasks = getRandomSubset(taskIds, 2);
+ final Set<TaskId> statefulTasks = statefulAndStatelessTasks.get(0);
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
+ tpSize, maxCapacity, partitionSize, false, statefulTasks);
+ final SortedMap<UUID, ClientState> clientStateMapCopy = copyClientStateMap(clientStateMap);
+
+ new HighAvailabilityTaskAssignor().assign(
+ clientStateMap,
+ taskIds,
+ statefulTasks,
+ rackAwareTaskAssignor,
+ configs
+ );
+
+ configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ replicaCount,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ 0,
+ 10,
+ StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE
+ );
+
+ rackAwareTaskAssignor = spy(new RackAwareTaskAssignor(
+ cluster,
+ taskTopicPartitionMap,
+ taskChangelogTopicPartitionMap,
+ subtopologySetMap,
+ processRackMap,
+ mockInternalTopicManager,
+ configs,
+ time
+ ));
+
+ new HighAvailabilityTaskAssignor().assign(
+ clientStateMapCopy,
+ taskIds,
+ statefulTasks,
+ rackAwareTaskAssignor,
+ configs
+ );
+
+ for (final Map.Entry<UUID, ClientState> entry : clientStateMap.entrySet()) {
+ assertThat(entry.getValue().statefulActiveTasks(), Matchers.equalTo(clientStateMapCopy.get(entry.getKey()).statefulActiveTasks()));
+ assertThat(entry.getValue().standbyTasks(), Matchers.equalTo(clientStateMapCopy.get(entry.getKey()).standbyTasks()));
+ }
}
private static void assertHasNoActiveTasks(final ClientState... clients) {
@@ -1334,33 +1434,5 @@ public class HighAvailabilityTaskAssignorTest {
return new ClientState(statefulActiveTasks, emptySet(), taskLags, EMPTY_CLIENT_TAGS, 1, processId);
}
- private static RackAwareTaskAssignor getRackAwareTaskAssignor(final AssignmentConfigs configs) {
- return spy(
- new RackAwareTaskAssignor(
- getClusterForAllTopics(),
- getTaskTopicPartitionMapForAllTasks(),
- getTaskChangelogMapForAllTasks(),
- new HashMap<>(),
- getProcessRacksForAllProcess(),
- mockInternalTopicManagerForChangelog(),
- configs
- )
- );
- }
-
- private void verifyTaskPlacementWithRackAwareAssignor(final RackAwareTaskAssignor rackAwareTaskAssignor,
- final Set<TaskId> allTaskIds,
- final Map<UUID, ClientState> clientStates,
- final boolean hasStandby) {
- // Verifies active and standby are in different clients
- verifyStandbySatisfyRackReplica(allTaskIds, rackAwareTaskAssignor.racksForProcess(), clientStates, null, true, null);
- if (enableRackAwareTaskAssignor) {
- verify(rackAwareTaskAssignor, times(2)).optimizeActiveTasks(any(), any(), anyInt(), anyInt());
- verify(rackAwareTaskAssignor, hasStandby ? times(1) : never()).optimizeStandbyTasks(any(), anyInt(), anyInt(), any());
- } else {
- verify(rackAwareTaskAssignor, never()).optimizeActiveTasks(any(), any(), anyInt(), anyInt());
- verify(rackAwareTaskAssignor, never()).optimizeStandbyTasks(any(), anyInt(), anyInt(), any());
- }
- }
}
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 55b0e3d0038..58dd4313dbd 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
@@ -179,7 +179,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
- getRackAwareDisabledConfig()
+ getRackAwareDisabledConfig(),
+ time
));
// False since partitionWithoutInfo10 is missing in cluster metadata
@@ -197,7 +198,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
));
// False since partitionWithoutInfo10 is missing in cluster metadata
@@ -216,7 +218,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
));
// False since nodeMissingRack has one node which doesn't have rack
@@ -235,7 +238,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(true),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
// False since process1 doesn't have rackId
assertFalse(assignor.validClientRack());
@@ -250,7 +254,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessWithNoConsumerRacks(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
// False since process1 doesn't have rackId
assertFalse(assignor.validClientRack());
@@ -266,7 +271,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final Map<UUID, String> racksForProcess = assignor.racksForProcess();
assertEquals(mkMap(mkEntry(UUID_1, RACK_1)), racksForProcess);
@@ -288,7 +294,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
processRacks,
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
assertFalse(assignor.validClientRack());
@@ -303,7 +310,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
// partitionWithoutInfo00 has rackInfo in cluster metadata
@@ -319,7 +327,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
));
// partitionWithoutInfo00 has rackInfo in cluster metadata
@@ -351,7 +360,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
assertTrue(assignor.canEnableRackAwareAssignor());
@@ -385,7 +395,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
- getRackAwareEnabledConfigWithStandby(1)
+ getRackAwareEnabledConfigWithStandby(1),
+ time
));
assertTrue(assignor.canEnableRackAwareAssignor());
@@ -422,7 +433,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
- getRackAwareEnabledConfigWithStandby(1)
+ getRackAwareEnabledConfigWithStandby(1),
+ time
));
assertFalse(assignor.canEnableRackAwareAssignor());
@@ -443,7 +455,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
));
assertFalse(assignor.canEnableRackAwareAssignor());
@@ -461,7 +474,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -492,7 +506,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -533,20 +548,22 @@ public class RackAwareTaskAssignorTest {
public void shouldOptimizeRandomActive() {
final int nodeSize = 30;
final int tpSize = 40;
+ final int partitionSize = 3;
final int clientSize = 30;
- final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, false);
+ final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, false);
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
- getRandomCluster(nodeSize, tpSize),
+ getRandomCluster(nodeSize, tpSize, partitionSize),
taskTopicPartitionMap,
mkMap(),
getTopologyGroupTaskMap(),
getRandomProcessRacks(clientSize, nodeSize),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
- final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize, 1);
final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize, partitionSize, 1, taskIds);
final Map<UUID, Integer> clientTaskCount = clientTaskCount(clientStateMap, ClientState::activeTaskCount);
@@ -564,20 +581,22 @@ public class RackAwareTaskAssignorTest {
public void shouldMaintainOriginalAssignment() {
final int nodeSize = 20;
final int tpSize = 40;
+ final int partitionSize = 3;
final int clientSize = 30;
- final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, false);
+ final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, false);
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
- getRandomCluster(nodeSize, tpSize),
+ getRandomCluster(nodeSize, tpSize, partitionSize),
taskTopicPartitionMap,
mkMap(),
getTopologyGroupTaskMap(),
getRandomProcessRacks(clientSize, nodeSize),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
- final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize, 1);
final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize, partitionSize, 1, taskIds);
final Map<TaskId, UUID> taskClientMap = new HashMap<>();
for (final Entry<UUID, ClientState> entry : clientStateMap.entrySet()) {
@@ -609,7 +628,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -654,7 +674,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -699,7 +720,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -741,7 +763,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -770,7 +793,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -801,7 +825,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManager,
- getRackAwareEnabledConfig()
+ getRackAwareEnabledConfig(),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
@@ -831,7 +856,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManagerForChangelog(),
- getRackAwareEnabledConfigWithStandby(1)
+ getRackAwareEnabledConfigWithStandby(1),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1, UUID_1);
@@ -866,7 +892,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManagerForChangelog(),
- getRackAwareEnabledConfigWithStandby(replicaCount)
+ getRackAwareEnabledConfigWithStandby(replicaCount),
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1, UUID_1);
@@ -928,7 +955,8 @@ public class RackAwareTaskAssignorTest {
getTopologyGroupTaskMap(),
getProcessRacksForAllProcess(),
mockInternalTopicManagerForChangelog(),
- assignorConfiguration
+ assignorConfiguration,
+ time
);
final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1, UUID_1);
@@ -983,26 +1011,28 @@ public class RackAwareTaskAssignorTest {
public void shouldOptimizeRandomStandby() {
final int nodeSize = 50;
final int tpSize = 60;
+ final int partionSize = 3;
final int clientSize = 50;
final int replicaCount = 3;
final int maxCapacity = 3;
final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
- tpSize, false);
+ tpSize, partionSize, false);
final AssignmentConfigs assignorConfiguration = getRackAwareEnabledConfigWithStandby(replicaCount);
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
- getRandomCluster(nodeSize, tpSize),
+ getRandomCluster(nodeSize, tpSize, partionSize),
taskTopicPartitionMap,
- getTaskTopicPartitionMap(tpSize, true),
+ getTaskTopicPartitionMap(tpSize, partionSize, true),
getTopologyGroupTaskMap(),
getRandomProcessRacks(clientSize, nodeSize),
- mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize),
- assignorConfiguration
+ mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize, partionSize),
+ assignorConfiguration,
+ time
);
- final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
- tpSize, maxCapacity);
final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
+ tpSize, partionSize, maxCapacity, taskIds);
final StandbyTaskAssignor standbyTaskAssignor = StandbyTaskAssignorFactory.create(
assignorConfiguration, assignor);
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignorTest.java
index e94f3aee8d7..b0049e5252f 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignorTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/StickyTaskAssignorTest.java
@@ -16,7 +16,20 @@
*/
package org.apache.kafka.streams.processor.internals.assignment;
+import java.util.Collection;
+import java.util.Optional;
+import java.util.SortedMap;
+import java.util.SortedSet;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.InternalTopicManager;
+import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
+import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentConfigs;
+import org.junit.Before;
import org.junit.Test;
import java.util.ArrayList;
@@ -29,6 +42,9 @@ import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.UUID;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
import static java.util.Arrays.asList;
import static java.util.Collections.singleton;
@@ -49,12 +65,31 @@ import static org.apache.kafka.streams.processor.internals.assignment.Assignment
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_1;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_2;
import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_3;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_3_0;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_3_1;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_3_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.UUID_6;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertValidAssignment;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.copyClientStateMap;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getClusterForAllTopics;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getProcessRacksForAllProcess;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRackAwareTaskAssignor;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomClientState;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomCluster;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomProcessRacks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomSubset;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskChangelogMapForAllTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMap;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMapForAllTasks;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTopologyGroupTaskMap;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForChangelog;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForRandomChangelog;
+import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.verifyTaskPlacementWithRackAwareAssignor;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.empty;
import static org.hamcrest.Matchers.equalTo;
@@ -64,13 +99,36 @@ import static org.hamcrest.Matchers.hasItems;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.Matchers.not;
+import static org.junit.Assert.assertFalse;
+import static org.mockito.Mockito.spy;
+@RunWith(Parameterized.class)
public class StickyTaskAssignorTest {
private final List<Integer> expectedTopicGroupIds = asList(1, 2);
-
+ private final Time time = new MockTime();
private final Map<UUID, ClientState> clients = new TreeMap<>();
+ @Parameter
+ public boolean enableRackAwareTaskAssignor;
+
+ private String rackAwareStrategy = StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE;
+
+ @Before
+ public void setUp() {
+ if (enableRackAwareTaskAssignor) {
+ rackAwareStrategy = StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC;
+ }
+ }
+
+ @Parameterized.Parameters(name = "enableRackAwareTaskAssignor={0}")
+ public static Collection<Object[]> data() {
+ return asList(new Object[][] {
+ {true},
+ {false}
+ });
+ }
+
@Test
public void shouldAssignOneActiveTaskToEachProcessWhenTaskCountSameAsProcessCount() {
createClient(UUID_1, 1);
@@ -361,15 +419,15 @@ public class StickyTaskAssignorTest {
TASK_0_0,
TASK_0_1,
TASK_0_2,
- new TaskId(1, 0),
- new TaskId(1, 1),
- new TaskId(1, 2),
- new TaskId(2, 0),
- new TaskId(2, 1),
- new TaskId(2, 2),
- new TaskId(3, 0),
- new TaskId(3, 1),
- new TaskId(3, 2)
+ TASK_1_0,
+ TASK_1_1,
+ TASK_1_2,
+ TASK_2_0,
+ TASK_2_1,
+ TASK_2_2,
+ TASK_3_0,
+ TASK_3_1,
+ TASK_3_2
);
assertThat(probingRebalanceNeeded, is(false));
@@ -387,7 +445,7 @@ public class StickyTaskAssignorTest {
final List<TaskId> taskIds = new ArrayList<>();
final TaskId[] taskIdArray = new TaskId[16];
- for (int i = 1; i <= 2; i++) {
+ for (int i = 0; i < 2; i++) {
for (int j = 0; j < 8; j++) {
taskIds.add(new TaskId(i, j));
}
@@ -396,7 +454,38 @@ public class StickyTaskAssignorTest {
Collections.shuffle(taskIds);
taskIds.toArray(taskIdArray);
- final boolean probingRebalanceNeeded = assign(taskIdArray);
+ final int nodeSize = 5;
+ final int topicSize = 2;
+ final int partitionSize = 8;
+ final int clientSize = 4;
+ final Cluster cluster = getRandomCluster(nodeSize, topicSize, partitionSize);
+ final Map<TaskId, Set<TopicPartition>> partitionsForTask = getTaskTopicPartitionMap(topicSize, partitionSize, false);
+ final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask = getTaskTopicPartitionMap(topicSize, partitionSize, true);
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = new HashMap<>();
+ final Map<UUID, Map<String, Optional<String>>> racksForProcessConsumer = getRandomProcessRacks(clientSize, nodeSize);
+ final InternalTopicManager internalTopicManager = mockInternalTopicManagerForRandomChangelog(nodeSize, topicSize, partitionSize);
+ final AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ 1,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ null,
+ null,
+ rackAwareStrategy
+ );
+ final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(
+ cluster,
+ partitionsForTask,
+ changelogPartitionsForTask,
+ tasksForTopicGroup,
+ racksForProcessConsumer,
+ internalTopicManager,
+ configs,
+ time
+ );
+
+ final boolean probingRebalanceNeeded = assign(configs, rackAwareTaskAssignor, taskIdArray);
assertThat(probingRebalanceNeeded, is(false));
Collections.sort(taskIds);
@@ -482,7 +571,6 @@ public class StickyTaskAssignorTest {
not(equalTo(taskIds)));
}
}
-
}
}
@@ -673,12 +761,45 @@ public class StickyTaskAssignorTest {
final List<TaskId> taskIds = asList(TASK_0_0, TASK_0_1, TASK_0_2);
Collections.shuffle(taskIds);
+
+ final int nodeSize = 5;
+ final int topicSize = 1;
+ final int partitionSize = 3;
+ final int clientSize = 2;
+ final Cluster cluster = getRandomCluster(nodeSize, topicSize, partitionSize);
+ final Map<TaskId, Set<TopicPartition>> partitionsForTask = getTaskTopicPartitionMap(topicSize, partitionSize, false);
+ final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask = getTaskTopicPartitionMap(topicSize, partitionSize, true);
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = new HashMap<>();
+ final Map<UUID, Map<String, Optional<String>>> racksForProcessConsumer = getRandomProcessRacks(clientSize, nodeSize);
+ final InternalTopicManager internalTopicManager = mockInternalTopicManagerForRandomChangelog(nodeSize, topicSize, partitionSize);
+
+ final AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ 0,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ null,
+ null,
+ rackAwareStrategy
+ );
+ final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(
+ cluster,
+ partitionsForTask,
+ changelogPartitionsForTask,
+ tasksForTopicGroup,
+ racksForProcessConsumer,
+ internalTopicManager,
+ configs,
+ time
+ );
+
final boolean probingRebalanceNeeded = new StickyTaskAssignor(true).assign(
clients,
new HashSet<>(taskIds),
new HashSet<>(taskIds),
- null,
- new AssignorConfiguration.AssignmentConfigs(0L, 1, 0, 60_000L, EMPTY_RACK_AWARE_ASSIGNMENT_TAGS)
+ rackAwareTaskAssignor,
+ configs
);
assertThat(probingRebalanceNeeded, is(false));
@@ -686,19 +807,257 @@ public class StickyTaskAssignorTest {
assertThat(c2.activeTasks(), empty());
}
+ @Test
+ public void shouldOptimizeStatefulAndStatelessTaskTraffic() {
+ final ClientState c1 = createClientWithPreviousActiveTasks(UUID_1, 1, TASK_0_0, TASK_0_1, TASK_0_2);
+ final ClientState c2 = createClientWithPreviousActiveTasks(UUID_2, 1, TASK_1_0, TASK_1_1, TASK_0_3, TASK_1_3);
+ final ClientState c3 = createClientWithPreviousActiveTasks(UUID_3, 1, TASK_1_2);
+
+ final List<TaskId> taskIds = asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3);
+ final List<TaskId> statefulTaskIds = asList(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
+ Collections.shuffle(taskIds);
+
+ final Cluster cluster = getClusterForAllTopics();
+ final Map<TaskId, Set<TopicPartition>> partitionsForTask = getTaskTopicPartitionMapForAllTasks();
+ final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask = getTaskChangelogMapForAllTasks();
+ final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = new HashMap<>();
+ final Map<UUID, Map<String, Optional<String>>> racksForProcessConsumer = getProcessRacksForAllProcess();
+ final InternalTopicManager internalTopicManager = mockInternalTopicManagerForChangelog();
+
+ final AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ 1,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ 10,
+ 1,
+ rackAwareStrategy
+ );
+ final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(
+ cluster,
+ partitionsForTask,
+ changelogPartitionsForTask,
+ tasksForTopicGroup,
+ racksForProcessConsumer,
+ internalTopicManager,
+ configs,
+ time
+ );
+
+ final boolean probingRebalanceNeeded = new StickyTaskAssignor().assign(
+ clients,
+ new HashSet<>(taskIds),
+ new HashSet<>(statefulTaskIds),
+ rackAwareTaskAssignor,
+ configs
+ );
+ assertThat(probingRebalanceNeeded, is(false));
+
+ if (enableRackAwareTaskAssignor) {
+ // Total cost for active stateful: 3
+ // Total cost for active stateless: 0
+ // Total cost for standby: 20
+ assertThat(c1.activeTasks(), equalTo(mkSet(TASK_0_3, TASK_1_0, TASK_1_2)));
+ assertThat(c1.standbyTasks(), equalTo(mkSet(TASK_0_0, TASK_0_1)));
+ assertThat(c2.activeTasks(), equalTo(mkSet(TASK_0_0, TASK_0_2, TASK_1_1)));
+ assertThat(c2.standbyTasks(), empty());
+ assertThat(c3.activeTasks(), equalTo(mkSet(TASK_0_1, TASK_1_3)));
+ assertThat(c3.standbyTasks(), equalTo(mkSet(TASK_1_0, TASK_1_1)));
+ } else {
+ // Total cost for active stateful: 30
+ // Total cost for active stateless: 40
+ // Total cost for standby: 10
+ assertThat(c1.activeTasks(), equalTo(mkSet(TASK_0_1, TASK_0_2, TASK_1_3)));
+ assertThat(c1.standbyTasks(), equalTo(mkSet(TASK_0_0)));
+ assertThat(c2.activeTasks(), equalTo(mkSet(TASK_0_3, TASK_1_0, TASK_1_1)));
+ assertThat(c2.standbyTasks(), equalTo(mkSet(TASK_0_1)));
+ assertThat(c3.activeTasks(), equalTo(mkSet(TASK_0_0, TASK_1_2)));
+ assertThat(c3.standbyTasks(), equalTo(mkSet(TASK_1_0, TASK_1_1)));
+ }
+
+ }
+
+ @Test
+ public void shouldAssignRandomInput() {
+ final int nodeSize = 50;
+ final int tpSize = 60;
+ final int partitionSize = 3;
+ final int clientSize = 50;
+ final int replicaCount = 1;
+ final int maxCapacity = 3;
+ final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
+ tpSize, partitionSize, false);
+ final AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ replicaCount,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ 10,
+ 1,
+ rackAwareStrategy
+ );
+
+ final RackAwareTaskAssignor rackAwareTaskAssignor = spy(new RackAwareTaskAssignor(
+ getRandomCluster(nodeSize, tpSize, partitionSize),
+ taskTopicPartitionMap,
+ getTaskTopicPartitionMap(tpSize, partitionSize, true),
+ getTopologyGroupTaskMap(),
+ getRandomProcessRacks(clientSize, nodeSize),
+ mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize, partitionSize),
+ configs,
+ time
+ ));
+
+ final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final List<Set<TaskId>> statefulAndStatelessTasks = getRandomSubset(taskIds, 2);
+ final Set<TaskId> statefulTasks = statefulAndStatelessTasks.get(0);
+ final Set<TaskId> statelessTasks = statefulAndStatelessTasks.get(1);
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
+ tpSize, maxCapacity, partitionSize, false, statefulTasks);
+
+
+ final boolean probing = new StickyTaskAssignor().assign(
+ clientStateMap,
+ taskIds,
+ statefulTasks,
+ rackAwareTaskAssignor,
+ configs
+ );
+
+ assertFalse(probing);
+ assertValidAssignment(
+ replicaCount,
+ statefulTasks,
+ statelessTasks,
+ clientStateMap,
+ new StringBuilder()
+ );
+ verifyTaskPlacementWithRackAwareAssignor(rackAwareTaskAssignor, taskIds, clientStateMap, true, enableRackAwareTaskAssignor);
+ }
+
+ @Test
+ public void shouldRemainOriginalAssignmentWithoutTrafficCost() {
+ // This test tests that if the traffic cost is 0, we should have same assignment with or without
+ // rack aware assignor enabled
+ final int nodeSize = 50;
+ final int tpSize = 60;
+ final int partitionSize = 3;
+ final int clientSize = 50;
+ final int replicaCount = 1;
+ final int maxCapacity = 3;
+ final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
+ tpSize, partitionSize, false);
+ final Cluster cluster = getRandomCluster(nodeSize, tpSize, partitionSize);
+ final Map<TaskId, Set<TopicPartition>> taskChangelogTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, true);
+ final Map<Subtopology, Set<TaskId>> subtopologySetMap = getTopologyGroupTaskMap();
+ final Map<UUID, Map<String, Optional<String>>> processRackMap = getRandomProcessRacks(clientSize, nodeSize);
+ final InternalTopicManager mockInternalTopicManager = mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize, partitionSize);
+
+ AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ replicaCount,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ 0, // Override traffic cost to 0 to maintain original assignment
+ 10,
+ rackAwareStrategy
+ );
+
+ RackAwareTaskAssignor rackAwareTaskAssignor = spy(new RackAwareTaskAssignor(
+ cluster,
+ taskTopicPartitionMap,
+ taskChangelogTopicPartitionMap,
+ subtopologySetMap,
+ processRackMap,
+ mockInternalTopicManager,
+ configs,
+ time
+ ));
+
+ final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
+ final List<Set<TaskId>> statefulAndStatelessTasks = getRandomSubset(taskIds, 2);
+ final Set<TaskId> statefulTasks = statefulAndStatelessTasks.get(0);
+ final SortedMap<UUID, ClientState> clientStateMap = getRandomClientState(clientSize,
+ tpSize, maxCapacity, partitionSize, false, statefulTasks);
+ final SortedMap<UUID, ClientState> clientStateMapCopy = copyClientStateMap(clientStateMap);
+
+ new StickyTaskAssignor().assign(
+ clientStateMap,
+ taskIds,
+ statefulTasks,
+ rackAwareTaskAssignor,
+ configs
+ );
+
+ configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ replicaCount,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ 0,
+ 10,
+ StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE
+ );
+
+ rackAwareTaskAssignor = spy(new RackAwareTaskAssignor(
+ cluster,
+ taskTopicPartitionMap,
+ taskChangelogTopicPartitionMap,
+ subtopologySetMap,
+ processRackMap,
+ mockInternalTopicManager,
+ configs,
+ time
+ ));
+
+ new StickyTaskAssignor().assign(
+ clientStateMapCopy,
+ taskIds,
+ statefulTasks,
+ rackAwareTaskAssignor,
+ configs
+ );
+
+ for (final Map.Entry<UUID, ClientState> entry : clientStateMap.entrySet()) {
+ assertThat(entry.getValue().statefulActiveTasks(), equalTo(clientStateMapCopy.get(entry.getKey()).statefulActiveTasks()));
+ assertThat(entry.getValue().standbyTasks(), equalTo(clientStateMapCopy.get(entry.getKey()).standbyTasks()));
+ }
+ }
+
private boolean assign(final TaskId... tasks) {
return assign(0, tasks);
}
private boolean assign(final int numStandbys, final TaskId... tasks) {
+ final List<TaskId> taskIds = asList(tasks);
+ Collections.shuffle(taskIds);
+ final AssignmentConfigs configs = new AssignorConfiguration.AssignmentConfigs(
+ 0L,
+ 1,
+ numStandbys,
+ 60_000L,
+ EMPTY_RACK_AWARE_ASSIGNMENT_TAGS,
+ null,
+ null,
+ rackAwareStrategy
+ );
+
+ return assign(configs, getRackAwareTaskAssignor(configs), tasks);
+ }
+
+ private boolean assign(final AssignmentConfigs configs, final RackAwareTaskAssignor rackAwareTaskAssignor, final TaskId... tasks) {
final List<TaskId> taskIds = asList(tasks);
Collections.shuffle(taskIds);
return new StickyTaskAssignor().assign(
clients,
new HashSet<>(taskIds),
new HashSet<>(taskIds),
- null,
- new AssignorConfiguration.AssignmentConfigs(0L, 1, numStandbys, 60_000L, EMPTY_RACK_AWARE_ASSIGNMENT_TAGS)
+ rackAwareTaskAssignor,
+ configs
);
}
@@ -725,7 +1084,7 @@ public class StickyTaskAssignorTest {
}
private ClientState createClientWithPreviousActiveTasks(final UUID processId, final int capacity, final TaskId... taskIds) {
- final ClientState clientState = new ClientState(capacity);
+ final ClientState clientState = new ClientState(processId, capacity);
clientState.addPreviousActiveTasks(mkSet(taskIds));
clients.put(processId, clientState);
return clientState;
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java
index dca25fcbb17..19f865f8cf7 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java
@@ -30,6 +30,7 @@ import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;
import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.internals.InternalTopicManager;
@@ -76,6 +77,7 @@ import static org.mockito.Mockito.spy;
@RunWith(Parameterized.class)
public class TaskAssignorConvergenceTest {
private static Random random;
+ private static final Time TIME = new MockTime();
@BeforeClass
public static void beforeClass() {
@@ -122,7 +124,7 @@ public class TaskAssignorConvergenceTest {
statelessTasks.add(taskId);
remainingStatelessTasks--;
- final Node[] replica = getRandomReplica(nodes, nodeIndex);
+ final Node[] replica = getRandomReplica(nodes, nodeIndex, i);
partitionInfoSet.add(new PartitionInfo(TOPIC_PREFIX + "_" + subtopology, i, replica[0], replica, replica));
nodeIndex++;
@@ -145,7 +147,7 @@ public class TaskAssignorConvergenceTest {
statefulTaskEndOffsetSums.put(taskId, 150000L);
remainingStatefulTasks--;
- Node[] replica = getRandomReplica(nodes, nodeIndex);
+ Node[] replica = getRandomReplica(nodes, nodeIndex, i);
partitionInfoSet.add(new PartitionInfo(TOPIC_PREFIX + "_" + subtopology, i, replica[0], replica, replica));
nodeIndex++;
@@ -154,7 +156,7 @@ public class TaskAssignorConvergenceTest {
tasksForTopicGroup.computeIfAbsent(new Subtopology(subtopology, null), k -> new HashSet<>()).add(taskId);
final int changelogNodeIndex = random.nextInt(nodes.size());
- replica = getRandomReplica(nodes, changelogNodeIndex);
+ replica = getRandomReplica(nodes, changelogNodeIndex, i);
final TopicPartitionInfo info = new TopicPartitionInfo(i, replica[0], Arrays.asList(replica[0], replica[1]), Collections.emptyList());
topicPartitionInfo.computeIfAbsent(changelogTopicName, tp -> new ArrayList<>()).add(info);
}
@@ -596,7 +598,8 @@ public class TaskAssignorConvergenceTest {
harness.tasksForTopicGroup,
harness.racksForProcessConsumer,
harness.internalTopicManager,
- configs
+ configs,
+ TIME
);
while (rebalancePending && iteration < iterationLimit) {
iteration++;