KAFKA-20116: Fix reporting task-(end)-offset sum for "streams" protocol

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

@@ -1243,6 +1243,10 @@ void runOnceWithoutProcessingThreads() {
             // regardless of streamsGroupReady, as these may throw exceptions that need to be handled.
             handleStreamsRebalanceData();
 
+            // The group coordinator places tasks based on the offsets we report to it, so publish them on every
+            // iteration -- including the ones that return early below, where a task may well still be restoring.
+            taskManager.maybeUpdateTaskOffsetSumSnapshot();
+
             if (!streamsGroupReady) {
                 return;
             }
@@ -1268,7 +1272,6 @@ void runOnceWithoutProcessingThreads() {
         if (isStartingRunningOrPartitionAssigned()) {
 
             taskManager.updateLags();
-            taskManager.maybeUpdateTaskOffsetSumSnapshot();
 
             /*
              * Within an iteration, after processing up to N (N initialized as 1 upon start up) records for each applicable tasks, check the current time:
@@ -1398,6 +1401,10 @@ void runOnceWithProcessingThreads() {
             // regardless of streamsGroupReady, as these may throw exceptions that need to be handled.
             handleStreamsRebalanceData();
 
+            // The group coordinator places tasks based on the offsets we report to it, so publish them on every
+            // iteration -- including the ones that return early below, where a task may well still be restoring.
+            taskManager.maybeUpdateTaskOffsetSumSnapshot();
+
             if (!streamsGroupReady) {
                 return;
             }
@@ -1416,7 +1423,6 @@ void runOnceWithProcessingThreads() {
         if (isRunning()) {
 
             taskManager.updateLags();
-            taskManager.maybeUpdateTaskOffsetSumSnapshot();
 
             checkStateUpdater();
 

streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java

@@ -109,6 +109,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.InOrder;
@@ -639,6 +640,48 @@ public void shouldAlsoPurgeBeforeTheCommitInterval(final boolean processingThrea
         verify(taskManager, times(2)).maybePurgeCommittedRecords();
     }
 
+    @ParameterizedTest
+    @CsvSource({"true, true", "true, false", "false, true", "false, false"})
+    public void shouldPublishTaskOffsetSumSnapshotOnEveryIteration(final boolean processingThreadsEnabled,
+                                                                   final boolean streamsGroupReady) {
+        // The group coordinator places tasks based on the offsets reported to it, so they have to be published on every
+        // iteration, including the ones that return early because the group is not ready -- which is exactly when a
+        // task may still be restoring.
+        final StreamsConfig config = new StreamsConfig(configProps(false, processingThreadsEnabled));
+        final TaskManager taskManager = mock(TaskManager.class);
+        when(mainConsumer.poll(Mockito.any())).thenReturn(ConsumerRecords.empty());
+        final ConsumerGroupMetadata consumerGroupMetadata = mock(ConsumerGroupMetadata.class);
+        when(mainConsumer.groupMetadata()).thenReturn(consumerGroupMetadata);
+        when(consumerGroupMetadata.groupInstanceId()).thenReturn(Optional.empty());
+
+        final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(
+            UUID.randomUUID(), Optional.empty(), Optional.empty(), Map.of(), Map.of(), Map::of, Map::of);
+        final TopologyMetadata topologyMetadata = new TopologyMetadata(internalTopologyBuilder, config);
+        topologyMetadata.buildAndRewriteTopology();
+        thread = new StreamThread(
+            mockTime, config, null,
+            mainConsumer, consumer,
+            changelogReader, null, taskManager, null,
+            new StreamsMetricsImpl(metrics, CLIENT_ID, mockTime),
+            topologyMetadata,
+            PROCESS_ID, CLIENT_ID, new LogContext(""),
+            new AtomicInteger(), new AtomicLong(Long.MAX_VALUE), new LinkedList<>(),
+            null, HANDLER, null,
+            Optional.of(streamsRebalanceData), mock(StreamsMetadataState.class), null, -1L
+        ).updateThreadMetadata(adminClientId(CLIENT_ID));
+        thread.setState(State.STARTING);
+        thread.setState(State.PARTITIONS_ASSIGNED);
+
+        thread.setStreamsGroupReady(streamsGroupReady);
+        if (processingThreadsEnabled) {
+            thread.runOnceWithProcessingThreads();
+        } else {
+            thread.runOnceWithoutProcessingThreads();
+        }
+
+        verify(taskManager).maybeUpdateTaskOffsetSumSnapshot();
+    }
+
     @Test
     public void shouldNotProcessWhenPartitionRevoked() {
         final Properties props = configProps(false, false);
@@ -3448,7 +3491,9 @@ public void shouldUpdateLagsAfterPolling(final boolean processingThreadsEnabled)
         final InOrder inOrder = Mockito.inOrder(mainConsumer, thread.taskManager());
         inOrder.verify(mainConsumer).poll(Mockito.any());
         inOrder.verify(thread.taskManager()).updateLags();
-        inOrder.verify(thread.taskManager()).maybeUpdateTaskOffsetSumSnapshot();
+        // The offset-sum snapshot is only read by the streams-protocol heartbeat thread, so under the classic protocol
+        // it is not published at all.
+        verify(thread.taskManager(), never()).maybeUpdateTaskOffsetSumSnapshot();
     }