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 1926075628d KAFKA-20665: Refine the assignment refiner scaffolding
(#22985)
1926075628d is described below
commit 1926075628d72b2c0dc82c0e8ec9c79a54f869cd
Author: Matthias J. Sax <[email protected]>
AuthorDate: Wed Jul 29 21:46:19 2026 -0700
KAFKA-20665: Refine the assignment refiner scaffolding (#22985)
The refiner needs a group-wide view similar to the assignor. This PR
changes the refiner setup and what parameters we pass into the it, to
provide all information it needs.
To avoid unnecessary recomputation of an "intermediate assignment" and
to freeze the current reconsolation milestone, we add a cache for it.
Reviewers: Sean Quah <[email protected]>
---
.../coordinator/group/GroupMetadataManager.java | 162 ++++++++++++++-------
.../group/streams/AssignmentRefiner.java | 99 +++++++++++++
.../coordinator/group/streams/StreamsGroup.java | 91 ++++++++++++
.../group/GroupMetadataManagerTest.java | 62 ++++++++
...sGroupStaticMemberGroupMetadataManagerTest.java | 55 +++++++
.../group/streams/AssignmentRefinerTest.java | 154 ++++++++++++++++++++
.../group/streams/StreamsGroupTest.java | 95 ++++++++++++
7 files changed, 668 insertions(+), 50 deletions(-)
diff --git
a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
index a437e4d10d4..367a3e9107b 100644
---
a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
+++
b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java
@@ -154,7 +154,7 @@ import
org.apache.kafka.coordinator.group.modern.share.ShareGroup.InitMapValue;
import
org.apache.kafka.coordinator.group.modern.share.ShareGroup.ShareGroupStatePartitionMetadataInfo;
import
org.apache.kafka.coordinator.group.modern.share.ShareGroupAssignmentBuilder;
import org.apache.kafka.coordinator.group.modern.share.ShareGroupMember;
-import org.apache.kafka.coordinator.group.streams.MemberTaskOffsets;
+import org.apache.kafka.coordinator.group.streams.AssignmentRefiner;
import
org.apache.kafka.coordinator.group.streams.StreamsCoordinatorRecordHelpers;
import org.apache.kafka.coordinator.group.streams.StreamsGroup;
import org.apache.kafka.coordinator.group.streams.StreamsGroupDescribeResult;
@@ -2114,7 +2114,14 @@ public class GroupMetadataManager {
} else {
StreamsGroupMember maybeOldStaticMember =
group.staticMember(instanceId);
if (maybeOldStaticMember != null &&
!maybeOldStaticMember.memberId().equals(memberId)) {
- replaceStaticOldMember = maybeOldStaticMember;
+ replaceStaticOldMember = maybeOldStaticMember;
+ // Replacing a static member relabels its target assignment
from the old to the new member ID without
+ // bumping the assignment epoch, so an intermediate assignment
derived for this epoch no longer matches
+ // the members it was derived for. Re-key it rather than
dropping it: the replacement copies the old
+ // member's state, so the decisions of this epoch still hold,
and deriving a new one here would re-plan
+ // mid-epoch and could revise the slice of a member that
already reconciled and is therefore not
+ // reconciled again within this epoch.
+
group.relabelRefinedAssignment(maybeOldStaticMember.memberId(), memberId);
}
member = getOrMaybeCreateStaticStreamsGroupMember(
group,
@@ -2229,18 +2236,12 @@ public class GroupMetadataManager {
assignmentUpdate = AssignmentUpdate.RECOMPUTE;
}
- TasksTuple refinedAssignment = null;
+ Map<String, TasksTuple> refinedGroupAssignment = null;
if (assignmentUpdate == AssignmentUpdate.NONE && group.state() ==
StreamsGroup.StreamsGroupState.STABLE) {
// We are not computing a new target assignment later, thus we try
to refine the current target assignment
// into an intermediate assignment (with warm-up tasks) the member
should be reconciled towards.
- refinedAssignment = maybeRefineAssignment(// no-op for now
- updatedMember,
- group.targetAssignment(),
- group.taskOffsets(),
- streamsGroupNumWarmupReplicas(groupId),
- streamsGroupAcceptableRecoveryLag(groupId)
- );
- if (!refinedAssignment.sameTasks(updatedMember.assignedTasks())) {
+ refinedGroupAssignment = refineAssignment(group,
group.targetAssignment(), updatedConfiguredTopology);
+ if (!refinedGroupAssignment.getOrDefault(memberId,
TasksTuple.EMPTY).sameTasks(updatedMember.assignedTasks())) {
assignmentUpdate = AssignmentUpdate.REFINED;
}
}
@@ -2294,16 +2295,21 @@ public class GroupMetadataManager {
// 4b. If we did not already refine above -- ie, we computed a new
target assignment, or the group is not
// yet reconciled (a rebalance is still in progress) -- refine the
target assignment into an intermediate
- // assignment (with warm-up tasks) the member should be reconciled
towards.
- if (refinedAssignment == null) {
- refinedAssignment = maybeRefineAssignment(// no-op for now
- updatedMember,
+ // assignment (with warm-up tasks) the member should be reconciled
towards. Otherwise, the assignment we
+ // refined above is the intermediate assignment of this epoch, whether
we bumped the epoch for it or not: it
+ // was derived from a settled group, so it also holds for the target
assignment of the bumped epoch, which a
+ // refinement step leaves unchanged.
+ if (refinedGroupAssignment == null) {
+ refinedGroupAssignment = refinedAssignmentForEpoch(
+ group,
+ updateTargetAssignmentResult.targetAssignmentEpoch(),
updateTargetAssignmentResult.targetAssignment,
- group.taskOffsets(),
- streamsGroupNumWarmupReplicas(group.groupId()),
- streamsGroupAcceptableRecoveryLag(group.groupId())
+ updatedConfiguredTopology
);
+ } else {
+ maybeCacheRefinedAssignment(group,
updateTargetAssignmentResult.targetAssignmentEpoch(), refinedGroupAssignment);
}
+ TasksTuple refinedAssignment =
refinedGroupAssignment.getOrDefault(memberId, TasksTuple.EMPTY);
// 5. Reconcile the member's assignment with the (refined) target
assignment if the member is not
// fully reconciled yet.
@@ -4507,40 +4513,96 @@ public class GroupMetadataManager {
/**
- * Refines the task assignor's target assignment into the
<em>intermediate</em> assignment that
- * the reconciler ({@link
org.apache.kafka.coordinator.group.streams.CurrentAssignmentBuilder}) converges
members toward.
- * <p>
- * The intermediate assignment is the current assignment with warm-up
tasks inserted (for later promotion to active,
- * based on per-member changelog lag), so that a task is moved to a new
owner only once that owner has caught up.
- * It is held in memory only and is never persisted: the assignor's target
assignment remains the persisted source of
- * truth (and the intermediate is reconstructed from the persisted
current-assignment records after a coordinator failover).
- * <p>
- * The refiner is invoked on every heartbeat, before reconciliation, so it
can react to all the inputs that can change
- * the intermediate assignment between reassignments — newly reported task
offsets (a warm-up may have become hot),
- * {@code num.warmup.replicas} / {@code acceptable.recovery.lag} config
changes, and members acknowledging task
- * revocation/restoration (advancing an in-flight migration).
- *
- * @param member
- * The member to produce the refined (intermediate) assignment for.
- * @param targetAssignment
- * All members' target assignments (group context).
- * @param taskOffsets
- * The latest per-member changelog offsets/end-offsets reported via
heartbeats (group context).
- * @param numWarmupReplicas
- * The configured maximum number of warm-up replicas.
- * @param acceptableRecoveryLag
- * The lag at or below which a warm-up is considered caught up and
can be promoted.
- *
- * @return The member's intermediate assignment tuple.
- */
- private static TasksTuple maybeRefineAssignment(
- final StreamsGroupMember member,
+ * Derives the <em>intermediate</em> assignment that the reconciler
+ * ({@link
org.apache.kafka.coordinator.group.streams.CurrentAssignmentBuilder}) converges
the members toward, for a
+ * refinement step that is not minted yet. It is compared against the
members' current assignments to decide whether
+ * a refinement step is due at all, and becomes the intermediate
assignment of the epoch if it is (see
+ * {@link #refinedAssignmentForEpoch}).
+ *
+ * @param group The streams group.
+ * @param targetAssignment All members' target assignments.
+ * @param configuredTopology The configured topology.
+ *
+ * @return The intermediate assignment keyed by member ID.
+ */
+ private Map<String, TasksTuple> refineAssignment(
+ final StreamsGroup group,
+ final Map<String, TasksTuple> targetAssignment,
+ final ConfiguredTopology configuredTopology
+ ) {
+ final int numWarmupReplicas =
streamsGroupNumWarmupReplicas(group.groupId());
+ if (numWarmupReplicas == 0) {
+ // Warm-up tasks are disabled, so there is nothing to refine and
no state to keep for the group.
+ return targetAssignment;
+ }
+ final Map<String, TasksTuple> refinedAssignment =
AssignmentRefiner.refine(
+ group.members(),
+ targetAssignment,
+ group.taskOffsets(),
+ configuredTopology,
+ numWarmupReplicas,
+ streamsGroupAcceptableRecoveryLag(group.groupId())
+ );
+ if (!AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment)) {
+ // Reconciling towards an intermediate assignment that lost or
duplicated an active task would leave input
+ // partitions unprocessed or processed twice, so fall back to the
target assignment. That is the same
+ // assignment the group reconciles towards with warm-up tasks
disabled, so no further handling is needed --
+ // beyond that tasks now move without warming up, which the
operator should know about.
+ log.error("[GroupId {}] The refined assignment does not hand out
as many active tasks as the target " +
+ "assignment. Reconciling towards the target assignment
instead, so tasks that have to move do so " +
+ "without warming up first. Target assignment: {}, refined
assignment: {}.",
+ group.groupId(), targetAssignment, refinedAssignment);
+ return targetAssignment;
+ }
+ return refinedAssignment;
+ }
+
+ /**
+ * Returns the intermediate assignment of the given assignment epoch,
deriving it once and caching it for the
+ * lifetime of the epoch. Every refinement step is an epoch of its own, so
the decisions of a step are taken when
+ * its epoch is minted and do not change while the members reconcile
toward it.
+ *
+ * @param group The streams group.
+ * @param assignmentEpoch The assignment epoch the members are being
reconciled to.
+ * @param targetAssignment All members' target assignments.
+ * @param configuredTopology The configured topology.
+ *
+ * @return The intermediate assignment keyed by member ID.
+ */
+ private Map<String, TasksTuple> refinedAssignmentForEpoch(
+ final StreamsGroup group,
+ final int assignmentEpoch,
final Map<String, TasksTuple> targetAssignment,
- final Map<String, MemberTaskOffsets> taskOffsets,
- final int numWarmupReplicas,
- final long acceptableRecoveryLag
+ final ConfiguredTopology configuredTopology
) {
- return targetAssignment.getOrDefault(member.memberId(),
TasksTuple.EMPTY);
+ final Map<String, TasksTuple> cached =
group.refinedAssignment(assignmentEpoch);
+ if (cached != null) {
+ return cached;
+ }
+ final Map<String, TasksTuple> refinedAssignment =
refineAssignment(group, targetAssignment, configuredTopology);
+ maybeCacheRefinedAssignment(group, assignmentEpoch, refinedAssignment);
+ return refinedAssignment;
+ }
+
+ /**
+ * Freezes the given intermediate assignment as the one of the given
assignment epoch, so that the decisions of a
+ * refinement step do not change while the members reconcile toward it.
+ * <p>
+ * Nothing is kept for a group that has warm-up tasks disabled: the
intermediate assignment is then just the target
+ * assignment, which the caller has anyway, so there is nothing to freeze.
+ *
+ * @param group The streams group.
+ * @param assignmentEpoch The assignment epoch the intermediate
assignment was derived for.
+ * @param refinedAssignment The intermediate assignment keyed by member
ID.
+ */
+ private void maybeCacheRefinedAssignment(
+ final StreamsGroup group,
+ final int assignmentEpoch,
+ final Map<String, TasksTuple> refinedAssignment
+ ) {
+ if (streamsGroupNumWarmupReplicas(group.groupId()) > 0) {
+ group.setRefinedAssignment(assignmentEpoch, refinedAssignment);
+ }
}
/**
diff --git
a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefiner.java
b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefiner.java
new file mode 100644
index 00000000000..1322f891016
--- /dev/null
+++
b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/AssignmentRefiner.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.streams;
+
+import org.apache.kafka.coordinator.group.streams.topics.ConfiguredTopology;
+
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Refines the task assignor's target assignment into the
<em>intermediate</em> assignment that the reconciler
+ * ({@link CurrentAssignmentBuilder}) converges the members toward.
+ * <p>
+ * The assignor decides <em>where</em> a task belongs; the refiner decides
<em>how</em> the group gets there. When a
+ * stateful task has to move to a member that does not hold its state yet,
handing it over right away would stall
+ * processing while the new owner restores the state from the changelog.
Instead, the refiner leaves the task with its
+ * current owner and hands the new owner a warm-up task, so it can restore the
state in the background. Once the warm-up
+ * task has caught up -- its lag is within {@code acceptable.recovery.lag} --
a later refinement step moves the task
+ * over. {@code num.warmup.replicas} bounds how many such warm-up tasks the
group runs at a time.
+ * <p>
+ * Properties of the intermediate assignment that callers rely on:
+ * <ul>
+ * <li>It is <b>derived for the whole group at once</b>, because both the
warm-up budget and the rule that a
+ * process must not hold a task twice are group-wide. Callers take the
individual member's slice out of the
+ * result.</li>
+ * <li>It is <b>held in memory only</b> and never persisted. The
assignor's target assignment remains the persisted
+ * source of truth, and the intermediate assignment is derived again from
the persisted current assignments after a
+ * coordinator failover.</li>
+ * <li>It is <b>frozen for the duration of an assignment epoch</b>: every
refinement step is an epoch of its own, so
+ * the decisions of a step are taken once, when the epoch is minted, and
do not change while the members reconcile
+ * towards it. See {@link StreamsGroup#refinedAssignment(int)}.</li>
+ * </ul>
+ */
+public class AssignmentRefiner {
+
+ private AssignmentRefiner() {
+ }
+
+ /**
+ * Derives the intermediate assignment for all members of the group.
+ *
+ * @param members All members of the group, with their
current assignments, the tasks they are still
+ * to revoke, and their process IDs.
+ * @param targetAssignment All members' target assignments, as
computed by the task assignor.
+ * @param taskOffsets The latest changelog offsets/end-offsets
reported by the members via heartbeats,
+ * from which the lag of a warm-up task is
derived. Not populated for a member that
+ * has not reported any offsets yet, for
example right after a coordinator failover.
+ * @param configuredTopology The configured topology, which tells
whether a subtopology is stateful. Only
+ * stateful tasks are warmed up; a stateless
task has no state to restore.
+ * @param numWarmupReplicas The maximum number of warm-up tasks the
group may run at a time. Zero disables
+ * warm-up tasks altogether, in which case
the intermediate assignment is the target
+ * assignment.
+ * @param acceptableRecoveryLag The lag at or below which a warm-up task
is considered caught up.
+ *
+ * @return The intermediate assignment, keyed by member ID.
+ */
+ public static Map<String, TasksTuple> refine(
+ Map<String, StreamsGroupMember> members,
+ Map<String, TasksTuple> targetAssignment,
+ Map<String, MemberTaskOffsets> taskOffsets,
+ ConfiguredTopology configuredTopology,
+ int numWarmupReplicas,
+ long acceptableRecoveryLag
+ ) {
+ // No warm-up tasks are inserted yet, so the intermediate assignment
is the target assignment.
+ return targetAssignment;
+ }
+
+ public static boolean preservesActiveTaskCount(
+ Map<String, TasksTuple> targetAssignment,
+ Map<String, TasksTuple> refinedAssignment
+ ) {
+ return countActiveTasks(targetAssignment) ==
countActiveTasks(refinedAssignment);
+ }
+
+ private static int countActiveTasks(Map<String, TasksTuple> assignment) {
+ int count = 0;
+ for (TasksTuple tasks : assignment.values()) {
+ for (Set<Integer> partitionIds : tasks.activeTasks().values()) {
+ count += partitionIds.size();
+ }
+ }
+ return count;
+ }
+}
diff --git
a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroup.java
b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroup.java
index 3a481459f1c..213e49a2eb1 100644
---
a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroup.java
+++
b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/StreamsGroup.java
@@ -225,6 +225,35 @@ public class StreamsGroup implements Group {
*/
private final Map<String, MemberTaskOffsets> taskOffsets = new HashMap<>();
+ /**
+ * The intermediate assignment the members are reconciled toward, as
derived by the {@link AssignmentRefiner} from
+ * the target assignment, together with the assignment epoch it was
derived for. Like {@link #taskOffsets}, this is
+ * held in memory only and never persisted; it is derived again from the
persisted assignments after a coordinator
+ * failover.
+ * <p>
+ * Every refinement step is an assignment epoch of its own, so the
intermediate assignment is derived once per
+ * epoch, when that epoch is minted, and stays fixed while the members
reconcile toward it. It is therefore cached
+ * together with the epoch it belongs to, and a cached value from another
epoch is never used.
+ * <p>
+ * Unlike {@link #taskOffsets} it is kept in a timeline container, because
it is derived from state that a failed
+ * write rolls back. Rolling back restores an earlier assignment epoch,
which the next epoch to be minted then
+ * reuses -- with a target assignment that may differ from the one this
was derived from. Left outside the rollback,
+ * the cache would be hit for that epoch and reconcile the members toward
an intermediate assignment of an
+ * assignment that no longer exists. The epoch and the assignment share
one container so that a rollback can never
+ * restore one without the other.
+ */
+ private final TimelineObject<RefinedAssignment> refinedAssignment;
+
+ /**
+ * An intermediate assignment together with the assignment epoch it was
derived for.
+ *
+ * @param assignmentEpoch The assignment epoch, or {@code -1} if nothing
was derived yet.
+ * @param assignment The intermediate assignment keyed by member ID.
+ */
+ private record RefinedAssignment(int assignmentEpoch, Map<String,
TasksTuple> assignment) {
+ private static final RefinedAssignment NONE = new
RefinedAssignment(-1, Map.of());
+ }
+
/**
* The Streams topology.
*/
@@ -299,6 +328,7 @@ public class StreamsGroup implements Group {
this.currentActiveTaskToProcessId = new
TimelineHashMap<>(snapshotRegistry, 0);
this.currentStandbyTaskToProcessIds = new
TimelineHashMap<>(snapshotRegistry, 0);
this.currentWarmupTaskToProcessIds = new
TimelineHashMap<>(snapshotRegistry, 0);
+ this.refinedAssignment = new TimelineObject<>(snapshotRegistry,
RefinedAssignment.NONE);
this.topology = new TimelineObject<>(snapshotRegistry,
Optional.empty());
this.configuredTopology = new TimelineObject<>(snapshotRegistry,
Optional.empty());
this.lastAssignmentConfigs = new TimelineHashMap<>(snapshotRegistry,
0);
@@ -586,6 +616,67 @@ public class StreamsGroup implements Group {
return Collections.unmodifiableMap(taskOffsets);
}
+ /**
+ * Returns the intermediate assignment that was derived for the given
assignment epoch, if any. A value derived for
+ * an earlier epoch is not returned: the intermediate assignment of an
epoch is fixed, but a new epoch means a new
+ * refinement step, which has to be derived anew.
+ *
+ * @param assignmentEpoch The assignment epoch to return the intermediate
assignment for.
+ *
+ * @return The intermediate assignment keyed by member ID, or {@code null}
if it was not derived for that epoch.
+ */
+ public Map<String, TasksTuple> refinedAssignment(int assignmentEpoch) {
+ final RefinedAssignment cached = refinedAssignment.get();
+ return cached.assignmentEpoch() == assignmentEpoch ?
cached.assignment() : null;
+ }
+
+ /**
+ * Caches the intermediate assignment derived for the given assignment
epoch. This is transient state; it must be
+ * set from the heartbeat path only, never while replaying records.
+ * <p>
+ * The given map is copied, so that the cached assignment is genuinely
fixed for its epoch rather than tracking
+ * whatever the caller hands over: a refiner that returns the target
assignment as-is would otherwise leave the
+ * cache holding a view of {@link #targetAssignment()}, which changes as
records are replayed. Only the map itself
+ * is copied -- a {@link TasksTuple} already owns its task sets, so there
is nothing else that could change
+ * underneath. The copy is free once the refiner returns a map it built
itself, because {@link Map#copyOf} hands
+ * back an already-immutable map unchanged.
+ *
+ * @param assignmentEpoch The assignment epoch the intermediate
assignment was derived for.
+ * @param refinedAssignment The intermediate assignment keyed by member
ID.
+ */
+ public void setRefinedAssignment(int assignmentEpoch, Map<String,
TasksTuple> refinedAssignment) {
+ this.refinedAssignment.set(new RefinedAssignment(
+ assignmentEpoch,
+ Map.copyOf(Objects.requireNonNull(refinedAssignment))
+ ));
+ }
+
+ /**
+ * Re-keys the cached intermediate assignment from an old to a new member
ID.
+ * <p>
+ * Replacing a static member relabels its target assignment the same way,
without advancing the assignment epoch.
+ * The cached intermediate assignment is therefore still the one the other
members are reconciling towards; only its
+ * key for the replaced member went stale. Re-keying it keeps the
decisions of the epoch fixed, whereas dropping it
+ * would derive a new one mid-epoch, which could revise the slice of a
member that already reconciled and is
+ * therefore not reconciled again within this epoch.
+ *
+ * @param oldMemberId The member ID the intermediate assignment was
derived for.
+ * @param newMemberId The member ID replacing it.
+ */
+ public void relabelRefinedAssignment(String oldMemberId, String
newMemberId) {
+ final RefinedAssignment cached = refinedAssignment.get();
+ final TasksTuple refinedTasks = cached.assignment().get(oldMemberId);
+ if (refinedTasks == null) {
+ // Either nothing was derived for the replaced member, or the
cached assignment belongs to another epoch --
+ // in which case refinedAssignment(int) does not hand it out
anyway.
+ return;
+ }
+ final Map<String, TasksTuple> relabeled = new
HashMap<>(cached.assignment());
+ relabeled.remove(oldMemberId);
+ relabeled.put(newMemberId, refinedTasks);
+ this.refinedAssignment.set(new
RefinedAssignment(cached.assignmentEpoch(), Map.copyOf(relabeled)));
+ }
+
/**
* Remove the static member mapping if the removed member is static.
*
diff --git
a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
index 7636eb7a202..c1bbff94398 100644
---
a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
+++
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java
@@ -19178,6 +19178,68 @@ public class GroupMetadataManagerTest {
.noneMatch(r -> r.key() instanceof
StreamsGroupTargetAssignmentMemberKey));
}
+ @Test
+ public void
testStreamsGroupKeepsNoRefinedAssignmentWhenWarmupsAreDisabled() {
+ String groupId = "fooup";
+ String memberId = Uuid.randomUuid().toString();
+ String subtopology1 = "subtopology1";
+ String fooTopicName = "foo";
+ Uuid fooTopicId = Uuid.randomUuid();
+ Topology topology = new Topology().setSubtopologies(List.of(
+ new
Subtopology().setSubtopologyId(subtopology1).setSourceTopics(List.of(fooTopicName))
+ ));
+
+ CoordinatorMetadataImage metadataImage = new MetadataImageBuilder()
+ .addTopic(fooTopicId, fooTopicName, 4)
+ .buildCoordinatorMetadataImage();
+ long groupMetadataHash = computeGroupHash(Map.of(
+ fooTopicName, computeTopicHash(fooTopicName, metadataImage)
+ ));
+
+ MockTaskAssignor assignor = new MockTaskAssignor("sticky");
+ GroupMetadataManagerTestContext context = new
GroupMetadataManagerTestContext.Builder()
+ .withStreamsGroupTaskAssignors(List.of(assignor))
+ .withMetadataImage(metadataImage)
+ .withStreamsGroup(new StreamsGroupBuilder(groupId, 10)
+ .withMember(streamsGroupMemberBuilderWithDefaults(memberId)
+ .setMemberEpoch(10)
+ .setPreviousMemberEpoch(10)
+
.setAssignedTasks(mkTasksTupleWithCommonEpoch(TaskRole.ACTIVE, 10,
+ TaskAssignmentTestUtil.mkTasks(subtopology1, 0, 1, 2)))
+ .build())
+ .withTopology(StreamsTopology.fromHeartbeatRequest(topology))
+ .withTargetAssignment(memberId,
TaskAssignmentTestUtil.mkTasksTuple(TaskRole.ACTIVE,
+ TaskAssignmentTestUtil.mkTasks(subtopology1, 0, 1, 2)))
+ .withTargetAssignmentEpoch(10)
+ .withMetadataHash(groupMetadataHash)
+ .withValidatedTopologyEpoch(0)
+ .withLastAssignmentConfigs(Map.of("num.standby.replicas", "0"))
+ )
+ .build();
+
+ Properties groupConfig = new Properties();
+
groupConfig.setProperty(GroupConfig.STREAMS_NUM_WARMUP_REPLICAS_CONFIG, "0");
+ context.updateGroupConfig(groupId, groupConfig);
+
+ context.streamsGroupHeartbeat(
+ new StreamsGroupHeartbeatRequestData()
+ .setGroupId(groupId)
+ .setMemberId(memberId)
+ .setMemberEpoch(10)
+ .setProcessId("process-id")
+ .setRebalanceTimeoutMs(1500)
+ .setActiveTasks(List.of(new
StreamsGroupHeartbeatRequestData.TaskIds()
+ .setSubtopologyId(subtopology1)
+ .setPartitions(List.of(0, 1, 2))))
+ .setStandbyTasks(List.of())
+ .setWarmupTasks(List.of()));
+
+ // With warm-up tasks disabled the intermediate assignment is just the
target assignment, which the caller has
+ // anyway, so nothing is frozen for the group -- the group keeps no
per-epoch state it would never read.
+ StreamsGroup group =
context.groupMetadataManager.streamsGroup(groupId);
+ assertNull(group.refinedAssignment(group.assignmentEpoch()));
+ }
+
@Test
public void testStreamsGroupHeartbeatStoresTaskOffsetsWithoutPersisting() {
String groupId = "fooup";
diff --git
a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/StreamsGroupStaticMemberGroupMetadataManagerTest.java
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/StreamsGroupStaticMemberGroupMetadataManagerTest.java
index 48fe6bb834a..68f1c763d23 100644
---
a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/StreamsGroupStaticMemberGroupMetadataManagerTest.java
+++
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/StreamsGroupStaticMemberGroupMetadataManagerTest.java
@@ -76,6 +76,7 @@ import static
org.apache.kafka.coordinator.group.StreamsGroupTestUtil.streamsGro
import static
org.apache.kafka.coordinator.group.StreamsGroupTestUtil.streamsTopicFixture;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -1524,6 +1525,60 @@ class StreamsGroupStaticMemberGroupMetadataManagerTest {
));
}
+ @Test
+ public void testStaticMemberReplacementRelabelsCachedRefinedAssignment() {
+ int groupEpoch = DEFAULT_GROUP_EPOCH;
+
+ String groupId = "fooup";
+ String oldMemberId = Uuid.randomUuid().toString();
+ String rejoinMemberId = Uuid.randomUuid().toString();
+ String instanceId = Uuid.randomUuid().toString();
+
+ StreamsTopicFixture topic = streamsTopicFixture("subtopology1", "foo",
4);
+ TasksTuple oldTargetAssignment = topic.targetAssignment(0, 1, 2, 3);
+ TasksTupleWithEpochs assignedTasks = topic.assignedTasks(groupEpoch,
0, 1, 2, 3);
+
+ StreamsGroupMember oldMember =
streamsGroupMemberBuilderWithDefaults(oldMemberId, instanceId)
+ .setMemberEpoch(LEAVE_GROUP_STATIC_MEMBER_EPOCH)
+ .setPreviousMemberEpoch(groupEpoch)
+ .setAssignedTasks(resetAssignedTasksEpochsToZero(assignedTasks))
+ .build();
+
+ GroupMetadataManagerTestContext context = new
GroupMetadataManagerTestContext.Builder()
+ .withStreamsGroupTaskAssignors(List.of(new
MockTaskAssignor("sticky")))
+ .withMetadataImage(topic.metadataImage())
+ .withStreamsGroup(new StreamsGroupBuilder(groupId, groupEpoch)
+ .withMember(oldMember)
+ .withTargetAssignment(oldMemberId, oldTargetAssignment)
+ .withTargetAssignmentEpoch(groupEpoch)
+
.withTopology(StreamsTopology.fromHeartbeatRequest(topic.topology()))
+ .withValidatedTopologyEpoch(0)
+ .withMetadataHash(topic.metadataHash())
+ .withLastAssignmentConfigs(getDefaultAssignmentConfigs()))
+
.withConfig(GroupCoordinatorConfig.STREAMS_GROUP_INITIAL_REBALANCE_DELAY_MS_CONFIG,
GroupCoordinatorConfig.STREAMS_GROUP_INITIAL_REBALANCE_DELAY_MS_DEFAULT)
+ .build();
+
+ // Prime the frozen intermediate assignment for the current assignment
epoch, keyed by the OLD member ID, as a
+ // previous heartbeat of the departed member would have left it.
+ StreamsGroup streamsGroup =
context.groupMetadataManager.streamsGroup(groupId);
+ streamsGroup.setRefinedAssignment(groupEpoch, Map.of(oldMemberId,
oldTargetAssignment));
+
+ CoordinatorResult<StreamsGroupHeartbeatResult, CoordinatorRecord>
result = context.streamsGroupHeartbeat(
+ staticJoinHeartbeat(groupId, rejoinMemberId, instanceId,
DEFAULT_PROCESS_ID)
+ );
+
+ // A static replacement relabels the target assignment oldMemberId ->
rejoinMemberId WITHOUT advancing the
+ // assignment epoch, so the epoch alone cannot tell the cached entry
apart from a still-valid one. Unless the
+ // cached entry is re-keyed too, the replacing member is sliced out of
a map that only knows the old ID and gets
+ // no tasks at all.
+ assertEquals(groupEpoch, result.response().data().memberEpoch());
+ assertEquals(
+ oldTargetAssignment.activeTasks(),
+
streamsGroup.refinedAssignment(groupEpoch).get(rejoinMemberId).activeTasks()
+ );
+
assertFalse(streamsGroup.refinedAssignment(groupEpoch).containsKey(oldMemberId));
+ }
+
@Test
public void
testStaticMemberRejoinsWithSameMemberIdAndDifferentInstanceId() {
int groupEpoch = DEFAULT_GROUP_EPOCH;
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
new file mode 100644
index 00000000000..eeb105d1fee
--- /dev/null
+++
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/AssignmentRefinerTest.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.streams;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class AssignmentRefinerTest {
+
+ private static TasksTuple active(final Map<String, Set<Integer>>
activeTasks) {
+ return new TasksTuple(activeTasks, Map.of(), Map.of());
+ }
+
+ @Test
+ public void shouldPreserveActiveTaskCountOfUnchangedAssignment() {
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1))),
+ "memberB", active(Map.of("0", Set.of(2)))
+ );
+
+
assertTrue(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
targetAssignment));
+ }
+
+ @Test
+ public void
shouldPreserveActiveTaskCountWhenATaskIsHeldBackWithItsCurrentOwner() {
+ // What a refinement step does to stage a migration: the target
assignment moves 0_2 to memberB, the refined
+ // assignment leaves it with memberA while memberB warms it up. The
active tasks themselves are unchanged.
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1))),
+ "memberB", active(Map.of("0", Set.of(2)))
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1, 2))),
+ "memberB", new TasksTuple(Map.of(), Map.of(), Map.of("0",
Set.of(2)))
+ );
+
+
assertTrue(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void shouldPreserveActiveTaskCountWhenStandbysAreDeferred() {
+ // A refinement step may defer a standby to a later step, which is not
a defect this check is about.
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", new TasksTuple(Map.of("0", Set.of(0)), Map.of("0",
Set.of(1)), Map.of()),
+ "memberB", new TasksTuple(Map.of("0", Set.of(1)), Map.of("0",
Set.of(0)), Map.of())
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0))),
+ "memberB", active(Map.of("0", Set.of(1)))
+ );
+
+
assertTrue(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void shouldNotPreserveActiveTaskCountWhenATaskWasDropped() {
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1))),
+ "memberB", active(Map.of("0", Set.of(2)))
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1))),
+ "memberB", active(Map.of())
+ );
+
+
assertFalse(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void
shouldNotPreserveActiveTaskCountWhenASubtopologyWasDroppedEntirely() {
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0), "1", Set.of(0)))
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0)))
+ );
+
+
assertFalse(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void
shouldNotPreserveActiveTaskCountWhenATaskWasHandedToTwoMembers() {
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0))),
+ "memberB", active(Map.of("0", Set.of(1)))
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1))),
+ "memberB", active(Map.of("0", Set.of(1)))
+ );
+
+
assertFalse(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void shouldNotPreserveActiveTaskCountWhenATaskWasInvented() {
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0)))
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0, 1)))
+ );
+
+
assertFalse(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void shouldNotDetectADropAndADuplicateCancellingEachOtherOut() {
+ // The accepted blind spot of counting: 0_0 was dropped and 0_1 handed
to both members, so the count still
+ // matches. It takes two coordinated mistakes in one derivation, and
the exhaustive invariant is covered by the
+ // refiner's own tests. If this check is ever strengthened, this test
is what should fail.
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0))),
+ "memberB", active(Map.of("0", Set.of(1)))
+ );
+ final Map<String, TasksTuple> refinedAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(1))),
+ "memberB", active(Map.of("0", Set.of(1)))
+ );
+
+
assertTrue(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
refinedAssignment));
+ }
+
+ @Test
+ public void
shouldPreserveActiveTaskCountWhenTheTargetAssignmentItselfDuplicatesATask() {
+ // A target assignment that places an active task twice is the
assignor's defect, not the refinement's, so a
+ // refinement that keeps it as-is is not blamed for it.
+ final Map<String, TasksTuple> targetAssignment = Map.of(
+ "memberA", active(Map.of("0", Set.of(0))),
+ "memberB", active(Map.of("0", Set.of(0)))
+ );
+
+
assertTrue(AssignmentRefiner.preservesActiveTaskCount(targetAssignment,
targetAssignment));
+ }
+}
diff --git
a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/StreamsGroupTest.java
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/StreamsGroupTest.java
index a4bc5bdad8f..44bc85689c1 100644
---
a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/StreamsGroupTest.java
+++
b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/StreamsGroupTest.java
@@ -136,6 +136,101 @@ public class StreamsGroupTest {
assertEquals(newerOffsets, streamsGroup.taskOffsets("member-id"));
}
+ @Test
+ public void testCacheAndRetrieveRefinedAssignment() {
+ StreamsGroup streamsGroup = createStreamsGroup("foo");
+
+ assertNull(streamsGroup.refinedAssignment(0));
+
+ Map<String, TasksTuple> refinedAssignment = Map.of(
+ "member-id", new TasksTuple(Map.of("sub-1", Set.of(0)), Map.of(),
Map.of())
+ );
+ streamsGroup.setRefinedAssignment(5, refinedAssignment);
+
+ assertEquals(refinedAssignment, streamsGroup.refinedAssignment(5));
+ // An intermediate assignment derived for one assignment epoch must
not be used for another one: a new epoch is
+ // a new refinement step.
+ assertNull(streamsGroup.refinedAssignment(4));
+ assertNull(streamsGroup.refinedAssignment(6));
+ }
+
+ @Test
+ public void testCachedRefinedAssignmentDoesNotTrackTheCallersMap() {
+ StreamsGroup streamsGroup = createStreamsGroup("foo");
+
+ TasksTuple memberTasks = new TasksTuple(Map.of("sub-1", Set.of(0)),
Map.of(), Map.of());
+ Map<String, TasksTuple> derived = new HashMap<>();
+ derived.put("member-id", memberTasks);
+
+ streamsGroup.setRefinedAssignment(5, derived);
+
+ // The intermediate assignment of an epoch is fixed once derived. A
refiner that returns the target assignment
+ // as-is hands over a view of it, which changes as records are
replayed, so the cache must copy rather than alias.
+ derived.put("late-member-id", new TasksTuple(Map.of("sub-1",
Set.of(1)), Map.of(), Map.of()));
+ derived.remove("member-id");
+
+ assertEquals(Map.of("member-id", memberTasks),
streamsGroup.refinedAssignment(5));
+ }
+
+ @Test
+ public void testRelabelRefinedAssignment() {
+ StreamsGroup streamsGroup = createStreamsGroup("foo");
+
+ TasksTuple oldMemberTasks = new TasksTuple(Map.of("sub-1", Set.of(0)),
Map.of(), Map.of());
+ TasksTuple otherMemberTasks = new TasksTuple(Map.of("sub-1",
Set.of(1)), Map.of(), Map.of());
+ streamsGroup.setRefinedAssignment(5, Map.of(
+ "old-member-id", oldMemberTasks,
+ "other-member-id", otherMemberTasks
+ ));
+
+ streamsGroup.relabelRefinedAssignment("old-member-id",
"new-member-id");
+
+ // The replaced member's slice moves to its new ID and every other
member's slice is left exactly as derived:
+ // the epoch's decisions must not change while the group reconciles
towards it.
+ assertEquals(
+ Map.of("new-member-id", oldMemberTasks, "other-member-id",
otherMemberTasks),
+ streamsGroup.refinedAssignment(5)
+ );
+ // Re-keying does not make the entry apply to a different epoch.
+ assertNull(streamsGroup.refinedAssignment(6));
+ }
+
+ @Test
+ public void testRelabelRefinedAssignmentOfUnknownMemberIsANoOp() {
+ StreamsGroup streamsGroup = createStreamsGroup("foo");
+
+ Map<String, TasksTuple> refinedAssignment = Map.of(
+ "member-id", new TasksTuple(Map.of("sub-1", Set.of(0)), Map.of(),
Map.of())
+ );
+ streamsGroup.setRefinedAssignment(5, refinedAssignment);
+
+ // Nothing was derived for that member, so there is nothing to re-key
-- and in particular the cached
+ // assignment of the epoch must not be disturbed.
+ streamsGroup.relabelRefinedAssignment("not-in-the-cache",
"new-member-id");
+
+ assertEquals(refinedAssignment, streamsGroup.refinedAssignment(5));
+ }
+
+ @Test
+ public void testRefinedAssignmentCacheHoldsOnlyTheLatestEpoch() {
+ StreamsGroup streamsGroup = createStreamsGroup("foo");
+
+ Map<String, TasksTuple> refinedAssignmentOfEpoch5 = Map.of(
+ "member-id", new TasksTuple(Map.of("sub-1", Set.of(0)), Map.of(),
Map.of())
+ );
+ Map<String, TasksTuple> refinedAssignmentOfEpoch6 = Map.of(
+ "member-id", new TasksTuple(Map.of("sub-1", Set.of(0, 1)),
Map.of(), Map.of())
+ );
+
+ streamsGroup.setRefinedAssignment(5, refinedAssignmentOfEpoch5);
+ streamsGroup.setRefinedAssignment(6, refinedAssignmentOfEpoch6);
+
+ // The cache is a single slot, so deriving for a new epoch evicts the
previous one rather than accumulating an
+ // entry per epoch. Only the epoch the members are currently
reconciling towards is ever asked for.
+ assertEquals(refinedAssignmentOfEpoch6,
streamsGroup.refinedAssignment(6));
+ assertNull(streamsGroup.refinedAssignment(5));
+ }
+
@Test
public void testRemoveMemberClearsTaskOffsets() {
StreamsGroup streamsGroup = createStreamsGroup("foo");
