Skip to content

Commit 3ad3020

Browse files
nicktelfordclaude
nicktelford
and
claude
committed
KAFKA-20711: Measure restore-remaining-records in offset slots
The restore-remaining-records-total metric is initialized from the changelog offset range (restoreEndOffset - startOffset), which counts every offset slot, but it was decremented by the number of records actually restored. These diverge whenever the changelog contains offsets the restore consumer never returns as records -- transaction markers or compacted-away records -- leaving the metric stuck above zero after restoration completed. The metric is now decremented in offset slots: each batch advances it by the change in store offset, and on completion any trailing slots between the last restored record and the end offset are accounted for, so it lands at exactly zero. Task#recordRestoration takes the offset-slot count separately from the record count so that restore-total and restore-rate continue to measure records restored. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent da3b5e7 commit 3ad3020

8 files changed

Lines changed: 126 additions & 20 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,7 @@ public void clearTaskTimeout() {
205205
}
206206

207207
@Override
208-
public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
208+
public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
209209
throw new UnsupportedOperationException("This task is read-only");
210210
}
211211

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ public boolean isActive() {
9696
}
9797

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

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

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,9 @@ static class ChangelogMetadata {
141141

142142
private long restoreStartTimeNs;
143143

144+
// the consumer position at which restoration started, so progress can be measured in offset slots
145+
private long restoreStartOffset;
146+
144147
private ChangelogMetadata(final StateStoreMetadata storeMetadata, final ProcessorStateManager stateManager) {
145148
this.changelogState = ChangelogState.REGISTERED;
146149
this.storeMetadata = storeMetadata;
@@ -668,6 +671,8 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
668671

669672
if (numRecords != 0) {
670673
final List<ConsumerRecord<byte[], byte[]>> records = changelogMetadata.bufferedRecords.subList(0, numRecords);
674+
// where restoration had reached before this batch; null until the first batch is restored
675+
final Long offsetBeforeRestore = storeMetadata.offset();
671676
final OptionalLong optionalLag = restoreConsumer.currentLag(partition);
672677
stateManager.restore(storeMetadata, records, optionalLag);
673678

@@ -680,9 +685,9 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
680685
changelogMetadata.bufferedRecords.clear();
681686
}
682687

683-
task.recordRestoration(time, numRecords, false);
684-
685688
final Long currentOffset = storeMetadata.offset();
689+
recordRestorationProgress(task, changelogMetadata, numRecords, offsetBeforeRestore, currentOffset);
690+
686691
log.trace("Restored {} records from changelog {} to store {}, end offset is {}, current offset is {}",
687692
numRecords, partition, storeName, recordEndOffset(changelogMetadata.restoreEndOffset), currentOffset);
688693

@@ -710,6 +715,10 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
710715
log.info("Finished restoring changelog {} to store {} with a total number of {} records",
711716
partition, storeName, changelogMetadata.totalRestored);
712717

718+
// account for any offset slots past the last restored record (e.g. trailing transaction
719+
// markers) so the remaining-records metric reaches exactly zero on completion
720+
recordRestorationProgress(task, changelogMetadata, 0, storeMetadata.offset(), changelogMetadata.restoreEndOffset - 1);
721+
713722
changelogMetadata.transitTo(ChangelogState.COMPLETED);
714723
pauseChangelogsFromRestoreConsumer(Collections.singleton(partition));
715724
if (storeMetadata.store() instanceof MeteredStateStore) {
@@ -730,6 +739,23 @@ private int restoreChangelog(final Task task, final ChangelogMetadata changelogM
730739
return numRecords;
731740
}
732741

742+
/**
743+
* Record restoration progress: restore-total/restore-rate advance by the records restored
744+
* ({@code numRecords}), while the remaining-records metric is decremented by the offset slots
745+
* between {@code previousOffset} (or {@code restoreStartOffset} if null) and {@code restoredToOffset}.
746+
* Measuring the latter in offset slots accounts for offsets the restore consumer never returns
747+
* (transaction markers, compacted records) so it reaches exactly zero on completion.
748+
*/
749+
private void recordRestorationProgress(final Task task,
750+
final ChangelogMetadata changelogMetadata,
751+
final long numRecords,
752+
final Long previousOffset,
753+
final long restoredToOffset) {
754+
final long restoredFrom = previousOffset == null ? changelogMetadata.restoreStartOffset - 1 : previousOffset;
755+
final long numOffsets = Math.max(restoredToOffset - restoredFrom, 0L);
756+
task.recordRestoration(time, numRecords, numOffsets, false);
757+
}
758+
733759
private Set<Task> getTasksFromPartitions(final Map<TaskId, Task> tasks,
734760
final Set<TopicPartition> partitions) {
735761
final Set<Task> result = new HashSet<>();
@@ -1033,6 +1059,8 @@ private void prepareChangelogs(final Map<TaskId, Task> tasks,
10331059
throw new StreamsException("Restore consumer get unexpected error trying to get the position " +
10341060
" of " + partition, e);
10351061
}
1062+
// remember where restoration began so progress can be measured in offset slots
1063+
changelogMetadata.restoreStartOffset = startOffset;
10361064
if (changelogMetadata.stateManager.taskType() == Task.TaskType.ACTIVE) {
10371065
try {
10381066
stateRestoreListener.onRestoreStart(partition, storeName, startOffset, changelogMetadata.restoreEndOffset);
@@ -1045,8 +1073,8 @@ private void prepareChangelogs(final Map<TaskId, Task> tasks,
10451073
// if the log is truncated between when we get the log end offset and when we get the
10461074
// consumer position, then it's possible that the difference become negative and there's actually
10471075
// no records to restore; in this case we just initialize the sensor to zero
1048-
final long recordsToRestore = Math.max(changelogMetadata.restoreEndOffset - startOffset, 0L);
1049-
task.recordRestoration(time, recordsToRestore, true);
1076+
final long offsetsToRestore = Math.max(changelogMetadata.restoreEndOffset - startOffset, 0L);
1077+
task.recordRestoration(time, 0, offsetsToRestore, true);
10501078
changelogMetadata.restoreStartTimeNs = time.nanoseconds();
10511079
} else if (changelogMetadata.stateManager.taskType() == TaskType.STANDBY) {
10521080
try {

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -254,12 +254,12 @@ public boolean isActive() {
254254
}
255255

256256
@Override
257-
public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
257+
public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
258258
if (initRemaining) {
259-
maybeRecordSensor(numRecords, time, restoreRemainingSensor);
259+
maybeRecordSensor(numOffsets, time, restoreRemainingSensor);
260260
} else {
261261
maybeRecordSensor(numRecords, time, restoreSensor);
262-
maybeRecordSensor(-1 * numRecords, time, restoreRemainingSensor);
262+
maybeRecordSensor(-1 * numOffsets, time, restoreRemainingSensor);
263263
}
264264
}
265265

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,11 @@ void maybeInitTaskTimeoutOrThrow(final long currentWallClockMs,
213213

214214
void clearTaskTimeout();
215215

216-
void recordRestoration(final Time time, final long numRecords, final boolean initRemaining);
216+
/**
217+
* {@code numRecords} feeds restore-total/restore-rate; {@code numOffsets} feeds the offset-based
218+
* remaining-records metric (or initialises it when {@code initRemaining} is true).
219+
*/
220+
void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining);
217221

218222
// task status inquiry
219223

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,12 +484,13 @@ public void shouldRecordRestoredRecords() {
484484
assertThat(totalMetric.metricValue(), equalTo(0.0));
485485
assertThat(rateMetric.metricValue(), equalTo(0.0));
486486

487-
task.recordRestoration(time, 25L, false);
487+
// standby tasks have no remaining-records metric, so the offset-slot argument is ignored
488+
task.recordRestoration(time, 25L, 30L, false);
488489

489490
assertThat(totalMetric.metricValue(), equalTo(25.0));
490491
assertThat(rateMetric.metricValue(), not(0.0));
491492

492-
task.recordRestoration(time, 50L, false);
493+
task.recordRestoration(time, 50L, 55L, false);
493494

494495
assertThat(totalMetric.metricValue(), equalTo(75.0));
495496
assertThat(rateMetric.metricValue(), not(0.0));

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

Lines changed: 74 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
import org.junit.jupiter.api.extension.ExtendWith;
5555
import org.junit.jupiter.params.ParameterizedTest;
5656
import org.junit.jupiter.params.provider.EnumSource;
57+
import org.mockito.ArgumentCaptor;
5758
import org.mockito.Mock;
5859
import org.mockito.Mockito;
5960
import org.mockito.junit.jupiter.MockitoExtension;
@@ -94,6 +95,7 @@
9495
import static org.mockito.ArgumentMatchers.any;
9596
import static org.mockito.ArgumentMatchers.anyBoolean;
9697
import static org.mockito.ArgumentMatchers.anyLong;
98+
import static org.mockito.Mockito.atLeastOnce;
9799
import static org.mockito.Mockito.mock;
98100
import static org.mockito.Mockito.never;
99101
import static org.mockito.Mockito.times;
@@ -663,6 +665,75 @@ public void shouldCheckCompletionIfPositionLargerThanEndOffset() {
663665
assertNull(callback.storeNameCalledStates.get(RESTORE_BATCH));
664666
}
665667

668+
@Test
669+
public void shouldDecrementRemainingRecordsToZeroWithOffsetGaps() {
670+
// end offset is 10 but the changelog holds only 8 data records (offsets 0..7); offsets 8 and 9
671+
// are transaction markers the restore consumer never returns, so remaining-records is initialized
672+
// to 10 offset slots and must still decrement to exactly zero once restoration completes
673+
setupActiveStateManager();
674+
setupStoreMetadata();
675+
setupStore();
676+
final TaskId taskId = new TaskId(0, 0);
677+
678+
// null while preparing (seek-to-beginning, so startOffset == 0) and before the first batch;
679+
// once a batch is applied it reflects the last restored record's offset (7)
680+
when(storeMetadata.offset()).thenReturn(null).thenReturn(null).thenReturn(7L);
681+
when(activeStateManager.taskId()).thenReturn(taskId);
682+
683+
consumer.updateBeginningOffsets(Collections.singletonMap(tp, 0L));
684+
adminClient.updateEndOffsets(Collections.singletonMap(tp, 10L));
685+
686+
final Task mockTask = mock(Task.class);
687+
688+
changelogReader.register(tp, activeStateManager);
689+
690+
// first restore initializes the changelog and records the initial remaining (= 10 slots)
691+
changelogReader.restore(Collections.singletonMap(taskId, mockTask));
692+
assertEquals(0L, consumer.position(tp));
693+
assertEquals(StoreChangelogReader.ChangelogState.RESTORING, changelogReader.changelogMetadata(tp).state());
694+
695+
for (int offset = 0; offset < 8; offset++) {
696+
consumer.addRecord(new ConsumerRecord<>(topicName, 0, offset, "key".getBytes(), "value".getBytes()));
697+
}
698+
699+
// second restore applies the 8 data records
700+
changelogReader.restore(Collections.singletonMap(taskId, mockTask));
701+
assertEquals(8L, changelogReader.changelogMetadata(tp).totalRestored());
702+
assertEquals(StoreChangelogReader.ChangelogState.RESTORING, changelogReader.changelogMetadata(tp).state());
703+
704+
// skip the trailing transaction-marker offsets (8, 9) so the consumer reaches the end
705+
consumer.seek(tp, 10L);
706+
707+
// third restore observes position >= endOffset and completes restoration
708+
changelogReader.restore(Collections.singletonMap(taskId, mockTask));
709+
assertEquals(StoreChangelogReader.ChangelogState.COMPLETED, changelogReader.changelogMetadata(tp).state());
710+
assertEquals(8L, changelogReader.changelogMetadata(tp).totalRestored());
711+
712+
// remaining-records is driven by numOffsets (init sets it, later calls decrement it); restore-total
713+
// is driven by numRecords, a distinct quantity when the changelog has offset gaps
714+
final ArgumentCaptor<Long> numRecordsCaptor = ArgumentCaptor.forClass(Long.class);
715+
final ArgumentCaptor<Long> numOffsetsCaptor = ArgumentCaptor.forClass(Long.class);
716+
final ArgumentCaptor<Boolean> initCaptor = ArgumentCaptor.forClass(Boolean.class);
717+
verify(mockTask, atLeastOnce())
718+
.recordRestoration(any(), numRecordsCaptor.capture(), numOffsetsCaptor.capture(), initCaptor.capture());
719+
720+
long initialRemaining = 0L;
721+
long decrementedRemaining = 0L;
722+
long totalRecords = 0L;
723+
for (int i = 0; i < initCaptor.getAllValues().size(); i++) {
724+
if (initCaptor.getAllValues().get(i)) {
725+
initialRemaining += numOffsetsCaptor.getAllValues().get(i);
726+
} else {
727+
decrementedRemaining += numOffsetsCaptor.getAllValues().get(i);
728+
totalRecords += numRecordsCaptor.getAllValues().get(i);
729+
}
730+
}
731+
732+
assertEquals(10L, initialRemaining, "remaining-records metric should be initialized from the offset range");
733+
assertEquals(initialRemaining, decrementedRemaining, "remaining-records metric should decrement to exactly zero");
734+
assertEquals(8L, totalRecords, "restore-total should count the records actually restored");
735+
}
736+
666737
@Test
667738
public void shouldRequestPositionAndHandleTimeoutException() {
668739
setupActiveStateManager();
@@ -700,7 +771,7 @@ public long position(final TopicPartition partition) {
700771
assertEquals(10L, (long) changelogReader.changelogMetadata(tp).endOffset());
701772
Mockito.verify(mockTask).clearTaskTimeout();
702773
Mockito.verify(mockTask).maybeInitTaskTimeoutOrThrow(anyLong(), any());
703-
Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyBoolean());
774+
Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyLong(), anyBoolean());
704775

705776
clearException.set(true);
706777
Mockito.reset(mockTask);
@@ -798,7 +869,7 @@ public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition
798869
assertEquals(10L, (long) changelogReader.changelogMetadata(tp).endOffset());
799870
assertEquals(6L, consumer.position(tp));
800871
Mockito.verify(mockTask).clearTaskTimeout();
801-
Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyBoolean());
872+
Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyLong(), anyBoolean());
802873
}
803874

804875
@Test
@@ -893,7 +964,7 @@ public synchronized ListConsumerGroupOffsetsResult listConsumerGroupOffsets(fina
893964
assertEquals(6L, consumer.position(tp));
894965
if (type == ACTIVE) {
895966
Mockito.verify(mockTask, times(2)).clearTaskTimeout();
896-
Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyBoolean());
967+
Mockito.verify(mockTask).recordRestoration(any(), anyLong(), anyLong(), anyBoolean());
897968
}
898969
}
899970

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -886,21 +886,23 @@ public void shouldRecordRestoredRecords() {
886886
assertThat(rateMetric.metricValue(), equalTo(0.0));
887887
assertThat(remainMetric.metricValue(), equalTo(0.0));
888888

889-
task.recordRestoration(time, 100L, true);
889+
// remaining-records is initialized from the offset range
890+
task.recordRestoration(time, 0L, 100L, true);
890891

891892
assertThat(remainMetric.metricValue(), equalTo(100.0));
892893

893-
task.recordRestoration(time, 25L, false);
894+
// restore-total counts records; remaining-records is decremented by offset slots (which may differ)
895+
task.recordRestoration(time, 25L, 30L, false);
894896

895897
assertThat(totalMetric.metricValue(), equalTo(25.0));
896898
assertThat(rateMetric.metricValue(), not(0.0));
897-
assertThat(remainMetric.metricValue(), equalTo(75.0));
899+
assertThat(remainMetric.metricValue(), equalTo(70.0));
898900

899-
task.recordRestoration(time, 50L, false);
901+
task.recordRestoration(time, 50L, 55L, false);
900902

901903
assertThat(totalMetric.metricValue(), equalTo(75.0));
902904
assertThat(rateMetric.metricValue(), not(0.0));
903-
assertThat(remainMetric.metricValue(), equalTo(25.0));
905+
assertThat(remainMetric.metricValue(), equalTo(15.0));
904906
}
905907

906908
@Test

0 commit comments

Comments
 (0)