[flink-xxx][runtime] support dynamic rate limit source (PoC)

docs/layouts/shortcodes/generated/source_configuration.html

@@ -0,0 +1,72 @@
+<table class="configuration table table-bordered">
+    <thead>
+        <tr>
+            <th class="text-left" style="width: 20%">Key</th>
+            <th class="text-left" style="width: 15%">Default</th>
+            <th class="text-left" style="width: 10%">Type</th>
+            <th class="text-left" style="width: 55%">Description</th>
+        </tr>
+    </thead>
+    <tbody>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.allocator</h5></td>
+            <td style="word-wrap: break-word;">"demand-aware"</td>
+            <td>String</td>
+            <td>How the global target rate is split across the source subtasks. 'demand-aware' is a max-min fair allocation based on the throughput each subtask reports, which copes with partition skew. Any other value is read as the fully qualified name of a RateAllocator with a public no-arg constructor.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.calculator</h5></td>
+            <td style="word-wrap: break-word;">"checkpoint-duration"</td>
+            <td>String</td>
+            <td>How the global target rate is derived after each checkpoint. 'checkpoint-duration' throttles down when checkpoint durations trend up and ramps up otherwise. Any other value is read as the fully qualified name of a RateLimitCalculator with a public no-arg constructor.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.enabled</h5></td>
+            <td style="word-wrap: break-word;">false</td>
+            <td>Boolean</td>
+            <td>Whether the source coordinator should track the rolling trend of checkpoint completion times and dynamically broadcast an adjusted target rate limit to source readers that opt in via RateLimiterStrategy#dynamic(double).</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.max</h5></td>
+            <td style="word-wrap: break-word;">0.5</td>
+            <td>Double</td>
+            <td>Upper bound of the checkpoint duration band, as a fraction of the checkpoint interval. While the latest checkpoint duration is above this value the rate is throttled down immediately.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.max-records-per-second</h5></td>
+            <td style="word-wrap: break-word;">1.7976931348623157E308</td>
+            <td>Double</td>
+            <td>The ceiling (records per second, across all subtasks) that the dynamic rate limit is allowed to ramp back up to.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.min</h5></td>
+            <td style="word-wrap: break-word;">0.1</td>
+            <td>Double</td>
+            <td>Lower bound of the checkpoint duration band, as a fraction of the checkpoint interval. While the latest checkpoint duration is below this value the rate is left unchanged.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.ramp-up-factor</h5></td>
+            <td style="word-wrap: break-word;">0.3</td>
+            <td>Double</td>
+            <td>The factor the current target rate is multiplied by when checkpoint durations are flat or improving, capped at source.dynamic-rate-limit.max-records-per-second, e.g. 1.1 raises the rate by 30%.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.rolling-window-size</h5></td>
+            <td style="word-wrap: break-word;">3</td>
+            <td>Integer</td>
+            <td>The number of most recent checkpoint durations to keep when computing the checkpoint-duration trend.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.target</h5></td>
+            <td style="word-wrap: break-word;">0.3</td>
+            <td>Double</td>
+            <td>The desired checkpoint duration, as a fraction of the checkpoint interval, between source.dynamic-rate-limit.min and source.dynamic-rate-limit.max. The rate is throttled down while the average duration over the rolling window is above it and ramped up otherwise. Only applies when the checkpoint interval is known; otherwise the rate follows the checkpoint duration trend.</td>
+        </tr>
+        <tr>
+            <td><h5>source.dynamic-rate-limit.throttle-down-factor</h5></td>
+            <td style="word-wrap: break-word;">0.8</td>
+            <td>Double</td>
+            <td>The factor the current target rate is multiplied by when checkpoint durations are trending up beyond the configured threshold, e.g. 0.8 cuts the rate by 20%.</td>
+        </tr>
+    </tbody>
+</table>

flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/SourceReaderBase.java

@@ -26,8 +26,10 @@
 import org.apache.flink.api.connector.source.SourceReader;
 import org.apache.flink.api.connector.source.SourceReaderContext;
 import org.apache.flink.api.connector.source.SourceSplit;
+import org.apache.flink.api.connector.source.util.ratelimit.DynamicRateLimiter;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiter;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
+import org.apache.flink.api.connector.source.util.ratelimit.SupportsDynamicRateLimit;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.base.source.reader.fetcher.SplitFetcherManager;
 import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
@@ -72,7 +74,7 @@
  */
 @PublicEvolving
 public abstract class SourceReaderBase<E, T, SplitT extends SourceSplit, SplitStateT>
-        implements SourceReader<T, SplitT> {
+        implements SourceReader<T, SplitT>, SupportsDynamicRateLimit {
     private static final Logger LOG = LoggerFactory.getLogger(SourceReaderBase.class);
 
     /** A queue to buffer the elements fetched by the fetcher thread. */
@@ -118,6 +120,8 @@ public abstract class SourceReaderBase<E, T, SplitT extends SourceSplit, SplitSt
     /** Future that tracks the result of acquiring permission from {@link #rateLimiter}. */
     @Nullable private CompletableFuture<Void> rateLimitPermissionFuture;
 
+    private boolean warnedAboutStaticRateLimiter;
+
     /**
      * The primary constructor for the source reader.
      *
@@ -385,6 +389,20 @@ public void handleSourceEvents(SourceEvent sourceEvent) {
         LOG.info("Received unhandled source event: {}", sourceEvent);
     }
 
+    @Override
+    public void updateTargetRate(double recordsPerSecond) {
+        if (rateLimiter instanceof DynamicRateLimiter) {
+            ((DynamicRateLimiter<SplitT>) rateLimiter).updateRate(recordsPerSecond);
+        } else if (!warnedAboutStaticRateLimiter) {
+            warnedAboutStaticRateLimiter = true;
+            LOG.warn(
+                    "Dynamic rate limit is enabled but the rate limiter {} does not support "
+                            + "updating its rate; use RateLimiterStrategy#dynamic. The new target "
+                            + "rate is ignored.",
+                    rateLimiter == null ? "(none)" : rateLimiter.getClass().getName());
+        }
+    }
+
     @Override
     public void pauseOrResumeSplits(
             Collection<String> splitsToPause, Collection<String> splitsToResume) {

flink-connectors/flink-connector-base/src/test/java/org/apache/flink/connector/base/source/reader/SourceReaderBaseTest.java

@@ -26,7 +26,9 @@
 import org.apache.flink.api.connector.source.SourceSplit;
 import org.apache.flink.api.connector.source.mocks.MockSourceSplit;
 import org.apache.flink.api.connector.source.mocks.MockSourceSplitSerializer;
+import org.apache.flink.api.connector.source.util.ratelimit.DynamicRateLimiter;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
+import org.apache.flink.api.connector.source.util.ratelimit.SupportsDynamicRateLimit;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.base.source.reader.fetcher.SingleThreadFetcherManager;
 import org.apache.flink.connector.base.source.reader.mocks.MockSourceReader;
@@ -70,12 +72,14 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static org.apache.flink.streaming.api.operators.source.TestingSourceOperator.createTestOperator;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** A unit test class for {@link SourceReaderBase}. */
@@ -176,6 +180,37 @@ void testLimitingRateInSplitReader() throws Exception {
                 .isLessThan(Duration.ofSeconds(35).toMillis());
     }
 
+    @Test
+    void testUpdateTargetRateReachesDynamicRateLimiter() throws Exception {
+        final AtomicReference<Double> updatedRate = new AtomicReference<>();
+        final RateLimiterStrategy<TestingSourceSplit> strategy =
+                parallelism -> new RecordingDynamicRateLimiter(updatedRate);
+        final SourceReader<?, ?> reader =
+                createReaderAndAwaitAvailable(
+                        Collections.singletonList("test-split"),
+                        Collections.singletonList(
+                                new TestingRecordsWithSplitIds<>("test-split", "value")),
+                        strategy);
+
+        assertThat(reader).isInstanceOf(SupportsDynamicRateLimit.class);
+        ((SupportsDynamicRateLimit) reader).updateTargetRate(42.0);
+
+        assertThat(updatedRate).hasValue(42.0);
+    }
+
+    @Test
+    void testUpdateTargetRateIgnoredForStaticRateLimiter() throws Exception {
+        final SourceReader<?, ?> reader =
+                createReaderAndAwaitAvailable(
+                        Collections.singletonList("test-split"),
+                        Collections.singletonList(
+                                new TestingRecordsWithSplitIds<>("test-split", "value")),
+                        RateLimiterStrategy.perSecond(2));
+
+        assertThatCode(() -> ((SupportsDynamicRateLimit) reader).updateTargetRate(42.0))
+                .doesNotThrowAnyException();
+    }
+
     @Test
     void testLimitingRatePerCheckpointInSplitReader() throws Exception {
         String[] recordArr = new String[30];
@@ -535,6 +570,25 @@ private Configuration getConfig() {
         return config;
     }
 
+    private static class RecordingDynamicRateLimiter
+            implements DynamicRateLimiter<TestingSourceSplit> {
+        private final AtomicReference<Double> updatedRate;
+
+        private RecordingDynamicRateLimiter(AtomicReference<Double> updatedRate) {
+            this.updatedRate = updatedRate;
+        }
+
+        @Override
+        public CompletionStage<Void> acquire(int numberOfEvents) {
+            return CompletableFuture.completedFuture(null);
+        }
+
+        @Override
+        public void updateRate(double recordsPerSecond) {
+            updatedRate.set(recordsPerSecond);
+        }
+    }
+
     // ------------------------------------------------------------------------
     //  Testing Setup Helpers
     // ------------------------------------------------------------------------

flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/CheckpointDurationRateLimitCalculator.java

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.api.connector.source.util.ratelimit;
+
+import org.apache.flink.annotation.Internal;
+
+import java.util.List;
+
+/**
+ * Computes a new target source rate (records per second, across all subtasks) from a rolling window
+ * of recent checkpoint durations, compared against fractions of the checkpoint interval. Below
+ * {@code min} nothing is adjusted and above {@code max} the rate is throttled down immediately.
+ * Within that band the average duration of the window is compared against {@code target}: above it
+ * the rate is throttled down, otherwise it is ramped up toward the configured maximum.
+ *
+ * <p>If the checkpoint interval is unknown, the fractions cannot be resolved and the average
+ * duration of the more recent half of the window is compared against the older half instead: a
+ * rising trend throttles the rate down, while a flat or improving trend ramps the rate back up.
+ *
+ * <p>The rate is scaled off the actually observed throughput rather than the previously broadcast
+ * target rate, so that the adjustment reacts to reality directly instead of compounding off a value
+ * that may already be stale, converging faster.
+ *
+ * <p>This is a pure function with no coordinator dependencies, making it independently
+ * unit-testable.
+ */
+@Internal
+public class CheckpointDurationRateLimitCalculator implements RateLimitCalculator {
+    private static final long serialVersionUID = 1L;
+
+    private static final double MIN_RATE = 1.0;
+
+    /**
+     * @param recentDurationsMillis rolling window of checkpoint durations, oldest first.
+     * @param currentThroughput the current observed throughput (records/sec, summed across
+     *     subtasks) to scale the adjustment off of; callers should fall back to the last target
+     *     rate when no throughput has been observed yet.
+     * @param checkpointIntervalMillis the checkpoint interval the duration band is a fraction of;
+     *     non-positive or unbounded values fall back to the trend comparison.
+     * @param params the configured thresholds and factors.
+     * @return the new target rate (records/sec), clamped to {@code [MIN_RATE,
+     *     params.getMaxRecordsPerSecond()]}.
+     */
+    @Override
+    public double computeNewRate(
+            List<Long> recentDurationsMillis,
+            double currentThroughput,
+            long checkpointIntervalMillis,
+            DynamicRateLimitParams params) {
+        if (recentDurationsMillis.isEmpty()) {
+            return currentThroughput;
+        }
+
+        final boolean shouldThrottle;
+        if (hasKnownInterval(checkpointIntervalMillis)) {
+            final long latestDurationMillis =
+                    recentDurationsMillis.get(recentDurationsMillis.size() - 1);
+            if (latestDurationMillis < params.getMinDurationFraction() * checkpointIntervalMillis) {
+                return currentThroughput;
+            }
+            shouldThrottle =
+                    latestDurationMillis
+                                    > params.getMaxDurationFraction() * checkpointIntervalMillis
+                            || average(recentDurationsMillis)
+                                    > params.getTargetDurationFraction() * checkpointIntervalMillis;
+        } else {
+            if (recentDurationsMillis.size() < 2) {
+                // not enough data points to establish a trend
+                return currentThroughput;
+            }
+            shouldThrottle = computeTrend(recentDurationsMillis) > 0;
+        }
+
+        final double newRate;
+        if (shouldThrottle) {
+            newRate = currentThroughput * params.getThrottleDownFactor();
+        } else {
+            // ramp up from the observed throughput, capped at the configured max
+            newRate =
+                    Math.min(
+                            currentThroughput * params.getRampUpFactor(),
+                            params.getMaxRecordsPerSecond());
+        }
+        return Math.max(newRate, MIN_RATE);
+    }
+
+    private boolean hasKnownInterval(long checkpointIntervalMillis) {
+        return checkpointIntervalMillis > 0 && checkpointIntervalMillis != Long.MAX_VALUE;
+    }
+
+    /**
+     * Returns the relative change between the average of the older half and the average of the more
+     * recent half of the window, e.g. {@code 0.1} means checkpoints got 10% slower.
+     */
+    private double computeTrend(List<Long> durationsMillis) {
+        int mid = durationsMillis.size() / 2;
+        double olderHalfAvg = average(durationsMillis.subList(0, mid));
+        double recentHalfAvg = average(durationsMillis.subList(mid, durationsMillis.size()));
+        if (olderHalfAvg == 0) {
+            return recentHalfAvg == 0 ? 0 : Double.POSITIVE_INFINITY;
+        }
+        return (recentHalfAvg - olderHalfAvg) / olderHalfAvg;
+    }
+
+    private double average(List<Long> values) {
+        return values.stream().mapToLong(Long::longValue).average().orElse(0);
+    }
+}