Repository: apache/kafka
File: streams/src/main/java/org/apache/kafka/streams/processor/internals/StoreChangelogReader.java (Java)

Code before:
                final TopicPartition partition = storeMetadata.changelogPartition();
                final String storeName = storeMetadata.store().name();

                long startOffset = 0L;
                try {
                    startOffset = restoreConsumer.position(partition);
                } catch (final TimeoutException swallow) {
                    // if we cannot find the starting position at the beginning, just use the default 0L
                } catch (final KafkaException e) {
                    // this also includes InvalidOffsetException, which should not happen under normal
                    // execution, hence it is also okay to wrap it as fatal StreamsException
                    throw new StreamsException("Restore consumer get unexpected error trying to get the position " +
                        " of " + partition, e);
                }

                try {
                    stateRestoreListener.onRestoreStart(partition, storeName, startOffset, changelogMetadata.restoreEndOffset);
                } catch (final Exception e) {
                    throw new StreamsException("State restore listener failed on batch restored", e);
                }
            }
        }
    }

    @Override
    public void unregister(final Collection<TopicPartition> revokedChangelogs) {
        // Only changelogs that are initialized have been added to the restore consumer's assignment
        final List<TopicPartition> revokedInitializedChangelogs = new ArrayList<>();

        for (final TopicPartition partition : revokedChangelogs) {
            final ChangelogMetadata changelogMetadata = changelogs.remove(partition);
            if (changelogMetadata != null) {
                // if the changelog is still in REGISTERED, it means it has not initialized and started
                // restoring yet, and hence we should not try to remove the changelog partition
                if (!changelogMetadata.state().equals(ChangelogState.REGISTERED)) {
                    revokedInitializedChangelogs.add(partition);

                    // if the changelog is not in RESTORING, it means
                    // the corresponding onRestoreStart was not called; in this case
                    // we should not call onRestoreSuspended either
                    if (changelogMetadata.stateManager.taskType() == Task.TaskType.ACTIVE &&
                        changelogMetadata.state().equals(ChangelogState.RESTORING)) {
                        try {
                            final String storeName = changelogMetadata.storeMetadata.store().name();
                            stateRestoreListener.onRestoreSuspended(partition, storeName, changelogMetadata.totalRestored);
                        } catch (final Exception e) {
                            throw new StreamsException("State restore listener failed on restore paused", e);
                        }
                    }
                }


Review comment (suggestion):
The logic for measuring remaining records is a bit complex: we first aggregate the total amount of records to restore across all changelog partitions at the beginning when initializing the changelogs; and then during restoration we keep decrementing by the number of restored records.

Code after:
                final StateStoreMetadata storeMetadata = changelogMetadata.storeMetadata;
                final TopicPartition partition = storeMetadata.changelogPartition();
                final String storeName = storeMetadata.store().name();

                long startOffset = 0L;
                try {
                    startOffset = restoreConsumer.position(partition);
                } catch (final TimeoutException swallow) {
                    // if we cannot find the starting position at the beginning, just use the default 0L
                } catch (final KafkaException e) {
                    // this also includes InvalidOffsetException, which should not happen under normal
                    // execution, hence it is also okay to wrap it as fatal StreamsException
                    throw new StreamsException("Restore consumer get unexpected error trying to get the position " +
                        " of " + partition, e);
                }

                try {
                    stateRestoreListener.onRestoreStart(partition, storeName, startOffset, changelogMetadata.restoreEndOffset);
                } catch (final Exception e) {
                    throw new StreamsException("State restore listener failed on batch restored", e);
                }

                final TaskId taskId = changelogs.get(partition).stateManager.taskId();
                final StreamTask task = (StreamTask) tasks.get(taskId);
                // 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.initRemainingRecordsToRestore(time, recordsToRestore);
            }
        }
    }

    @Override
    public void unregister(final Collection<TopicPartition> revokedChangelogs) {
        // Only changelogs that are initialized have been added to the restore consumer's assignment
        final List<TopicPartition> revokedInitializedChangelogs = new ArrayList<>();

        for (final TopicPartition partition : revokedChangelogs) {
            final ChangelogMetadata changelogMetadata = changelogs.remove(partition);
            if (changelogMetadata != null) {
                // if the changelog is still in REGISTERED, it means it has not initialized and started
                // restoring yet, and hence we should not try to remove the changelog partition
                if (!changelogMetadata.state().equals(ChangelogState.REGISTERED)) {
                    revokedInitializedChangelogs.add(partition);

                    // if the changelog is not in RESTORING, it means
                    // the corresponding onRestoreStart was not called; in this case
                    // we should not call onRestoreSuspended either
                    if (changelogMetadata.stateManager.taskType() == Task.TaskType.ACTIVE &&
                        changelogMetadata.state().equals(ChangelogState.RESTORING)) {