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 8a74ae6df248c..d6a78fc67b524 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 @@ -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(); diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java index 13d4d3631da26..606690ce15745 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java +++ b/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(); }