Skip to content

Commit e2aa6d6

Browse files
nicktelfordnicktelford
authored
KAFKA-20711: Measure restore-remaining-records in offset slots (#22614)
The `restore-remaining-records-total` task metric overshoots and never reaches zero. The code does initialize the metric based on offset difference (`endOffset - startOffset`), but counts down based on number of records restored, not taking "missing offsets" from log compaction, TX markers, or aborted messages (under EOS) into account. This fix makes the decrements consistent with the initialization by expressing the remaining-records metric's progress in changelog offset slots rather than record counts. Reviewers: Matthias J. Sax <matthias@confluent.io>
1 parent e7e947e commit e2aa6d6

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
@@ -210,7 +210,7 @@ public void clearTaskTimeout() {
210210
}
211211

212212
@Override
213-
public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
213+
public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
214214
throw new UnsupportedOperationException("This task is read-only");
215215
}
216216

‎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);
688+
final long currentOffset = storeMetadata.offset();
689+
recordRestorationProgress(task, changelogMetadata, numRecords, offsetBeforeRestore, currentOffset + 1);
684690

685-
final Long currentOffset = storeMetadata.offset();
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);
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 lastRestoredOffset} (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 lastRestoredOffset, // this is not a "position" so we need to correct it below
753+
final long restoredToOffset) {
754+
final long restoredFromOffset = lastRestoredOffset == null ? changelogMetadata.restoreStartOffset : lastRestoredOffset + 1;
755+
final long numOffsets = Math.max(restoredToOffset - restoredFromOffset, 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
@@ -258,12 +258,12 @@ public boolean isActive() {
258258
}
259259

260260
@Override
261-
public void recordRestoration(final Time time, final long numRecords, final boolean initRemaining) {
261+
public void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining) {
262262
if (initRemaining) {
263-
maybeRecordSensor(numRecords, time, restoreRemainingSensor);
263+
maybeRecordSensor(numOffsets, time, restoreRemainingSensor);
264264
} else {
265265
maybeRecordSensor(numRecords, time, restoreSensor);
266-
maybeRecordSensor(-1 * numRecords, time, restoreRemainingSensor);
266+
maybeRecordSensor(-1 * numOffsets, time, restoreRemainingSensor);
267267
}
268268
}
269269

‎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
@@ -223,7 +223,11 @@ void maybeInitTaskTimeoutOrThrow(final long currentWallClockMs,
223223

224224
void clearTaskTimeout();
225225

226-
void recordRestoration(final Time time, final long numRecords, final boolean initRemaining);
226+
/**
227+
* {@code numRecords} feeds restore-total/restore-rate; {@code numOffsets} feeds the offset-based
228+
* remaining-records metric (or initialises it when {@code initRemaining} is true).
229+
*/
230+
void recordRestoration(final Time time, final long numRecords, final long numOffsets, final boolean initRemaining);
227231

228232
// task status inquiry
229233

‎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
@@ -1142,21 +1142,23 @@ public void shouldRecordRestoredRecords() {
11421142
assertThat(rateMetric.metricValue(), equalTo(0.0));
11431143
assertThat(remainMetric.metricValue(), equalTo(0.0));
11441144

1145-
task.recordRestoration(time, 100L, true);
1145+
// remaining-records is initialized from the offset range
1146+
task.recordRestoration(time, 0L, 100L, true);
11461147

11471148
assertThat(remainMetric.metricValue(), equalTo(100.0));
11481149

1149-
task.recordRestoration(time, 25L, false);
1150+
// restore-total counts records; remaining-records is decremented by offset slots (which may differ)
1151+
task.recordRestoration(time, 25L, 30L, false);
11501152

11511153
assertThat(totalMetric.metricValue(), equalTo(25.0));
11521154
assertThat(rateMetric.metricValue(), not(0.0));
1153-
assertThat(remainMetric.metricValue(), equalTo(75.0));
1155+
assertThat(remainMetric.metricValue(), equalTo(70.0));
11541156

1155-
task.recordRestoration(time, 50L, false);
1157+
task.recordRestoration(time, 50L, 55L, false);
11561158

11571159
assertThat(totalMetric.metricValue(), equalTo(75.0));
11581160
assertThat(rateMetric.metricValue(), not(0.0));
1159-
assertThat(remainMetric.metricValue(), equalTo(25.0));
1161+
assertThat(remainMetric.metricValue(), equalTo(15.0));
11601162
}
11611163

11621164
@Test

0 commit comments

Comments
 (0)