diff --git a/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/GatedRateLimiter.java b/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/GatedRateLimiter.java index 2fa991af0138ba..75fb0b9303a858 100644 --- a/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/GatedRateLimiter.java +++ b/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/GatedRateLimiter.java @@ -28,8 +28,9 @@ /** * 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 implements RateLimiter { @@ -37,6 +38,9 @@ public class GatedRateLimiter implements RateLimiter< private final int capacityPerCycle; private int capacityLeft; + /** Completed while the current cycle has capacity left, incomplete once it has run out. */ + private CompletableFuture gatingFuture = CompletableFuture.completedFuture(null); + /** * Instantiates a new GatedRateLimiter. * @@ -48,14 +52,10 @@ public GatedRateLimiter(int capacityPerCycle) { this.capacityLeft = capacityPerCycle; } - transient CompletableFuture gatingFuture = null; - @Override public CompletionStage 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); diff --git a/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/RateLimiter.java b/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/RateLimiter.java index 06de7bc7269db8..bc2a14ab191ebf 100644 --- a/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/RateLimiter.java +++ b/flink-core/src/main/java/org/apache/flink/api/connector/source/util/ratelimit/RateLimiter.java @@ -48,7 +48,7 @@ default CompletionStage 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 acquire(int numberOfEvents); diff --git a/flink-tests/src/test/java/org/apache/flink/api/connector/source/lib/util/GatedRateLimiterTest.java b/flink-tests/src/test/java/org/apache/flink/api/connector/source/lib/util/GatedRateLimiterTest.java index 15a034e0f7d5b1..21c905a63747ee 100644 --- a/flink-tests/src/test/java/org/apache/flink/api/connector/source/lib/util/GatedRateLimiterTest.java +++ b/flink-tests/src/test/java/org/apache/flink/api/connector/source/lib/util/GatedRateLimiterTest.java @@ -18,6 +18,7 @@ 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; @@ -25,6 +26,7 @@ import java.util.concurrent.CompletionStage; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; class GatedRateLimiterTest { @@ -32,7 +34,8 @@ class GatedRateLimiterTest { void testCapacityNotExceededOnCheckpoint() { int capacityPerCycle = 5; - final GatedRateLimiter gatedRateLimiter = new GatedRateLimiter(capacityPerCycle); + final GatedRateLimiter gatedRateLimiter = + new GatedRateLimiter<>(capacityPerCycle); for (int x = 0; x < capacityPerCycle; x++) { assertThat(gatedRateLimiter.acquire()).isCompleted(); } @@ -50,4 +53,75 @@ void testCapacityNotExceededOnCheckpoint() { CompletionStage postCheckpoint = gatedRateLimiter.acquire(); assertThat(postCheckpoint).isNotCompleted(); } + + @Test + void testCapacityNotExceededWhenAcquiringMultipleEvents() { + int capacityPerCycle = 5; + + final GatedRateLimiter 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 exceedsRemainingCapacity = gatedRateLimiter.acquire(3); + assertThat(exceedsRemainingCapacity).isNotCompleted(); + + gatedRateLimiter.notifyCheckpointComplete(0); + + assertThat(exceedsRemainingCapacity).isCompleted(); + } + + @Test + void testRequestLargerThanCapacityIsReleasedByNextCycle() { + final GatedRateLimiter 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 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 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 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 gatedRateLimiter = new GatedRateLimiter<>(5); + + assertThatThrownBy(() -> gatedRateLimiter.acquire(0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("positive"); + assertThatThrownBy(() -> gatedRateLimiter.acquire(-1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("positive"); + } }