[FLINK-36947][Connectors/Kinesis] Fix issue where excessive GetRecords in PollingKinesisShardSplitReader calls are made on idle source causing throttling

[FLINK-36939][Connectors/Kinesis] Fix issue where excessive BlockingQueue.poll() in FanOutKinesisShardSplitReader are made on idle source causing high CPU utilisation

Author: leekeiabstraction

flink-connector-aws/flink-connector-aws-kinesis-streams/pom.xml

@@ -153,6 +153,12 @@ under the License.
             <artifactId>flink-architecture-tests-test</artifactId>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.awaitility</groupId>
+            <artifactId>awaitility</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/KinesisStreamsSource.java

@@ -81,7 +81,6 @@
 import java.util.function.Supplier;
 
 import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_CONSUMER_NAME;
-import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_CONSUMER_SUBSCRIPTION_TIMEOUT;
 import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_DESCRIBE_CONSUMER_RETRY_STRATEGY_MAX_ATTEMPTS_OPTION;
 import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_DESCRIBE_CONSUMER_RETRY_STRATEGY_MAX_DELAY_OPTION;
 import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_DESCRIBE_CONSUMER_RETRY_STRATEGY_MIN_DELAY_OPTION;
@@ -228,7 +227,7 @@ private Supplier<SplitReader<Record, KinesisShardSplit>> getKinesisShardSplitRea
                                 createKinesisAsyncStreamProxy(streamArn, sourceConfig),
                                 consumerArn,
                                 shardMetricGroupMap,
-                                sourceConfig.get(EFO_CONSUMER_SUBSCRIPTION_TIMEOUT));
+                                sourceConfig);
             default:
                 throw new IllegalArgumentException("Unsupported reader type: " + readerType);
         }

flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/config/KinesisSourceConfigOptions.java

@@ -85,6 +85,13 @@ public enum ConsumerLifecycle {
                     .defaultValue(ReaderType.POLLING)
                     .withDescription("The type of reader used to read from the Kinesis stream.");
 
+    public static final ConfigOption<Duration> READER_EMPTY_RECORDS_FETCH_INTERVAL =
+            ConfigOptions.key("source.reader.empty-records-fetch-interval")
+                    .durationType()
+                    .defaultValue(Duration.ofMillis(250))
+                    .withDescription(
+                            "The interval in milliseconds between fetches with empty records");
+
     public static final ConfigOption<ConsumerLifecycle> EFO_CONSUMER_LIFECYCLE =
             ConfigOptions.key("source.efo.lifecycle")
                     .enumType(ConsumerLifecycle.class)

flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/reader/KinesisShardSplitReaderBase.java

@@ -19,9 +19,11 @@
 package org.apache.flink.connector.kinesis.source.reader;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitsChange;
+import org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions;
 import org.apache.flink.connector.kinesis.source.metrics.KinesisShardMetrics;
 import org.apache.flink.connector.kinesis.source.split.KinesisShardSplit;
 import org.apache.flink.connector.kinesis.source.split.KinesisShardSplitState;
@@ -38,12 +40,14 @@
 import java.util.ArrayDeque;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Date;
 import java.util.Deque;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.WeakHashMap;
 
 import static java.util.Collections.singleton;
 
@@ -60,16 +64,30 @@ public abstract class KinesisShardSplitReaderBase
     private final Set<String> pausedSplitIds = new HashSet<>();
     private final Map<String, KinesisShardMetrics> shardMetricGroupMap;
 
-    protected KinesisShardSplitReaderBase(Map<String, KinesisShardMetrics> shardMetricGroupMap) {
+    private final long emptyRecordsIntervalMillis;
+
+    private final Map<KinesisShardSplitState, Long> scheduledFetchTimes = new WeakHashMap<>();
+
+    protected KinesisShardSplitReaderBase(
+            Map<String, KinesisShardMetrics> shardMetricGroupMap, Configuration configuration) {
         this.shardMetricGroupMap = shardMetricGroupMap;
+        this.emptyRecordsIntervalMillis =
+                configuration
+                        .get(KinesisSourceConfigOptions.READER_EMPTY_RECORDS_FETCH_INTERVAL)
+                        .toMillis();
     }
 
     @Override
     public RecordsWithSplitIds<Record> fetch() throws IOException {
         KinesisShardSplitState splitState = assignedSplits.poll();
 
         // When there are no assigned splits, return quickly
-        if (splitState == null) {
+        if (skipWhenNoAssignedSplit(splitState)) {
+            return INCOMPLETE_SHARD_EMPTY_RECORDS;
+        }
+
+        if (skipUntilScheduledFetchTime(splitState)) {
+            assignedSplits.add(splitState);
             return INCOMPLETE_SHARD_EMPTY_RECORDS;
         }
 
@@ -82,6 +100,7 @@ public RecordsWithSplitIds<Record> fetch() throws IOException {
         RecordBatch recordBatch;
         try {
             recordBatch = fetchRecords(splitState);
+            scheduleNextFetchTime(splitState, recordBatch);
         } catch (ResourceNotFoundException e) {
             LOG.warn(
                     "Failed to fetch records from shard {}: shard no longer exists. Marking split as complete",
@@ -125,6 +144,66 @@ public RecordsWithSplitIds<Record> fetch() throws IOException {
                 recordBatch.isCompleted());
     }
 
+    private boolean skipWhenNoAssignedSplit(KinesisShardSplitState splitState) throws IOException {
+        if (splitState == null) {
+            try {
+                // Small sleep to prevent busy polling
+                Thread.sleep(1);
+                return true;
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new IOException("Sleep was interrupted while skipping no assigned split", e);
+            }
+        }
+
+        return false;
+    }
+
+    private boolean skipUntilScheduledFetchTime(KinesisShardSplitState splitState)
+            throws IOException {
+        if (scheduledFetchTimes.containsKey(splitState)
+                && scheduledFetchTimes.get(splitState) > System.currentTimeMillis()) {
+            try {
+                // Small sleep to prevent busy polling
+                Thread.sleep(1);
+                return true;
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new IOException(
+                        "Sleep was interrupted while skipping until scheduled fetch record time",
+                        e);
+            }
+        }
+
+        return false;
+    }
+
+    /**
+     * Schedules next fetch time, to be called immediately on the result of a fetchRecords() call.
+     *
+     * <p>If recordBatch does not contain records, next fetchRecords() is scheduled. Before
+     * scheduled time, fetcher thread will skip fetching (and have small sleep) for the split.
+     *
+     * <p>If recordBatch is not empty, next fetchRecords() time is not scheduled resulting in next
+     * fetch on the split is performed at first opportunity.
+     *
+     * @param splitState splitState on which the fetchRecords() was called on
+     * @param recordBatch recordBatch returned by fetchRecords()
+     */
+    private void scheduleNextFetchTime(KinesisShardSplitState splitState, RecordBatch recordBatch) {
+        if (recordBatch == null || recordBatch.getRecords().isEmpty()) {
+            long scheduledGetRecordTimeMillis =
+                    System.currentTimeMillis() + emptyRecordsIntervalMillis;
+            this.scheduledFetchTimes.put(splitState, scheduledGetRecordTimeMillis);
+            if (LOG.isDebugEnabled()) {
+                LOG.debug(
+                        "Fetched zero records from split {}, scheduling next fetch to {}",
+                        splitState.getSplitId(),
+                        new Date(scheduledGetRecordTimeMillis).toInstant());
+            }
+        }
+    }
+
     /**
      * Main method implementations must implement to fetch records from Kinesis.
      *

flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/reader/fanout/FanOutKinesisShardSplitReader.java

@@ -19,6 +19,7 @@
 package org.apache.flink.connector.kinesis.source.reader.fanout;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitsChange;
 import org.apache.flink.connector.kinesis.source.metrics.KinesisShardMetrics;
 import org.apache.flink.connector.kinesis.source.proxy.AsyncStreamProxy;
@@ -32,6 +33,8 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.apache.flink.connector.kinesis.source.config.KinesisSourceConfigOptions.EFO_CONSUMER_SUBSCRIPTION_TIMEOUT;
+
 /**
  * An implementation of the KinesisShardSplitReader that consumes from Kinesis using Enhanced
  * Fan-Out and HTTP/2.
@@ -48,11 +51,11 @@ public FanOutKinesisShardSplitReader(
             AsyncStreamProxy asyncStreamProxy,
             String consumerArn,
             Map<String, KinesisShardMetrics> shardMetricGroupMap,
-            Duration subscriptionTimeout) {
-        super(shardMetricGroupMap);
+            Configuration configuration) {
+        super(shardMetricGroupMap, configuration);
         this.asyncStreamProxy = asyncStreamProxy;
         this.consumerArn = consumerArn;
-        this.subscriptionTimeout = subscriptionTimeout;
+        this.subscriptionTimeout = configuration.get(EFO_CONSUMER_SUBSCRIPTION_TIMEOUT);
     }
 
     @Override

flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/reader/polling/PollingKinesisShardSplitReader.java

@@ -44,7 +44,7 @@ public PollingKinesisShardSplitReader(
             StreamProxy kinesisProxy,
             Map<String, KinesisShardMetrics> shardMetricGroupMap,
             Configuration configuration) {
-        super(shardMetricGroupMap);
+        super(shardMetricGroupMap, configuration);
         this.kinesis = kinesisProxy;
         this.configuration = configuration;
         this.maxRecordsToGet = configuration.get(KinesisSourceConfigOptions.SHARD_GET_RECORDS_MAX);