[FLINK-40326][core] Do not over-grant capacity in GatedRateLimiter

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

@@ -28,15 +28,19 @@
 
 /**
  * An implementation of {@link RateLimiter} that completes defined number of futures in-between the
- * external notification events. The first cycle completes immediately, without waiting for the
- * external notifications.
+ * external notification events. The first cycle does not wait for an external notification: its
+ * capacity is available from the start, so requests complete immediately for as long as the
+ * capacity left covers them.
  */
 @Internal
 public class GatedRateLimiter<Split extends SourceSplit> implements RateLimiter<Split> {
 
     private final int capacityPerCycle;
     private int capacityLeft;
 
+    /** Completed while the current cycle has capacity left, incomplete once it has run out. */
+    private CompletableFuture<Void> gatingFuture = CompletableFuture.completedFuture(null);
+
     /**
      * Instantiates a new GatedRateLimiter.
      *
@@ -48,14 +52,10 @@ public GatedRateLimiter(int capacityPerCycle) {
         this.capacityLeft = capacityPerCycle;
     }
 
-    transient CompletableFuture<Void> gatingFuture = null;
-
     @Override
     public CompletionStage<Void> acquire(int numberOfEvents) {
-        if (gatingFuture == null) {
-            gatingFuture = CompletableFuture.completedFuture(null);
-        }
-        if (capacityLeft <= 0) {
+        checkArgument(numberOfEvents > 0, "Number of events has to be a positive number.");
+        if (capacityLeft < numberOfEvents) {
             gatingFuture = new CompletableFuture<>();
         }
         return gatingFuture.thenRun(() -> capacityLeft -= numberOfEvents);

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

@@ -48,7 +48,7 @@ default CompletionStage<Void> acquire() {
      * correct functioning, the next invocation of this method should only happen after the
      * previously returned future has been completed.
      *
-     * @param numberOfEvents The number of events.
+     * @param numberOfEvents The number of events, which has to be a positive number.
      */
     CompletionStage<Void> acquire(int numberOfEvents);
 

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

@@ -18,21 +18,24 @@
 
 package org.apache.flink.api.connector.source.lib.util;
 
+import org.apache.flink.api.connector.source.lib.NumberSequenceSource.NumberSequenceSplit;
 import org.apache.flink.api.connector.source.util.ratelimit.GatedRateLimiter;
 
 import org.junit.jupiter.api.Test;
 
 import java.util.concurrent.CompletionStage;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 class GatedRateLimiterTest {
 
     @Test
     void testCapacityNotExceededOnCheckpoint() {
         int capacityPerCycle = 5;
 
-        final GatedRateLimiter gatedRateLimiter = new GatedRateLimiter(capacityPerCycle);
+        final GatedRateLimiter<NumberSequenceSplit> gatedRateLimiter =
+                new GatedRateLimiter<>(capacityPerCycle);
         for (int x = 0; x < capacityPerCycle; x++) {
             assertThat(gatedRateLimiter.acquire()).isCompleted();
         }
@@ -50,4 +53,75 @@ void testCapacityNotExceededOnCheckpoint() {
         CompletionStage<Void> postCheckpoint = gatedRateLimiter.acquire();
         assertThat(postCheckpoint).isNotCompleted();
     }
+
+    @Test
+    void testCapacityNotExceededWhenAcquiringMultipleEvents() {
+        int capacityPerCycle = 5;
+
+        final GatedRateLimiter<NumberSequenceSplit> gatedRateLimiter =
+                new GatedRateLimiter<>(capacityPerCycle);
+        assertThat(gatedRateLimiter.acquire(3)).isCompleted();
+
+        // Only two permits are left in this cycle, so a request for three events has to wait even
+        // though the remaining capacity is still greater than zero.
+        CompletionStage<Void> exceedsRemainingCapacity = gatedRateLimiter.acquire(3);
+        assertThat(exceedsRemainingCapacity).isNotCompleted();
+
+        gatedRateLimiter.notifyCheckpointComplete(0);
+
+        assertThat(exceedsRemainingCapacity).isCompleted();
+    }
+
+    @Test
+    void testRequestLargerThanCapacityIsReleasedByNextCycle() {
+        final GatedRateLimiter<NumberSequenceSplit> gatedRateLimiter = new GatedRateLimiter<>(2);
+
+        // A single request may legitimately exceed the capacity of an entire cycle. It must not
+        // deadlock: resetting the capacity on the next completed checkpoint releases it.
+        CompletionStage<Void> exceedsWholeCycle = gatedRateLimiter.acquire(3);
+        assertThat(exceedsWholeCycle).isNotCompleted();
+
+        gatedRateLimiter.notifyCheckpointComplete(0);
+
+        assertThat(exceedsWholeCycle).isCompleted();
+
+        // Because completing it took 3 events from a cycle that only had 2, a further checkpoint is
+        // needed before requests are allowed again.
+        CompletionStage<Void> followingRequest = gatedRateLimiter.acquire(1);
+        assertThat(followingRequest).isNotCompleted();
+
+        gatedRateLimiter.notifyCheckpointComplete(1);
+
+        assertThat(followingRequest).isCompleted();
+        assertThat(gatedRateLimiter.acquire(1)).isCompleted();
+    }
+
+    @Test
+    void testCheckpointCompleteBeforeFirstAcquire() {
+        int capacityPerCycle = 5;
+
+        final GatedRateLimiter<NumberSequenceSplit> gatedRateLimiter =
+                new GatedRateLimiter<>(capacityPerCycle);
+
+        // A checkpoint can complete before the reader has emitted anything, for instance while it
+        // is still waiting for its first split assignment.
+        gatedRateLimiter.notifyCheckpointComplete(0);
+
+        for (int x = 0; x < capacityPerCycle; x++) {
+            assertThat(gatedRateLimiter.acquire()).isCompleted();
+        }
+        assertThat(gatedRateLimiter.acquire()).isNotCompleted();
+    }
+
+    @Test
+    void testNonPositiveNumberOfEventsIsRejected() {
+        final GatedRateLimiter<NumberSequenceSplit> gatedRateLimiter = new GatedRateLimiter<>(5);
+
+        assertThatThrownBy(() -> gatedRateLimiter.acquire(0))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("positive");
+        assertThatThrownBy(() -> gatedRateLimiter.acquire(-1))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("positive");
+    }
 }