diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java index 4828f0238b266..7d1bfe7aee2e9 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java @@ -133,6 +133,43 @@ private void initialize(final GroupSpec groupSpec, final TopologyDescriber topol } } } + + maybeAddOffsetBasedPrevMembers(groupSpec); + } + + /** + * A member that rejoins after a restart gets a fresh member ID and an empty target assignment, so the previous + * ownership maps cannot capture that it owned any tasks before the restart. Such members report cumulative + * offsets for tasks with state on local disk; treat them like previous standbys of those tasks, so that the + * stickiness phases assign the tasks back to the local state, most caught-up state first. Members already known + * as previous standbys keep precedence over offset-based candidates. + */ + private void maybeAddOffsetBasedPrevMembers(final GroupSpec groupSpec) { + final Map> tasksWithReportedOffsets = new HashMap<>(); + for (final Map.Entry memberEntry : groupSpec.members().entrySet()) { + final AssignmentMemberSpec memberSpec = memberEntry.getValue(); + final Member member = new Member(memberSpec.processId(), memberEntry.getKey()); + for (final Map.Entry taskOffset : memberSpec.taskOffsets().entrySet()) { + if (taskOffset.getValue() < 0) { + continue; + } + tasksWithReportedOffsets + .computeIfAbsent(taskOffset.getKey(), task -> new ArrayList<>()) + .add(new MemberAndOffsetSum(member, taskOffset.getValue())); + } + } + + for (final Map.Entry> entry : tasksWithReportedOffsets.entrySet()) { + final ArrayList candidates = + localState.standbyTaskToPrevMember.computeIfAbsent(entry.getKey(), task -> new ArrayList<>()); + final Set knownMemberIds = candidates.stream().map(candidate -> candidate.memberId).collect(Collectors.toSet()); + entry.getValue().sort(Comparator.comparingLong(MemberAndOffsetSum::offsetSum).reversed()); + for (final MemberAndOffsetSum memberAndOffsetSum : entry.getValue()) { + if (knownMemberIds.add(memberAndOffsetSum.member().memberId)) { + candidates.add(memberAndOffsetSum.member()); + } + } + } } private GroupAssignment buildGroupAssignment(final Set members) { @@ -409,6 +446,9 @@ public Member(final String processId, final String memberId) { } } + private record MemberAndOffsetSum(Member member, long offsetSum) { + } + private static class LocalState { // helper data structures: Map activeTaskToPrevMember; diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java index 161f3e6a5d808..c6a76c858d083 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java @@ -1301,6 +1301,164 @@ public void shouldAssignStandbyTaskToPreviousOwnerBasedOnBelowQuotaCondition() { assertTrue(member1TotalTasks <= 2, "Member1 should have <= 2 total tasks (quota), but has " + member1TotalTasks); } + @Test + public void shouldAssignTasksToMembersReportingTaskOffsetsOnColdStart() { + // Cold start (KAFKA-20719): all members have fresh member IDs and empty target assignments, but report + // cumulative offsets for the tasks whose state they hold on local disk. + final AssignmentMemberSpec memberSpec1 = createAssignmentMemberSpecWithOffsets("process1", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 100L), mkEntry(new TaskId("test-subtopology", 1), 100L))); + final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpecWithOffsets("process2", + mkMap(mkEntry(new TaskId("test-subtopology", 2), 100L), mkEntry(new TaskId("test-subtopology", 3), 100L))); + final Map members = mkMap( + mkEntry("member1", memberSpec1), + mkEntry("member2", memberSpec2)); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, new HashMap<>()), + new TopologyDescriberImpl(4, true, List.of("test-subtopology")) + ); + + assertEquals(Sets.newSet(0, 1), getActiveTasks(result, "test-subtopology", "member1")); + assertEquals(Sets.newSet(2, 3), getActiveTasks(result, "test-subtopology", "member2")); + } + + @Test + public void shouldAssignTasksToProcessesReportingTaskOffsetsOnColdStartWithMultipleMembersPerProcess() { + // Members of the same process share the state directory and report the same offsets; tasks should stay + // on the process holding the state, spread across its members. + final Map process1Offsets = mkMap( + mkEntry(new TaskId("test-subtopology", 0), 100L), mkEntry(new TaskId("test-subtopology", 1), 100L)); + final Map process2Offsets = mkMap( + mkEntry(new TaskId("test-subtopology", 2), 100L), mkEntry(new TaskId("test-subtopology", 3), 100L)); + final Map members = mkMap( + mkEntry("member1_1", createAssignmentMemberSpecWithOffsets("process1", process1Offsets)), + mkEntry("member1_2", createAssignmentMemberSpecWithOffsets("process1", process1Offsets)), + mkEntry("member2_1", createAssignmentMemberSpecWithOffsets("process2", process2Offsets)), + mkEntry("member2_2", createAssignmentMemberSpecWithOffsets("process2", process2Offsets))); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, new HashMap<>()), + new TopologyDescriberImpl(4, true, List.of("test-subtopology")) + ); + + assertEquals(mkMap(mkEntry("test-subtopology", Sets.newSet(0, 1))), + mergeAllActiveTasks(result, "member1_1", "member1_2")); + assertEquals(mkMap(mkEntry("test-subtopology", Sets.newSet(2, 3))), + mergeAllActiveTasks(result, "member2_1", "member2_2")); + } + + @Test + public void shouldPreferPreviousActiveOwnerOverReportedTaskOffsets() { + // A live member owning the task in the current target assignment beats any reported on-disk offsets. + final AssignmentMemberSpec memberSpec1 = createAssignmentMemberSpec("process1", + mkMap(mkEntry("test-subtopology", Set.of(0))), Map.of()); + final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpecWithOffsets("process2", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 1_000_000L))); + final Map members = mkMap( + mkEntry("member1", memberSpec1), + mkEntry("member2", memberSpec2)); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, new HashMap<>()), + new TopologyDescriberImpl(2, true, List.of("test-subtopology")) + ); + + assertEquals(Set.of(0), getActiveTasks(result, "test-subtopology", "member1")); + assertEquals(Set.of(1), getActiveTasks(result, "test-subtopology", "member2")); + } + + @Test + public void shouldPreferPreviousStandbyOverReportedTaskOffsets() { + // A member that is standby in the current target assignment beats an offset-based candidate, even if the + // standby member also reports offsets for the task (which must not register it as a duplicate candidate). + final AssignmentMemberSpec memberSpec1 = new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Map.of(), + mkMap(mkEntry("test-subtopology", Set.of(0))), + Map.of(), + "process1", + Map.of(), + mkMap(mkEntry(new TaskId("test-subtopology", 0), 500L)), + Map.of()); + final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpecWithOffsets("process2", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 1_000_000L))); + final Map members = mkMap( + mkEntry("member1", memberSpec1), + mkEntry("member2", memberSpec2)); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, new HashMap<>()), + new TopologyDescriberImpl(2, true, List.of("test-subtopology")) + ); + + assertEquals(Set.of(0), getActiveTasks(result, "test-subtopology", "member1")); + assertEquals(Set.of(1), getActiveTasks(result, "test-subtopology", "member2")); + } + + @Test + public void shouldPreferMemberReportingHigherTaskOffsetSum() { + // When several processes hold state for the same task, the most caught-up one wins. + final AssignmentMemberSpec memberSpec1 = createAssignmentMemberSpecWithOffsets("process1", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 50L))); + final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpecWithOffsets("process2", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 100L))); + final Map members = mkMap( + mkEntry("member1", memberSpec1), + mkEntry("member2", memberSpec2)); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, new HashMap<>()), + new TopologyDescriberImpl(2, true, List.of("test-subtopology")) + ); + + assertEquals(Set.of(1), getActiveTasks(result, "test-subtopology", "member1")); + assertEquals(Set.of(0), getActiveTasks(result, "test-subtopology", "member2")); + } + + @Test + public void shouldIgnoreNegativeReportedTaskOffsetSums() { + // Negative offset sums are sentinel/unknown values and must not grant stickiness over a valid report. + final AssignmentMemberSpec memberSpec1 = createAssignmentMemberSpecWithOffsets("process1", + mkMap(mkEntry(new TaskId("test-subtopology", 0), -2L), mkEntry(new TaskId("test-subtopology", 1), 100L))); + final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpecWithOffsets("process2", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 100L))); + final Map members = mkMap( + mkEntry("member1", memberSpec1), + mkEntry("member2", memberSpec2)); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, new HashMap<>()), + new TopologyDescriberImpl(2, true, List.of("test-subtopology")) + ); + + assertEquals(Set.of(1), getActiveTasks(result, "test-subtopology", "member1")); + assertEquals(Set.of(0), getActiveTasks(result, "test-subtopology", "member2")); + } + + @Test + public void shouldAssignStandbysAwayFromProcessHoldingStateOnColdStart() { + // On cold start with standbys enabled, actives follow the reported state and standbys land on the + // other process (a standby may not be co-located with its active). + final AssignmentMemberSpec memberSpec1 = createAssignmentMemberSpecWithOffsets("process1", + mkMap(mkEntry(new TaskId("test-subtopology", 0), 100L))); + final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpecWithOffsets("process2", + mkMap(mkEntry(new TaskId("test-subtopology", 1), 100L))); + final Map members = mkMap( + mkEntry("member1", memberSpec1), + mkEntry("member2", memberSpec2)); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, mkMap(mkEntry(NUM_STANDBY_REPLICAS_CONFIG, "1"))), + new TopologyDescriberImpl(2, true, List.of("test-subtopology")) + ); + + assertEquals(Set.of(0), getActiveTasks(result, "test-subtopology", "member1")); + assertEquals(Set.of(1), getActiveTasks(result, "test-subtopology", "member2")); + assertEquals(Set.of(1), getStandbyTasks(result, "test-subtopology", "member1")); + assertEquals(Set.of(0), getStandbyTasks(result, "test-subtopology", "member2")); + } + private int getAllActiveTaskCount(GroupAssignment result, String... memberIds) { int size = 0; @@ -1330,6 +1488,19 @@ private Set getActiveTasks(GroupAssignment result, final String topolog return res; } + private Set getStandbyTasks(GroupAssignment result, final String topologyId, String... memberIds) { + Set res = new HashSet<>(); + for (String memberId : memberIds) { + final MemberAssignment testMember = result.members().get(memberId); + assertNotNull(testMember); + assertNotNull(testMember.standbyTasks()); + if (testMember.standbyTasks().get(topologyId) != null) { + res.addAll(testMember.standbyTasks().get(topologyId)); + } + } + return res; + } + private List getAllActiveTaskIds(GroupAssignment result, String... memberIds) { List res = new ArrayList<>(); for (String memberId : memberIds) { @@ -1442,6 +1613,19 @@ private AssignmentMemberSpec createAssignmentMemberSpec(final String processId) Map.of()); } + private AssignmentMemberSpec createAssignmentMemberSpecWithOffsets(final String processId, final Map taskOffsets) { + return new AssignmentMemberSpec( + Optional.empty(), + Optional.empty(), + Map.of(), + Map.of(), + Map.of(), + processId, + Map.of(), + taskOffsets, + Map.of()); + } + private AssignmentMemberSpec createAssignmentMemberSpec(final String processId, final Map> prevActiveTasks, final Map> prevStandbyTasks) { return new AssignmentMemberSpec( diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/ColdStartStickinessIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/ColdStartStickinessIntegrationTest.java index e9d406cb150ef..271a07a511f62 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/ColdStartStickinessIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/ColdStartStickinessIntegrationTest.java @@ -45,7 +45,6 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.TestInfo; import org.junit.jupiter.api.Timeout; @@ -129,7 +128,6 @@ public void cleanup() throws Exception { CLUSTER.deleteAllTopics(); } - @Disabled("Reproduces KAFKA-20719; enable once fixed") @ParameterizedTest @ValueSource(booleans = {false, true}) public void shouldStickToLocalStateOnColdStart(final boolean streamsProtocol) throws Exception { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java index bb85da626b0d8..d211a701f8011 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java @@ -967,6 +967,10 @@ public void run() { * @throws StreamsException if the store's change log does not contain the partition */ boolean runLoop() { + // Populate the task-offset-sum snapshot before subscribing: subscribing triggers the join heartbeat, + // which must already carry the offset sums of tasks discovered in the local state directory, so that + // the broker-side sticky assignor can assign those tasks back to this client on a cold start. + taskManager.maybeUpdateTaskOffsetSumSnapshot(); subscribeConsumer(); // if the thread is still in the middle of a rebalance, we should keep polling