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 2026/09/18 17:49:02 UTC
(kafka) branch trunk updated: KAFKA-20665: extend refiner to compute patch-set and intermediate assignment (#23484)
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 1f19eac23ed KAFKA-20665: extend refiner to compute patch-set and intermediate assignment (#23484)
1f19eac23ed is described below
commit 1f19eac23ed292c908616bf468c9c59288b5f31d
Author: Matthias J. Sax <ma...@confluent.io>
AuthorDate: Fri Sep 18 10:48:54 2026 -0700
KAFKA-20665: extend refiner to compute patch-set and intermediate assignment (#23484)
Part of KIP-1071.
Reviewers: Lucas Brutschy <lb...@confluent.io>
---
.../group/streams/AssignmentRefinerImpl.java | 301 +++++++++-
.../group/streams/AssignmentRefinerTest.java | 640 ++++++++++++++++++++-
2 files changed, 937 insertions(+), 4 deletions(-)
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerImpl.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerImpl.java
index ccf339313a9..65ce8d9d598 100644
--- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerImpl.java
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerImpl.java
@@ -22,6 +22,7 @@ import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -30,6 +31,7 @@ import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
+import java.util.function.BinaryOperator;
import java.util.function.Consumer;
/**
@@ -38,7 +40,7 @@ import java.util.function.Consumer;
* state.
*
* <p>{@link #refine} returns the target assignment unchanged, like {@link NoOpAssignmentRefiner} for now,
- * because this class is WIP is not used yet.
+ * because this class is still a work in progress and not used yet.
*/
public class AssignmentRefinerImpl implements AssignmentRefiner {
@@ -218,7 +220,7 @@ public class AssignmentRefinerImpl implements AssignmentRefiner {
// Because of assignment offloading and member fencing, the target assignment could contain a member which
// was removed from the group in the meantime. For this case, all previously owned tasks of this member
- // (which did not get move to a new owner) will be "dandling" which will be fixed by the next assignor run.
+ // (which were not moved to a new owner) will be dangling, which the next assignor run fixes.
// Furthermore, we stage all tasks the assignor moves to this member on their old owners to keep them
// "online".
final StreamsGroupMember targetMember = members.get(targetOwner);
@@ -289,7 +291,7 @@ public class AssignmentRefinerImpl implements AssignmentRefiner {
/**
* Decides which of the staged migrations get a warm-up task, under the warmup budget.
*
- * <p>There is different scenarios:
+ * <p>There are several scenarios:
* <ul>
* <li>A warm-up task already restoring keeps its warm-up slot if the target assignment didn't change, and the
* warm-up task is not caught up yet. It could also get revoked if the warmup budget was reduced and keeping
@@ -521,6 +523,217 @@ public class AssignmentRefinerImpl implements AssignmentRefiner {
.loadWith(newWarmupsByProcess.getOrDefault(candidate.currentProcessId(), 0));
}
+ /**
+ * Decides which of the target assignment's standby placements this step has to hold back.
+ *
+ * <p>Nothing is invented or dropped permanently: every placement comes from the target assignment, and one held
+ * back here is emitted by a later step once its reason is gone. A placement of task {@code t} on member {@code m}
+ * of process {@code p} is withheld when:
+ * <ol>
+ * <li>a staged migration keeps {@code t} running as an active task on {@code p}, and a process cannot hold
+ * {@code t} twice.</li>
+ * <li>{@code t}'s migration onto {@code p} borrowed an existing standby on {@code p}. To not run
+ * {@code num.standby.replicas + 1} standbys, we hold back one assignment of the standby to a new owner.</li>
+ * </ol>
+ *
+ * @param targetAssignment
+ * All members' target assignments, as computed by the task assignor.
+ * @param currentAssignment
+ * The indexed current assignment, from {@link #indexCurrentAssignment}.
+ * @param decisions
+ * What the case analysis decided, from {@link #analyzeTasks}.
+ * @param warmupPlan
+ * How each staged migration is being warmed, from {@link #planWarmups}.
+ * @param members
+ * All members of the group, used to resolve which process a member runs in.
+ * @param subtopologies
+ * The resolved subtopologies, which tell whether a subtopology is stateful.
+ *
+ * @return The standby placements to withhold, as the tasks to drop from each member's target assignment, in
+ * canonical order. A member with nothing withheld does not appear.
+ */
+ static SortedMap<String, SortedSet<TaskId>> filterStandbys(
+ final Map<String, TasksTuple> targetAssignment,
+ final CurrentAssignmentIndex currentAssignment,
+ final TaskDecisions decisions,
+ final WarmupPlan warmupPlan,
+ final Map<String, StreamsGroupMember> members,
+ final SortedMap<String, ConfiguredSubtopology> subtopologies
+ ) {
+ final StandbyConflicts conflicts = indexStandbyConflicts(
+ targetAssignment,
+ currentAssignment,
+ decisions,
+ warmupPlan,
+ members,
+ subtopologies
+ );
+ final SortedMap<String, SortedSet<TaskId>> withheld = new TreeMap<>();
+
+ targetAssignment.forEach((memberId, tasks) -> {
+ final StreamsGroupMember member = members.get(memberId);
+ if (member == null) {
+ // The target assignment can name a member the group has already removed. Its tasks reach nobody, so
+ // there is nothing to hold back and no process to resolve it against.
+ return;
+ }
+
+ forEachStatefulTask(tasks.standbyTasks(), subtopologies, task -> {
+ if (isStandbyWithheld(memberId, member.processId(), task, conflicts)) {
+ withheld.computeIfAbsent(memberId, __ -> new TreeSet<>()).add(task);
+ }
+ });
+ });
+
+ return Collections.unmodifiableSortedMap(withheld);
+ }
+
+ /**
+ * Builds the {@link StandbyConflicts} lookups, ie, each standby task which cannot be relocated yet, because
+ * (1) its target process owns the corresponding active task from a staged migration, or (2) the standby must stay
+ * on its current owner to fulfill its role of being borrowed.
+ */
+ private static StandbyConflicts indexStandbyConflicts(
+ final Map<String, TasksTuple> targetAssignment,
+ final CurrentAssignmentIndex currentAssignment,
+ final TaskDecisions decisions,
+ final WarmupPlan warmupPlan,
+ final Map<String, StreamsGroupMember> members,
+ final SortedMap<String, ConfiguredSubtopology> subtopologies
+ ) {
+ // Conflicts from staged migrations
+ final Map<TaskId, String> activeStagedOnProcess = new HashMap<>();
+ for (final StagedMigration migration : decisions.stagedMigrations()) {
+ activeStagedOnProcess.put(migration.task(), members.get(migration.currentOwner()).processId());
+ }
+
+ // Conflicts from borrows
+ //
+ // If we have more than one standby (ie, `num.standby.replicas >= 2`), we need to ensure to only hold back
+ // one standby task migration (for the single borrow of the active task migration)
+ // If there is a parallel sibling-move for the same standby task on a different process we leave it alone;
+ // we are looking for a process with a new standby being assigned to (there might be multiple, so we pick one
+ // deterministically based on memberId order)
+ final BinaryOperator<String> firstInMemberOrder = (left, right) -> left.compareTo(right) <= 0 ? left : right;
+ final Map<TaskId, String> standbyKeptOnMember = new HashMap<>();
+ final Set<TaskId> alreadyUndelivered = new HashSet<>();
+ targetAssignment.forEach((targetMemberId, tasks) ->
+ forEachStatefulTask(tasks.standbyTasks(), subtopologies, task -> {
+ if (warmupPlan.borrowedMigrations().contains(task)) {
+ final StreamsGroupMember targetMember = members.get(targetMemberId);
+ final String currentOwnerProcessId = activeStagedOnProcess.get(task);
+ if (targetMember == null // dropped out of the group
+ || targetMember.processId().equals(currentOwnerProcessId)) { // stage migration; tracked above
+ alreadyUndelivered.add(task);
+ } else if (findCopyOnProcess(currentAssignment, task, targetMember.processId()).isEmpty()) { // only withhold if not a sibling move
+ standbyKeptOnMember.merge(task, targetMemberId, firstInMemberOrder);
+ }
+ }
+ }));
+ alreadyUndelivered.forEach(standbyKeptOnMember::remove);
+
+ return new StandbyConflicts(activeStagedOnProcess, standbyKeptOnMember);
+ }
+
+ /**
+ * Whether this step has to hold the standby placement back.
+ */
+ private static boolean isStandbyWithheld(
+ final String memberId,
+ final String processId,
+ final TaskId task,
+ final StandbyConflicts conflicts
+ ) {
+ // Rule 1: a process cannot hold `task` twice, so a standby on the process a staged migration keeps the
+ // active running on waits for that migration to complete.
+ if (processId.equals(conflicts.activeStagedOnProcess().get(task))) {
+ return true;
+ }
+
+ // Rule 2: a borrowed migration keeps its existing copy as one of the group's entitled replicas, so one
+ // relocated placement of that standby waits, which keeps the group at `num.standby.replicas`.
+ return memberId.equals(conflicts.standbyKeptOnMember().get(task));
+ }
+
+ /**
+ * Builds the intermediate assignment: the target assignment, with a patch applied.
+ *
+ * <p><b>A granted task needs no patch:</b> The target assignment already places the task on its new owner and
+ * already omits it from the old one, so letting it through unchanged <em>is</em> the grant.
+ * Only a migration this step holds back has to be written down.
+ *
+ * <p>Five kinds of patches.
+ * (Note: multiple patches might apply at once, eg, a regular warmup plant is the first three patches combined):
+ * <ul>
+ * <li>Keep the active task on its old owner.</li>
+ * <li>Withhold the active task from its new owner.</li>
+ * <li>Place a warm-up task.</li>
+ * <li>Keep a standby task on its old owner.</li>
+ * <li>Withhold a standby task from its new owner.</li>
+ * </ul>
+ *
+ * @param targetAssignment
+ * All members' target assignments, as computed by the task assignor.
+ * @param currentAssignment
+ * The indexed current assignment, from {@link #indexCurrentAssignment}, used to find where a borrowed
+ * copy sits.
+ * @param decisions
+ * What the case analysis decided, from {@link #analyzeTasks}.
+ * @param warmupPlan
+ * How each staged migration is being warmed, from {@link #planWarmups}.
+ * @param withheldStandbys
+ * The standby placements to hold back, from {@link #filterStandbys}.
+ *
+ * @return The intermediate assignment, keyed by member ID. The target assignment itself when nothing diverges.
+ */
+ static Map<String, TasksTuple> assemble(
+ final Map<String, TasksTuple> targetAssignment,
+ final CurrentAssignmentIndex currentAssignment,
+ final TaskDecisions decisions,
+ final WarmupPlan warmupPlan,
+ final SortedMap<String, SortedSet<TaskId>> withheldStandbys
+ ) {
+ final Map<String, PatchedTasks> patches = new HashMap<>();
+
+ for (final StagedMigration migration : decisions.stagedMigrations()) {
+ final TaskId task = migration.task();
+ patchFor(patches, targetAssignment, migration.currentOwner()).addActive(task);
+ patchFor(patches, targetAssignment, migration.targetOwner()).removeActive(task);
+
+ if (warmupPlan.borrowedMigrations().contains(task)) {
+ findCopyOnProcess(currentAssignment, task, migration.targetProcessId().orElseThrow())
+ .ifPresent(copy -> patchFor(patches, targetAssignment, copy.memberId()).addStandby(task));
+ }
+ }
+
+ warmupPlan.warmupTasks().forEach((task, memberId) ->
+ patchFor(patches, targetAssignment, memberId).addWarmup(task));
+
+ withheldStandbys.forEach((memberId, tasks) -> {
+ final PatchedTasks patch = patchFor(patches, targetAssignment, memberId);
+ tasks.forEach(patch::removeStandby);
+ });
+
+ if (patches.isEmpty()) {
+ return targetAssignment;
+ }
+
+ final Map<String, TasksTuple> intermediateAssignment = new HashMap<>(targetAssignment);
+ patches.forEach((memberId, patch) -> intermediateAssignment.put(memberId, patch.toTasksTuple()));
+ return Collections.unmodifiableMap(intermediateAssignment);
+ }
+
+ private static PatchedTasks patchFor(
+ final Map<String, PatchedTasks> patches,
+ final Map<String, TasksTuple> targetAssignment,
+ final String memberId
+ ) {
+ return patches.computeIfAbsent(
+ memberId,
+ __ -> new PatchedTasks(targetAssignment.getOrDefault(memberId, TasksTuple.EMPTY))
+ );
+ }
+
/**
* Whether the member has restored the task closely enough to take it over as an active task. Mirrors the client's
* own predicate, so that both ends agree on when a warm-up task is caught up.
@@ -869,6 +1082,88 @@ public class AssignmentRefinerImpl implements AssignmentRefiner {
}
}
+ /**
+ * The per-task lookups {@link #filterStandbys} consults, each built once so that a rule is a map lookup rather
+ * than a fresh scan of the group.
+ *
+ * @param activeStagedOnProcess
+ * The process each staged migration keeps its task running on.
+ * @param standbyKeptOnMember
+ * For each migration warmed by a standby that stays where it is, the one member whose relocated placement
+ * of that task waits in the borrowed copy's stead. Tasks warmed some other way are absent, as are those
+ * with no placement left to hold back.
+ */
+ private record StandbyConflicts(
+ Map<TaskId, String> activeStagedOnProcess,
+ Map<TaskId, String> standbyKeptOnMember
+ ) {
+ }
+
+ /**
+ * One member's tasks in the target assignment, made mutable so that the assembly can patch them.
+ *
+ * <p><b>A subtopology whose partition set becomes empty is dropped, and that is load-bearing rather than
+ * tidiness.</b> The coordinator decides whether a refinement step is due by asking whether a member's assigned
+ * tasks still match what it already holds, and that comparison is a plain map equality: a subtopology key mapped
+ * to an empty set is <em>not</em> equal to the same map without the key. A patch that removed a member's last
+ * task for some subtopology and left the key behind would therefore compare unequal forever, and the group would
+ * mint a fresh refinement step on every heartbeat without anything changing. Pruning is what makes a patch that
+ * ends up holding the target assignment's tasks read as the target assignment.
+ */
+ private static final class PatchedTasks {
+
+ private final Map<String, Set<Integer>> activeTasks;
+ private final Map<String, Set<Integer>> standbyTasks;
+ private final Map<String, Set<Integer>> warmupTasks;
+
+ private PatchedTasks(final TasksTuple tasks) {
+ this.activeTasks = mutableCopy(tasks.activeTasks());
+ this.standbyTasks = mutableCopy(tasks.standbyTasks());
+ this.warmupTasks = mutableCopy(tasks.warmupTasks());
+ }
+
+ private void addActive(final TaskId task) {
+ add(activeTasks, task);
+ }
+
+ private void removeActive(final TaskId task) {
+ remove(activeTasks, task);
+ }
+
+ private void addStandby(final TaskId task) {
+ add(standbyTasks, task);
+ }
+
+ private void removeStandby(final TaskId task) {
+ remove(standbyTasks, task);
+ }
+
+ private void addWarmup(final TaskId task) {
+ add(warmupTasks, task);
+ }
+
+ private TasksTuple toTasksTuple() {
+ return new TasksTuple(activeTasks, standbyTasks, warmupTasks);
+ }
+
+ private static Map<String, Set<Integer>> mutableCopy(final Map<String, Set<Integer>> tasks) {
+ final Map<String, Set<Integer>> copy = new HashMap<>();
+ tasks.forEach((subtopologyId, partitionIds) -> copy.put(subtopologyId, new HashSet<>(partitionIds)));
+ return copy;
+ }
+
+ private static void add(final Map<String, Set<Integer>> tasks, final TaskId task) {
+ tasks.computeIfAbsent(task.subtopologyId(), __ -> new HashSet<>()).add(task.partition());
+ }
+
+ private static void remove(final Map<String, Set<Integer>> tasks, final TaskId task) {
+ final Set<Integer> partitionIds = tasks.get(task.subtopologyId());
+ if (partitionIds != null && partitionIds.remove(task.partition()) && partitionIds.isEmpty()) {
+ tasks.remove(task.subtopologyId());
+ }
+ }
+ }
+
/**
* Which warm-up tasks the intermediate assignment places, and how each staged migration is being warmed.
*
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerTest.java
index e6ea92fa066..1f5ccf207a4 100644
--- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerTest.java
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerTest.java
@@ -23,17 +23,21 @@ import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology;
import org.junit.jupiter.api.Test;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.SortedMap;
+import java.util.SortedSet;
import java.util.TreeMap;
import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasksTuple;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
public class AssignmentRefinerTest {
@@ -669,7 +673,7 @@ public class AssignmentRefinerTest {
@Test
public void shouldDecideNothingForATaskTheTargetAssignmentNoLongerContains() {
- // A topology change removed the task. It belongs in nobody's slice, and its current holder revokes it the
+ // A topology change removed the task. It belongs to nobody now, and its current holder revokes it the
// ordinary way.
final Map<String, StreamsGroupMember> members = Map.of(
"memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)))
@@ -1348,6 +1352,563 @@ public class AssignmentRefinerTest {
assertEquals(Set.of(), plan.parkedMigrations());
}
+ // ---------------------------------------------------------------------------------------------------------------
+ // filterStandbys
+ // ---------------------------------------------------------------------------------------------------------------
+
+ @Test
+ public void shouldWithholdAStandbyOnAProcessThatStillRunsTheTaskAsActive() {
+ // The swap shape: the assignor moves the active to memberB and leaves a standby behind on memberA. While the
+ // migration is staged the task keeps running on memberA, so the standby cannot be placed there as well.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of("memberA", Set.of(STATEFUL_0)), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldEmitTheStandbyOnTheMemberGrantingTheActiveAwayInTheSameStep() {
+ // The efficient half of the swap: memberA hands the active over and keeps a standby in its place, which the
+ // client does by relabelling the task it already has. Nothing is staged, so no rule holds the placement
+ // back, and the relabel happens now rather than a step later when that state is already gone.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+ // memberB's standby is caught up, so the migration is granted rather than staged.
+ final Map<String, MemberTaskOffsets> taskOffsets = Map.of("memberB", offsets(100, 100));
+
+ assertEquals(Map.of(), filter(members, targetAssignment, taskOffsets, 1));
+ }
+
+ @Test
+ public void shouldEmitAStandbyOnASiblingOfTheMemberGrantingTheActiveAway() {
+ // memberA2 would be a second copy on processA until memberA1's hand-over finishes, and it does have to wait
+ // for it -- but in the reconciler, which holds the placement back while the process still runs the task.
+ // Withholding it here as well would only add an epoch.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA1", member("memberA1", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberA2", member("memberA2", "processA", TasksTuple.EMPTY),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA1", TasksTuple.EMPTY,
+ "memberA2", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+ final Map<String, MemberTaskOffsets> taskOffsets = Map.of("memberB", offsets(100, 100));
+
+ assertEquals(Map.of(), filter(members, targetAssignment, taskOffsets, 1));
+ }
+
+ @Test
+ public void shouldEmitAStandbyBlockedOnlyByAPendingRevocation() {
+ // The filter reads the tasks members have been granted, never the ones they were told to give up: indexing
+ // revocations for this rule alone would duplicate what the reconciler already enforces, which refuses to
+ // grant a role for a task the process still physically holds. So this is emitted and the hand-over
+ // serializes itself, at the cost of an extra heartbeat or two before the group settles.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member(
+ "memberA",
+ "processA",
+ TasksTuple.EMPTY,
+ mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ )
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of(), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldWithholdAStandbyOnTheProcessAMigrationIsStagedOn() {
+ // The placement rule 1 protects is the one the staged migration makes: the task runs on memberB for this
+ // step, so the target assignment's standby of it cannot land on memberB's process as well. Nothing holds
+ // the task as an active task here, so the current assignment says nothing about where it runs.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", TasksTuple.EMPTY),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+ final AssignmentRefinerImpl.TaskDecisions decisions = new AssignmentRefinerImpl.TaskDecisions(
+ List.of(new AssignmentRefinerImpl.StagedMigration(
+ STATEFUL_0, "memberB", "memberA", Optional.of("processA"), Optional.empty())),
+ List.of()
+ );
+
+ assertEquals(
+ Map.of("memberB", Set.of(STATEFUL_0)),
+ filter(members, targetAssignment, Map.of(), decisions, 1)
+ );
+ }
+
+ @Test
+ public void shouldEmitTheStandbyOnTheProcessAMigrationIsStagedAwayFrom() {
+ // The migration is staged from memberB, so the intermediate assignment runs the task on processB and not on
+ // processA. That leaves memberA revoking the active it holds, and the standby placement there is how it
+ // recycles that state, so only processB's placement waits.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+ final AssignmentRefinerImpl.TaskDecisions decisions = new AssignmentRefinerImpl.TaskDecisions(
+ List.of(new AssignmentRefinerImpl.StagedMigration(
+ STATEFUL_0, "memberB", "memberC", Optional.of("processC"), Optional.empty())),
+ List.of()
+ );
+
+ assertEquals(
+ Map.of("memberB", Set.of(STATEFUL_0)),
+ filter(members, targetAssignment, Map.of(), decisions, 1)
+ );
+ }
+
+ @Test
+ public void shouldWithholdTheRelocatedStandbyOfABorrowedMigration() {
+ // Borrowing keeps memberB's standby where it is and lets it serve as the warmer too. The relocated placement
+ // the target assignment wants on memberC is what makes that free: granting it as well would leave three
+ // copies where the target assignment asks for two.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of("memberC", Set.of(STATEFUL_0)), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldWithholdOnlyOneRelocatedStandbyOfABorrowedMigration() {
+ // The borrowed copy stands in for one placement, not for every placement the target assignment relocated.
+ // With two standby replicas configured, memberC's placement waits and memberE's is emitted, which leaves
+ // the group on the two replicas it is entitled to.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY),
+ "memberD", member("memberD", "processD", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberE", member("memberE", "processE", TasksTuple.EMPTY)
+ );
+ // The active moves onto memberB, which holds one of the two standbys, so the target assignment has to
+ // relocate both of them: memberB's because it runs the active now, memberD's for balance.
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberD", TasksTuple.EMPTY,
+ "memberE", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of("memberC", Set.of(STATEFUL_0)), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldNotWithholdARelocatedStandbyTheProcessAlreadyHolds() {
+ // The placement that waits for a borrowed copy is one whose process holds nothing of the task. processC
+ // does hold the task, on memberC1, and the target assignment moves that standby onto its sibling memberC2,
+ // which reopens the state the process already has on disk. So memberE's placement is the one that waits,
+ // and processC holds a copy throughout the migration.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC1", member("memberC1", "processC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC2", member("memberC2", "processC", TasksTuple.EMPTY),
+ "memberE", member("memberE", "processE", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC1", TasksTuple.EMPTY,
+ "memberC2", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberE", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of("memberE", Set.of(STATEFUL_0)), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldNotWithholdARelocationWhenRuleOneAlreadyHoldsOneBack() {
+ // memberZ's placement is held back anyway, because the staged migration keeps the task running as an active
+ // task on processZ -- and that active task is the replica the placement would have been. One placement is
+ // all the borrowed copy stands in for, so memberA's cold placement is emitted rather than held back too.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberZ", member("memberZ", "processZ", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC1", member("memberC1", "processC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberA", member("memberA", "processA", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberZ", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC1", TasksTuple.EMPTY,
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of("memberZ", Set.of(STATEFUL_0)), filter(members, targetAssignment, Map.of(), 1));
+ assertEquals(
+ countCopies(targetAssignment, STATEFUL_0),
+ countCopies(assemble(members, targetAssignment, Map.of(), 1), STATEFUL_0)
+ );
+ }
+
+ @Test
+ public void shouldNotWithholdARelocationAssignedToAMemberThatIsGone() {
+ // The target assignment still names memberD, which the group has since removed, so that placement reaches
+ // nobody and the borrowed copy already stands in for it. memberC's placement is emitted -- it sorts ahead
+ // of memberD, so picking by member order alone would have held it back.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberD", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of(), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldEmitTheRelocatedStandbyOfASiblingMove() {
+ // The mirror image of the borrow, and what the sibling move's slot pays for: the copy is moving off memberB1
+ // onto memberB2 as a warm-up, so it stops being the replica the group is entitled to, and the relocated
+ // placement is emitted to backfill it.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB1", member("memberB1", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberB2", member("memberB2", "processB", TasksTuple.EMPTY),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB1", TasksTuple.EMPTY,
+ "memberB2", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of(), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldNotWithholdARelocatedStandbyTheMemberAlreadyHolds() {
+ // Only a placement the member does not have yet adds a replica. memberC already holds this one, so emitting
+ // it changes nothing about the replica count and the borrowing rule has no reason to hold it back.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(Map.of(), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldNotWithholdStandbysOfStatelessTasks() {
+ // A stateless task has no state to restore, so it is never staged and never collides with anything the
+ // refiner decides. Its placements flow through from the target assignment untouched.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATELESS, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATELESS, 0))
+ );
+
+ assertEquals(Map.of(), filter(members, targetAssignment, Map.of(), 1));
+ }
+
+ // ---------------------------------------------------------------------------------------------------------------
+ // assemble
+ // ---------------------------------------------------------------------------------------------------------------
+
+ @Test
+ public void shouldReturnTheTargetAssignmentItselfWhenNothingDiverges() {
+ // A converged group is the overwhelmingly common case, and it costs nothing: with no migration to hold back
+ // there is no patch, so the target assignment is handed straight back.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+
+ assertSame(targetAssignment, assemble(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldKeepTheActiveWithItsCurrentOwnerAndWithholdItFromTheTargetOwner() {
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+
+ // memberB is withheld the active and planted with the warm-up instead; memberA keeps running the task.
+ assertEquals(
+ Map.of(
+ "memberA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.WARMUP, mkTasks(STATEFUL, 0))
+ ),
+ assemble(members, targetAssignment, Map.of(), 1)
+ );
+ }
+
+ @Test
+ public void shouldApplyNoPatchForAGrantedTask() {
+ // The target assignment already places the task on its new owner and omits it from the old one, so letting
+ // it through unchanged is the grant. Nothing is written down for it.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", TasksTuple.EMPTY),
+ "memberB", member("memberB", "processB", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+
+ assertSame(targetAssignment, assemble(members, targetAssignment, Map.of(), 1));
+ }
+
+ @Test
+ public void shouldKeepABorrowedStandbyWhereHistoryLeftIt() {
+ // The target assignment is relocating memberB's standby to memberC, which is exactly why borrowing it is
+ // free -- but that means the target assignment does not give it to memberB. Without patching it back in,
+ // the reconciler would revoke the very copy warming memberB, and the migration would finish cold.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(
+ Map.of(
+ "memberA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ // the borrowed copy, kept: still the standby the group is entitled to, and the warmer as well
+ "memberB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ // the relocated placement, withheld so that the replica count does not move
+ "memberC", TasksTuple.EMPTY
+ ),
+ assemble(members, targetAssignment, Map.of(), 1)
+ );
+ }
+
+ @Test
+ public void shouldMoveASiblingStandbyOntoTheTargetOwnerAndBackfillTheRelocatedOne() {
+ // The mirror of the borrow. memberB1's copy is not kept, because it is moving onto memberB2 as a warm-up;
+ // in exchange the relocated placement on memberC is emitted, which is what the spent slot pays for.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB1", member("memberB1", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberB2", member("memberB2", "processB", TasksTuple.EMPTY),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB1", TasksTuple.EMPTY,
+ "memberB2", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(
+ Map.of(
+ "memberA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberB1", TasksTuple.EMPTY,
+ "memberB2", mkTasksTuple(TaskRole.WARMUP, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ ),
+ assemble(members, targetAssignment, Map.of(), 1)
+ );
+ }
+
+ @Test
+ public void shouldEmitTheSwapAsOneStepOnceTheWarmerIsCaughtUp() {
+ // The shape the whole design is built around: one step hands memberB the active and memberA the standby, so
+ // both sides relabel what they already hold and no restore work is wasted.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.WARMUP, mkTasks(STATEFUL, 0)))
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+ final Map<String, MemberTaskOffsets> taskOffsets = Map.of("memberB", offsets(100, 100));
+
+ assertSame(targetAssignment, assemble(members, targetAssignment, taskOffsets, 1));
+ }
+
+ @Test
+ public void shouldDropASubtopologyKeyWhoseLastTaskWasPatchedAway() {
+ // Pruning is not tidiness: the coordinator decides whether a refinement step is due with a plain map
+ // comparison, so a key left behind mapping to an empty set would read as a change on every heartbeat and
+ // mint refinement steps forever.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))
+ );
+
+ final TasksTuple memberB = assemble(members, targetAssignment, Map.of(), 1).get("memberB");
+
+ // The withheld active was memberB's only task for the subtopology, so the key goes with it.
+ assertEquals(Map.of(), memberB.activeTasks());
+ assertTrue(memberB.sameTasks(withEpochs(mkTasksTuple(TaskRole.WARMUP, mkTasks(STATEFUL, 0)))));
+ }
+
+ @Test
+ public void shouldPreserveTheActiveTaskCountThroughEveryDerivation() {
+ // The wrapper ignores a refined assignment that drops or duplicates an active task, so a derivation that
+ // trips this check would ship a refiner the coordinator silently discards.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0, 1))),
+ "memberB1", member("memberB1", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberB2", member("memberB2", "processB", TasksTuple.EMPTY),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB1", TasksTuple.EMPTY,
+ "memberB2", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 1))
+ );
+
+ for (int numWarmupReplicas = 0; numWarmupReplicas <= 3; numWarmupReplicas++) {
+ assertTrue(
+ AssignmentRefiner.preservesActiveTaskCount(
+ targetAssignment,
+ assemble(members, targetAssignment, Map.of(), numWarmupReplicas)
+ ),
+ "active task count not preserved at numWarmupReplicas=" + numWarmupReplicas
+ );
+ }
+ }
+
+ @Test
+ public void shouldPlaceEveryActiveTaskOnExactlyOneMember() {
+ // The invariant the reconciler cannot recover from if it is broken: a task active on two members at once.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0, 1))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 1))
+ );
+
+ final Map<TaskId, String> activeOwners = new HashMap<>();
+ assemble(members, targetAssignment, Map.of(), 1).forEach((memberId, tasks) ->
+ tasks.activeTasks().forEach((subtopologyId, partitionIds) -> partitionIds.forEach(partitionId -> {
+ final String previous = activeOwners.put(new TaskId(subtopologyId, partitionId), memberId);
+ assertNull(previous, "task active on both " + previous + " and " + memberId);
+ })));
+
+ assertEquals(Set.of(STATEFUL_0, STATEFUL_1), activeOwners.keySet());
+ }
+
+ @Test
+ public void shouldPreserveTheReplicaCountOfABorrowedMigration() {
+ // What makes a borrow free is that it spends a replica the group already has rather than a warm-up slot, so
+ // the copy count has to come out exactly as the target assignment has it -- one active and two standbys
+ // here, with the borrowed copy standing in for the placement that waits.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0))),
+ "memberB", member("memberB", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY),
+ "memberD", member("memberD", "processD", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberE", member("memberE", "processE", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", TasksTuple.EMPTY,
+ "memberB", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0)),
+ "memberD", TasksTuple.EMPTY,
+ "memberE", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))
+ );
+
+ assertEquals(
+ countCopies(targetAssignment, STATEFUL_0),
+ countCopies(assemble(members, targetAssignment, Map.of(), 1), STATEFUL_0)
+ );
+ }
+
+ @Test
+ public void shouldNeverPlaceATaskTwiceOnOneProcess() {
+ // A process holds a given task in at most one role. The refiner relies on this rather than enforcing it, so
+ // the derivation has to avoid producing an intermediate assignment that breaks it.
+ final Map<String, StreamsGroupMember> members = Map.of(
+ "memberA", member("memberA", "processA", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0, 1))),
+ "memberB1", member("memberB1", "processB", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 0))),
+ "memberB2", member("memberB2", "processB", TasksTuple.EMPTY),
+ "memberC", member("memberC", "processC", TasksTuple.EMPTY)
+ );
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", mkTasksTuple(TaskRole.STANDBY, mkTasks(STATEFUL, 1)),
+ "memberB1", TasksTuple.EMPTY,
+ "memberB2", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 0)),
+ "memberC", mkTasksTuple(TaskRole.ACTIVE, mkTasks(STATEFUL, 1))
+ );
+
+ final Set<String> seen = new HashSet<>();
+ assemble(members, targetAssignment, Map.of(), 2).forEach((memberId, tasks) -> {
+ final String processId = members.get(memberId).processId();
+ Map.of(
+ TaskRole.ACTIVE, tasks.activeTasks(),
+ TaskRole.STANDBY, tasks.standbyTasks(),
+ TaskRole.WARMUP, tasks.warmupTasks()
+ ).forEach((role, byRole) -> byRole.forEach((subtopologyId, partitionIds) -> partitionIds.forEach(
+ partitionId -> assertTrue(
+ seen.add(processId + "/" + subtopologyId + "/" + partitionId),
+ "process " + processId + " holds " + subtopologyId + "_" + partitionId + " more than once"
+ ))));
+ });
+ }
+
// ---------------------------------------------------------------------------------------------------------------
// Fixtures
// ---------------------------------------------------------------------------------------------------------------
@@ -1397,6 +1958,83 @@ public class AssignmentRefinerTest {
);
}
+ private static SortedMap<String, SortedSet<TaskId>> filter(
+ final Map<String, StreamsGroupMember> members,
+ final Map<String, TasksTuple> targetAssignment,
+ final Map<String, MemberTaskOffsets> taskOffsets,
+ final int numWarmupReplicas
+ ) {
+ return filter(
+ members,
+ targetAssignment,
+ taskOffsets,
+ analyze(members, targetAssignment, taskOffsets),
+ numWarmupReplicas
+ );
+ }
+
+ /**
+ * Filters against decisions the caller supplies, for a staged migration this group's case analysis would not
+ * produce.
+ */
+ private static SortedMap<String, SortedSet<TaskId>> filter(
+ final Map<String, StreamsGroupMember> members,
+ final Map<String, TasksTuple> targetAssignment,
+ final Map<String, MemberTaskOffsets> taskOffsets,
+ final AssignmentRefinerImpl.TaskDecisions decisions,
+ final int numWarmupReplicas
+ ) {
+ final AssignmentRefinerImpl.CurrentAssignmentIndex currentAssignment = index(members, taskOffsets);
+ return AssignmentRefinerImpl.filterStandbys(
+ targetAssignment,
+ currentAssignment,
+ decisions,
+ AssignmentRefinerImpl.planWarmups(decisions, members, load(members), numWarmupReplicas),
+ members,
+ subtopologies()
+ );
+ }
+
+ /**
+ * The whole derivation, end to end -- which is what the go-live change will wire into {@code refine()}.
+ */
+ private static Map<String, TasksTuple> assemble(
+ final Map<String, StreamsGroupMember> members,
+ final Map<String, TasksTuple> targetAssignment,
+ final Map<String, MemberTaskOffsets> taskOffsets,
+ final int numWarmupReplicas
+ ) {
+ final AssignmentRefinerImpl.CurrentAssignmentIndex currentAssignment = index(members, taskOffsets);
+ final AssignmentRefinerImpl.TaskDecisions decisions =
+ AssignmentRefinerImpl.analyzeTasks(currentAssignment, targetAssignment, members, subtopologies());
+ final AssignmentRefinerImpl.WarmupPlan warmupPlan =
+ AssignmentRefinerImpl.planWarmups(decisions, members, load(members), numWarmupReplicas);
+ return AssignmentRefinerImpl.assemble(
+ targetAssignment,
+ currentAssignment,
+ decisions,
+ warmupPlan,
+ AssignmentRefinerImpl.filterStandbys(
+ targetAssignment, currentAssignment, decisions, warmupPlan, members, subtopologies())
+ );
+ }
+
+ /**
+ * How many copies of the task the assignment holds, counting every role.
+ */
+ private static int countCopies(final Map<String, TasksTuple> assignment, final TaskId task) {
+ int copies = 0;
+ for (final TasksTuple tasks : assignment.values()) {
+ for (final Map<String, Set<Integer>> byRole
+ : List.of(tasks.activeTasks(), tasks.standbyTasks(), tasks.warmupTasks())) {
+ if (byRole.getOrDefault(task.subtopologyId(), Set.of()).contains(task.partition())) {
+ copies++;
+ }
+ }
+ }
+ return copies;
+ }
+
private static SortedMap<String, ConfiguredSubtopology> subtopologies() {
final SortedMap<String, ConfiguredSubtopology> subtopologies = new TreeMap<>();
subtopologies.put(STATEFUL, new ConfiguredSubtopology(