KAFKA-20719: "Streams" group sticky assignor not sticky at startup

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<TaskId, ArrayList<MemberAndOffsetSum>> tasksWithReportedOffsets = new HashMap<>();
+        for (final Map.Entry<String, AssignmentMemberSpec> memberEntry : groupSpec.members().entrySet()) {
+            final AssignmentMemberSpec memberSpec = memberEntry.getValue();
+            final Member member = new Member(memberSpec.processId(), memberEntry.getKey());
+            for (final Map.Entry<TaskId, Long> 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<TaskId, ArrayList<MemberAndOffsetSum>> entry : tasksWithReportedOffsets.entrySet()) {
+            final ArrayList<Member> candidates =
+                localState.standbyTaskToPrevMember.computeIfAbsent(entry.getKey(), task -> new ArrayList<>());
+            final Set<String> 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<String> 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<TaskId, Member> activeTaskToPrevMember;

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<String, AssignmentMemberSpec> 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<TaskId, Long> process1Offsets = mkMap(
+            mkEntry(new TaskId("test-subtopology", 0), 100L), mkEntry(new TaskId("test-subtopology", 1), 100L));
+        final Map<TaskId, Long> process2Offsets = mkMap(
+            mkEntry(new TaskId("test-subtopology", 2), 100L), mkEntry(new TaskId("test-subtopology", 3), 100L));
+        final Map<String, AssignmentMemberSpec> 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<String, AssignmentMemberSpec> 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<String, AssignmentMemberSpec> 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<String, AssignmentMemberSpec> 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<String, AssignmentMemberSpec> 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<String, AssignmentMemberSpec> 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<Integer> getActiveTasks(GroupAssignment result, final String topolog
         return res;
     }
 
+    private Set<Integer> getStandbyTasks(GroupAssignment result, final String topologyId, String... memberIds) {
+        Set<Integer> 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<Integer> getAllActiveTaskIds(GroupAssignment result, String... memberIds) {
         List<Integer> 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<TaskId, Long> 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<String, Set<Integer>> prevActiveTasks,
                                                             final Map<String, Set<Integer>> prevStandbyTasks) {
         return new AssignmentMemberSpec(

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 {

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