MINOR: ChangelogReader should poll for duration 0 for standby restore

clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java

@@ -65,6 +65,7 @@ public class MockConsumer<K, V> implements Consumer<K, V> {
     private KafkaException pollException;
     private KafkaException offsetsException;
     private AtomicBoolean wakeup;
+    private Duration lastPollTimeout;
     private boolean closed;
     private boolean shouldRebalance;
 
@@ -157,13 +158,15 @@ public synchronized void unsubscribe() {
     @Deprecated
     @Override
     public synchronized ConsumerRecords<K, V> poll(long timeout) {
-        return poll(Duration.ZERO);
+        return poll(Duration.ofMillis(timeout));
     }
 
     @Override
     public synchronized ConsumerRecords<K, V> poll(final Duration timeout) {
         ensureNotClosed();
 
+        lastPollTimeout = timeout;
+
         // Synchronize around the entire execution so new tasks to be triggered on subsequent poll calls can be added in
         // the callback
         synchronized (pollTasks) {
@@ -556,6 +559,10 @@ public void resetShouldRebalance() {
         shouldRebalance = false;
     }
 
+    public Duration lastPollTimeout() {
+        return lastPollTimeout;
+    }
+
     @Override
     public void close(Duration timeout) {
         close();

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

@@ -216,9 +216,6 @@ public StoreChangelogReader(final Time time,
         this.restoreConsumer = restoreConsumer;
         this.stateRestoreListener = stateRestoreListener;
 
-        // NOTE for restoring active and updating standby we may prefer different poll time
-        // in order to make sure we call the main consumer#poll in time.
-        // TODO: once both of these are moved to a separate thread this may no longer be a concern
         this.pollTime = Duration.ofMillis(config.getLong(StreamsConfig.POLL_MS_CONFIG));
         this.updateOffsetIntervalMs = config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) == Long.MAX_VALUE ?
             DEFAULT_OFFSET_UPDATE_MS : config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG);
@@ -415,7 +412,10 @@ public void restore() {
             final ConsumerRecords<byte[], byte[]> polledRecords;
 
             try {
-                polledRecords = restoreConsumer.poll(pollTime);
+                // for restoring active and updating standby we may prefer different poll time
+                // in order to make sure we call the main consumer#poll in time.
+                // TODO: once we move ChangelogReader to a separate thread this may no longer be a concern
+                polledRecords = restoreConsumer.poll(state.equals(ChangelogReaderState.STANDBY_UPDATING) ? Duration.ZERO : pollTime);
             } catch (final InvalidOffsetException e) {
                 log.warn("Encountered {} fetching records from restore consumer for partitions {}, it is likely that " +
                     "the consumer's position has fallen out of the topic partition offset range because the topic was " +

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

@@ -46,6 +46,7 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 
+import java.time.Duration;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
@@ -226,6 +227,37 @@ public Map<TopicPartition, Long> endOffsets(final Collection<TopicPartition> par
         }
     }
 
+    @Test
+    public void shouldPollWithRightTimeout() {
+        EasyMock.expect(storeMetadata.offset()).andReturn(null).andReturn(9L).anyTimes();
+        EasyMock.replay(stateManager, storeMetadata, store);
+
+        final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {
+            @Override
+            public Map<TopicPartition, Long> endOffsets(final Collection<TopicPartition> partitions) {
+                return partitions.stream().collect(Collectors.toMap(Function.identity(), partition -> 11L));
+            }
+        };
+        consumer.updateBeginningOffsets(Collections.singletonMap(tp, 5L));
+
+        final StoreChangelogReader changelogReader =
+                new StoreChangelogReader(time, config, logContext, consumer, callback);
+
+        changelogReader.register(tp, stateManager);
+
+        if (type == STANDBY) {
+            changelogReader.transitToUpdateStandby();
+        }
+
+        changelogReader.restore();
+
+        if (type == ACTIVE) {
+            assertEquals(Duration.ofMillis(config.getLong(StreamsConfig.POLL_MS_CONFIG)), consumer.lastPollTimeout());
+        } else {
+            assertEquals(Duration.ZERO, consumer.lastPollTimeout());
+        }
+    }
+
     @Test
     public void shouldRestoreFromPositionAndCheckForCompletion() {
         EasyMock.expect(storeMetadata.offset()).andReturn(5L).anyTimes();