Address comments.

Author: lvyanquan

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

@@ -290,7 +290,7 @@ record -> {
                     };
         }
         if (isRateLimited) {
-            rateLimiter.notifyStatusChange(
+            rateLimiter.notifyAddingSplit(
                     this.toSplitType(
                             this.currentSplitContext.splitId, this.currentSplitContext.state));
         }
@@ -516,7 +516,7 @@ private static final class RateLimitingSourceOutputWrapper<T> implements SourceO
         final SourceOutput<T> sourceOutput;
 
         /** Count of records handled during the current rate-limiting window. */
-        int currentWindowRecordCount;
+        private int currentWindowRecordCount;
 
         /**
          * Creates a new RecordCountingSourceOutputWrapper.

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

@@ -205,7 +205,7 @@ void testLimitingRatePerCheckpointInSplitReader() throws Exception {
     }
 
     @Test
-    void testLimitingRateWithStatusChangeInSplitReader() throws Exception {
+    void testLimitingRateWithAddingSplitInSplitReader() throws Exception {
         String[] recordArr = new String[60];
         for (int i = 0; i < recordArr.length; i++) {
             recordArr[i] = "value" + i;
@@ -237,7 +237,7 @@ void testLimitingRateWithStatusChangeInSplitReader() throws Exception {
     /** A rate limiter that reduce the maxPerSecond for specific splits. */
     private static class SplitAwaredRateLimiter
             implements org.apache.flink.api.connector.source.util.ratelimit.RateLimiter<
-                    org.apache.flink.connector.base.source.reader.mocks.TestingSourceSplit> {
+                    TestingSourceSplit> {
 
         private final Executor limiter =
                 Executors.newSingleThreadExecutor(new ExecutorThreadFactory("flink-rate-limiter"));
@@ -255,8 +255,7 @@ public CompletionStage<Void> acquire(int requestSize) {
         }
 
         @Override
-        public void notifyStatusChange(
-                org.apache.flink.connector.base.source.reader.mocks.TestingSourceSplit split) {
+        public void notifyAddingSplit(TestingSourceSplit split) {
             if (!split.splitId().equals("test-split1")) {
                 this.rateLimiter = RateLimiter.create(maxPerSecond / 2);
             }

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

@@ -19,6 +19,7 @@
 package org.apache.flink.api.connector.source.util.ratelimit;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.connector.source.SourceSplit;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -31,7 +32,7 @@
  * external notifications.
  */
 @Internal
-public class GatedRateLimiter<S> implements RateLimiter<S> {
+public class GatedRateLimiter<Split extends SourceSplit> implements RateLimiter<Split> {
 
     private final int capacityPerCycle;
     private int capacityLeft;

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

@@ -19,6 +19,7 @@
 package org.apache.flink.api.connector.source.util.ratelimit;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.connector.source.SourceSplit;
 import org.apache.flink.util.concurrent.ExecutorThreadFactory;
 
 import org.apache.flink.shaded.guava33.com.google.common.util.concurrent.RateLimiter;
@@ -30,8 +31,8 @@
 
 /** An implementation of {@link RateLimiter} based on Guava's RateLimiter. */
 @Internal
-public class GuavaRateLimiter<S>
-        implements org.apache.flink.api.connector.source.util.ratelimit.RateLimiter<S> {
+public class GuavaRateLimiter<Split extends SourceSplit>
+        implements org.apache.flink.api.connector.source.util.ratelimit.RateLimiter<Split> {
 
     private final Executor limiter =
             Executors.newSingleThreadExecutor(new ExecutorThreadFactory("flink-rate-limiter"));

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

@@ -19,13 +19,14 @@
 package org.apache.flink.api.connector.source.util.ratelimit;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.connector.source.SourceSplit;
 import org.apache.flink.util.concurrent.FutureUtils;
 
 import java.util.concurrent.CompletionStage;
 
 /** A convenience implementation of {@link RateLimiter} that does not throttle requests. */
 @Internal
-public class NoOpRateLimiter<S> implements RateLimiter<S> {
+public class NoOpRateLimiter<Split extends SourceSplit> implements RateLimiter<Split> {
 
     @Override
     public CompletionStage<Void> acquire(int requestSize) {

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

@@ -19,15 +19,20 @@
 package org.apache.flink.api.connector.source.util.ratelimit;
 
 import org.apache.flink.annotation.Experimental;
+import org.apache.flink.api.connector.source.SourceSplit;
 
 import javax.annotation.concurrent.NotThreadSafe;
 
 import java.util.concurrent.CompletionStage;
 
-/** The interface to rate limit execution of methods. */
+/**
+ * The interface to rate limit execution of methods.
+ *
+ * @param <SplitT> The type of the source splits.
+ */
 @NotThreadSafe
 @Experimental
-public interface RateLimiter<S> {
+public interface RateLimiter<SplitT extends SourceSplit> {
 
     /**
      * Returns a future that is completed once another event would not exceed the rate limit. For
@@ -57,10 +62,8 @@ default CompletionStage<Void> acquire() {
     default void notifyCheckpointComplete(long checkpointId) {}
 
     /**
-     * Notifies this {@code RateLimiter} that the status has changed. This can be used to adjust the
-     * rate limiting behavior based on the new status.
-     *
-     * @param status The new status.
+     * Notifies this {@code RateLimiter} that a new split has been added. the result of before
+     * {@link #acquire(int)} method call should be ensured to be completed when calling this method.
      */
-    default void notifyStatusChange(S status) {}
+    default void notifyAddingSplit(SplitT split) {}
 }

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

@@ -18,23 +18,26 @@
 package org.apache.flink.api.connector.source.util.ratelimit;
 
 import org.apache.flink.annotation.Experimental;
+import org.apache.flink.api.connector.source.SourceSplit;
 
 import java.io.Serializable;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 
 /**
  * A factory for {@link RateLimiter RateLimiters} which apply rate-limiting to a source sub-task.
+ *
+ * @param <SplitT> The type of the source splits.
  */
 @Experimental
-public interface RateLimiterStrategy<S> extends Serializable {
+public interface RateLimiterStrategy<SplitT extends SourceSplit> extends Serializable {
 
     /**
      * Creates a {@link RateLimiter} that limits the rate of records going through. When there is
      * parallelism, the limiting rate is evenly reduced per subtask, such that all the sub-tasks
      * limiting rates equals the cumulative limiting rate.
      */
-    RateLimiter<S> createRateLimiter(int parallelism);
+    RateLimiter<SplitT> createRateLimiter(int parallelism);
 
     /**
      * Creates a {@code RateLimiterStrategy} that is limiting the number of records per second.

flink-tests/src/test/java/org/apache/flink/api/connector/source/lib/util/RateLimitedSourceReaderITCase.java

@@ -20,6 +20,7 @@
 
 import org.apache.flink.api.common.eventtime.WatermarkStrategy;
 import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.api.connector.source.lib.NumberSequenceSource;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiter;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
 import org.apache.flink.connector.datagen.source.DataGeneratorSource;
@@ -87,7 +88,8 @@ private List<Long> range(int startInclusive, int endInclusive) {
                 .collect(Collectors.toList());
     }
 
-    private static final class MockRateLimiter implements RateLimiter {
+    private static final class MockRateLimiter
+            implements RateLimiter<NumberSequenceSource.NumberSequenceSplit> {
 
         int callCount;
 
@@ -102,13 +104,15 @@ public int getCallCount() {
         }
     }
 
-    private static class MockRateLimiterStrategy implements RateLimiterStrategy {
+    private static class MockRateLimiterStrategy
+            implements RateLimiterStrategy<NumberSequenceSource.NumberSequenceSplit> {
 
         private static final List<MockRateLimiter> rateLimiters =
                 Collections.synchronizedList(new ArrayList<>());
 
         @Override
-        public RateLimiter createRateLimiter(int parallelism) {
+        public RateLimiter<NumberSequenceSource.NumberSequenceSplit> createRateLimiter(
+                int parallelism) {
             MockRateLimiter mockRateLimiter = new MockRateLimiter();
             rateLimiters.add(mockRateLimiter);
             return mockRateLimiter;

flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/SinkV2ITCase.java

@@ -25,6 +25,7 @@
 import org.apache.flink.api.common.typeinfo.IntegerTypeInfo;
 import org.apache.flink.api.connector.sink2.Committer;
 import org.apache.flink.api.connector.source.Source;
+import org.apache.flink.api.connector.source.lib.NumberSequenceSource;
 import org.apache.flink.api.connector.source.util.ratelimit.GatedRateLimiter;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiter;
 import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
@@ -365,7 +366,7 @@ private void assertUnalignedCheckpointInNonSink(StreamGraph streamGraph) {
      * for another two checkpoints and 5) exiting.
      */
     private Source<Integer, ?, ?> createStreamingSource() {
-        RateLimiterStrategy<Void> rateLimiterStrategy =
+        RateLimiterStrategy<NumberSequenceSource.NumberSequenceSplit> rateLimiterStrategy =
                 parallelism -> new BurstingRateLimiter(SOURCE_DATA.size() / 4, 2);
         return new DataGeneratorSource<>(
                 l -> SOURCE_DATA.get(l.intValue() % SOURCE_DATA.size()),
@@ -374,8 +375,9 @@ private void assertUnalignedCheckpointInNonSink(StreamGraph streamGraph) {
                 IntegerTypeInfo.INT_TYPE_INFO);
     }
 
-    private static class BurstingRateLimiter implements RateLimiter<Void> {
-        private final RateLimiter<Void> rateLimiter;
+    private static class BurstingRateLimiter
+            implements RateLimiter<NumberSequenceSource.NumberSequenceSplit> {
+        private final RateLimiter<NumberSequenceSource.NumberSequenceSplit> rateLimiter;
         private final int numCheckpointCooldown;
         private int cooldown;