KAFKA-20711: Measure restore-remaining-records in offset slots

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

@@ -205,7 +205,7 @@ public void clearTaskTimeout() {
     }
 
     @Override
-    public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
+    public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
         throw new UnsupportedOperationException("This task is read-only");
     }
 

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

@@ -96,7 +96,7 @@ public boolean isActive() {
     }
 
     @Override
-    public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
+    public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
         if (initRemaining) {
             throw new IllegalStateException("Standby task would not record remaining records to restore");
         }

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

@@ -141,6 +141,9 @@ static class ChangelogMetadata {
 
         private long restoreStartTimeNs;
 
+        // the consumer position at which restoration started, so progress can be measured in offset slots
+        private long restoreStartOffset;
+
         private ChangelogMetadata(final StateStoreMetadata storeMetadata, final ProcessorStateManager stateManager) {
             this.changelogState = ChangelogState.REGISTERED;
             this.storeMetadata = storeMetadata;
@@ -668,6 +671,8 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
 
         if (numRecords != 0) {
             final List<ConsumerRecord<byte[], byte[]>> records = changelogMetadata.bufferedRecords.subList(0, numRecords);
+            // where restoration had reached before this batch; null until the first batch is restored
+            final Long offsetBeforeRestore = storeMetadata.offset();
             final OptionalLong optionalLag = restoreConsumer.currentLag(partition);
             stateManager.restore(storeMetadata, records, optionalLag);
 
@@ -680,9 +685,9 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
                 changelogMetadata.bufferedRecords.clear();
             }
 
-            task.recordRestoration(time, numRecords, false);
-
             final Long currentOffset = storeMetadata.offset();
+            recordRestorationProgress(task, changelogMetadata, numRecords, offsetBeforeRestore, currentOffset);
+
             log.trace("Restored {} records from changelog {} to store {}, end offset is {}, current offset is {}",
                 numRecords, partition, storeName, recordEndOffset(changelogMetadata.restoreEndOffset), currentOffset);
 
@@ -710,6 +715,10 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
             log.info("Finished restoring changelog {} to store {} with a total number of {} records",
                 partition, storeName, changelogMetadata.totalRestored);
 
+            // account for any offset slots past the last restored record (e.g. trailing transaction
+            // markers) so the remaining-records metric reaches exactly zero on completion
+            recordRestorationProgress(task, changelogMetadata, 0, storeMetadata.offset(), changelogMetadata.restoreEndOffset - 1);
+
             changelogMetadata.transitTo(ChangelogState.COMPLETED);
             pauseChangelogsFromRestoreConsumer(Collections.singleton(partition));
             if (storeMetadata.store() instanceof MeteredStateStore) {
@@ -730,6 +739,23 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
         return numRecords;
     }
 
+    /**
+     * Record restoration progress: restore-total/restore-rate advance by the records restored
+     * ({@code numRecords}), while the remaining-records metric is decremented by the offset slots
+     * between {@code previousOffset} (or {@code restoreStartOffset} if null) and {@code restoredToOffset}.
+     * Measuring the latter in offset slots accounts for offsets the restore consumer never returns
+     * (transaction markers, compacted records) so it reaches exactly zero on completion.
+     */
+    private void recordRestorationProgress(final Task task,
+                                           final ChangelogMetadata changelogMetadata,
+                                           final long numRecords,
+                                           final Long previousOffset,
+                                           final long restoredToOffset) {
+        final long restoredFrom = previousOffset == null ? changelogMetadata.restoreStartOffset - 1 : previousOffset;
+        final long numOffsets = Math.max(restoredToOffset - restoredFrom, 0L);
+        task.recordRestoration(time, numRecords, numOffsets, false);
+    }
+
     private Set<Task> getTasksFromPartitions(final Map<TaskId, Task> tasks,
                                              final Set<TopicPartition> partitions) {
         final Set<Task> result = new HashSet<>();
@@ -1033,6 +1059,8 @@ private void prepareChangelogs(final Map<TaskId, Task> tasks,
                 throw new StreamsException("Restore consumer get unexpected error trying to get the position " +
                         " of " + partition, e);
             }
+            // remember where restoration began so progress can be measured in offset slots
+            changelogMetadata.restoreStartOffset = startOffset;
             if (changelogMetadata.stateManager.taskType() == Task.TaskType.ACTIVE) {
                 try {
                     stateRestoreListener.onRestoreStart(partition, storeName, startOffset, changelogMetadata.restoreEndOffset);
@@ -1045,8 +1073,8 @@ private void prepareChangelogs(final Map<TaskId, Task> tasks,
                 // if the log is truncated between when we get the log end offset and when we get the
                 // consumer position, then it's possible that the difference become negative and there's actually
                 // no records to restore; in this case we just initialize the sensor to zero
-                final long recordsToRestore = Math.max(changelogMetadata.restoreEndOffset - startOffset, 0L);
-                task.recordRestoration(time, recordsToRestore, true);
+                final long offsetsToRestore = Math.max(changelogMetadata.restoreEndOffset - startOffset, 0L);
+                task.recordRestoration(time, 0, offsetsToRestore, true);
                 changelogMetadata.restoreStartTimeNs = time.nanoseconds();
             }  else if (changelogMetadata.stateManager.taskType() == TaskType.STANDBY) {
                 try {

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

@@ -254,12 +254,12 @@ public boolean isActive() {
     }
 
     @Override
-    public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
+    public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
         if (initRemaining) {
-            maybeRecordSensor(numRecords, time, restoreRemainingSensor);
+            maybeRecordSensor(numOffsets, time, restoreRemainingSensor);
         } else {
             maybeRecordSensor(numRecords, time, restoreSensor);
-            maybeRecordSensor(-1 * numRecords, time, restoreRemainingSensor);
+            maybeRecordSensor(-1 * numOffsets, time, restoreRemainingSensor);
         }
     }
 

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

@@ -213,7 +213,11 @@ void maybeInitTaskTimeoutOrThrow(final long currentWallClockMs,
 
     void clearTaskTimeout();
 
-    void recordRestoration(final Time time, final long numRecords, final boolean initRemaining);
+    /**
+     * {@code numRecords} feeds restore-total/restore-rate; {@code numOffsets} feeds the offset-based
+     * remaining-records metric (or initialises it when {@code initRemaining} is true).
+     */
+    void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining);
 
     // task status inquiry
 

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

@@ -484,12 +484,13 @@ public void shouldRecordRestoredRecords() {
         assertThat(totalMetric.metricValue(), equalTo(0.0));
         assertThat(rateMetric.metricValue(), equalTo(0.0));
 
-        task.recordRestoration(time, 25L, false);
+        // standby tasks have no remaining-records metric, so the offset-slot argument is ignored
+        task.recordRestoration(time, 25L, 30L, false);
 
         assertThat(totalMetric.metricValue(), equalTo(25.0));
         assertThat(rateMetric.metricValue(), not(0.0));
 
-        task.recordRestoration(time, 50L, false);
+        task.recordRestoration(time, 50L, 55L, false);
 
         assertThat(totalMetric.metricValue(), equalTo(75.0));
         assertThat(rateMetric.metricValue(), not(0.0));

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

@@ -54,6 +54,7 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -94,6 +95,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -663,6 +665,75 @@ public void shouldCheckCompletionIfPositionLargerThanEndOffset() {
         assertNull(callback.storeNameCalledStates.get(RESTORE_BATCH));
     }
 
+    @Test
+    public void shouldDecrementRemainingRecordsToZeroWithOffsetGaps() {
+        // end offset is 10 but the changelog holds only 8 data records (offsets 0..7); offsets 8 and 9
+        // are transaction markers the restore consumer never returns, so remaining-records is initialized
+        // to 10 offset slots and must still decrement to exactly zero once restoration completes
+        setupActiveStateManager();
+        setupStoreMetadata();
+        setupStore();
+        final TaskId taskId = new TaskId(0, 0);
+
+        // null while preparing (seek-to-beginning, so startOffset == 0) and before the first batch;
+        // once a batch is applied it reflects the last restored record's offset (7)
+        when(storeMetadata.offset()).thenReturn(null).thenReturn(null).thenReturn(7L);
+        when(activeStateManager.taskId()).thenReturn(taskId);
+
+        consumer.updateBeginningOffsets(Collections.singletonMap(tp, 0L));
+        adminClient.updateEndOffsets(Collections.singletonMap(tp, 10L));
+
+        final Task mockTask = mock(Task.class);
+
+        changelogReader.register(tp, activeStateManager);
+
+        // first restore initializes the changelog and records the initial remaining (= 10 slots)
+        changelogReader.restore(Collections.singletonMap(taskId, mockTask));
+        assertEquals(0L, consumer.position(tp));
+        assertEquals(StoreChangelogReader.ChangelogState.RESTORING, changelogReader.changelogMetadata(tp).state());
+
+        for (int offset = 0; offset < 8; offset++) {
+            consumer.addRecord(new ConsumerRecord<>(topicName, 0, offset, "key".getBytes(), "value".getBytes()));
+        }
+
+        // second restore applies the 8 data records
+        changelogReader.restore(Collections.singletonMap(taskId, mockTask));
+        assertEquals(8L, changelogReader.changelogMetadata(tp).totalRestored());
+        assertEquals(StoreChangelogReader.ChangelogState.RESTORING, changelogReader.changelogMetadata(tp).state());
+
+        // skip the trailing transaction-marker offsets (8, 9) so the consumer reaches the end
+        consumer.seek(tp, 10L);
+
+        // third restore observes position >= endOffset and completes restoration
+        changelogReader.restore(Collections.singletonMap(taskId, mockTask));
+        assertEquals(StoreChangelogReader.ChangelogState.COMPLETED, changelogReader.changelogMetadata(tp).state());
+        assertEquals(8L, changelogReader.changelogMetadata(tp).totalRestored());
+
+        // remaining-records is driven by numOffsets (init sets it, later calls decrement it); restore-total
+        // is driven by numRecords, a distinct quantity when the changelog has offset gaps
+        final ArgumentCaptor<Long> numRecordsCaptor = ArgumentCaptor.forClass(Long.class);
+        final ArgumentCaptor<Long> numOffsetsCaptor = ArgumentCaptor.forClass(Long.class);
+        final ArgumentCaptor<Boolean> initCaptor = ArgumentCaptor.forClass(Boolean.class);
+        verify(mockTask, atLeastOnce())
+            .recordRestoration(any(), numRecordsCaptor.capture(), numOffsetsCaptor.capture(), initCaptor.capture());
+
+        long initialRemaining = 0L;
+        long decrementedRemaining = 0L;
+        long totalRecords = 0L;
+        for (int i = 0; i < initCaptor.getAllValues().size(); i++) {
+            if (initCaptor.getAllValues().get(i)) {
+                initialRemaining += numOffsetsCaptor.getAllValues().get(i);
+            } else {
+                decrementedRemaining += numOffsetsCaptor.getAllValues().get(i);
+                totalRecords += numRecordsCaptor.getAllValues().get(i);
+            }
+        }
+
+        assertEquals(10L, initialRemaining, "remaining-records metric should be initialized from the offset range");
+        assertEquals(initialRemaining, decrementedRemaining, "remaining-records metric should decrement to exactly zero");
+        assertEquals(8L, totalRecords, "restore-total should count the records actually restored");
+    }
+
     @Test
     public void shouldRequestPositionAndHandleTimeoutException() {
         setupActiveStateManager();
@@ -700,7 +771,7 @@ public long position(final TopicPartition partition) {
         assertEquals(10L, (long) changelogReader.changelogMetadata(tp).endOffset());
         Mockito.verify(mockTask).clearTaskTimeout();
         Mockito.verify(mockTask).maybeInitTaskTimeoutOrThrow(anyLong(), any());
-        Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyBoolean());
+        Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyLong(), anyBoolean());
 
         clearException.set(true);
         Mockito.reset(mockTask);
@@ -798,7 +869,7 @@ public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition
         assertEquals(10L, (long) changelogReader.changelogMetadata(tp).endOffset());
         assertEquals(6L, consumer.position(tp));
         Mockito.verify(mockTask).clearTaskTimeout();
-        Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyBoolean());
+        Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyLong(), anyBoolean());
     }
 
     @Test
@@ -893,7 +964,7 @@ public synchronized ListConsumerGroupOffsetsResult listConsumerGroupOffsets(fina
         assertEquals(6L, consumer.position(tp));
         if (type == ACTIVE) {
             Mockito.verify(mockTask, times(2)).clearTaskTimeout();
-            Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyBoolean());
+            Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyLong(), anyBoolean());
         }
     }
 

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

@@ -886,21 +886,23 @@ public void shouldRecordRestoredRecords() {
         assertThat(rateMetric.metricValue(), equalTo(0.0));
         assertThat(remainMetric.metricValue(), equalTo(0.0));
 
-        task.recordRestoration(time, 100L, true);
+        // remaining-records is initialized from the offset range
+        task.recordRestoration(time, 0L, 100L, true);
 
         assertThat(remainMetric.metricValue(), equalTo(100.0));
 
-        task.recordRestoration(time, 25L, false);
+        // restore-total counts records; remaining-records is decremented by offset slots (which may differ)
+        task.recordRestoration(time, 25L, 30L, false);
 
         assertThat(totalMetric.metricValue(), equalTo(25.0));
         assertThat(rateMetric.metricValue(), not(0.0));
-        assertThat(remainMetric.metricValue(), equalTo(75.0));
+        assertThat(remainMetric.metricValue(), equalTo(70.0));
 
-        task.recordRestoration(time, 50L, false);
+        task.recordRestoration(time, 50L, 55L, false);
 
         assertThat(totalMetric.metricValue(), equalTo(75.0));
         assertThat(rateMetric.metricValue(), not(0.0));
-        assertThat(remainMetric.metricValue(), equalTo(25.0));
+        assertThat(remainMetric.metricValue(), equalTo(15.0));
     }
 
     @Test