From 236b1d6cd70965ef123a9e708d248d5b57133def Mon Sep 17 00:00:00 2001 From: Dongie Agnir Date: Tue, 14 Jul 2026 13:37:58 -0700 Subject: [PATCH] Add acquireAsync This commit replaces tryAcquire with acquireAsync to better support async codepaths. acquireAsync returns a CompletableFuture that is completed when a token is successfully acquired; unlike tryAcquire, the caller is not expected to do any of their own waiting once future is complete. --- core/retries/pom.xml | 5 + .../DefaultAdaptiveRetryStrategy.java | 7 +- .../ratelimiter/RateLimiterTokenBucket.java | 145 +++++++++++++++--- .../RateLimiterTokenBucketStore.java | 20 ++- .../RateLimiterTokenBucketStoreTest.java | 37 +++++ .../RateLimiterTokenBucketTest.java | 89 +++++++++-- 6 files changed, 255 insertions(+), 48 deletions(-) create mode 100644 core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java diff --git a/core/retries/pom.xml b/core/retries/pom.xml index aba7d96cd5d8..623bd4a00e6a 100644 --- a/core/retries/pom.xml +++ b/core/retries/pom.xml @@ -70,5 +70,10 @@ assertj-core test + + org.mockito + mockito-core + test + diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java index c52c7859526f..e63e683a18ab 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java @@ -43,15 +43,12 @@ public final class DefaultAdaptiveRetryStrategy @Override protected Duration computeInitialBackoff(AcquireInitialTokenRequest request) { - RateLimiterTokenBucket bucket = rateLimiterTokenBucketStore.tokenBucketForScope(request.scope()); - return bucket.tryAcquire().delay(); + throw new UnsupportedOperationException("TODO"); } @Override protected Duration computeBackoff(RefreshRetryTokenRequest request, DefaultRetryToken token) { - Duration backoff = super.computeBackoff(request, token); - RateLimiterTokenBucket bucket = rateLimiterTokenBucketStore.tokenBucketForScope(token.scope()); - return backoff.plus(bucket.tryAcquire().delay()); + throw new UnsupportedOperationException("TODO"); } @Override diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java index 2d8743af9924..19eb196b7603 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java @@ -16,42 +16,114 @@ package software.amazon.awssdk.retries.internal.ratelimiter; import java.time.Duration; +import java.util.Deque; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.LinkedBlockingDeque; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.annotations.SdkTestInternalApi; +import software.amazon.awssdk.annotations.ThreadSafe; +import software.amazon.awssdk.utils.CompletableFutureUtils; +import software.amazon.awssdk.utils.SdkAutoCloseable; /** * The {@link RateLimiterTokenBucket} keeps track of past throttling responses and adapts to slow down the send rate to adapt to - * the service. It does this by suggesting a delay amount as result of a {@link #tryAcquire()} call. Callers must update its - * internal state by calling {@link #updateRateAfterThrottling()} when getting a throttling response or - * {@link #updateRateAfterSuccess()} when getting successful response. + * the service. It does this by delaying the completion of the future returned by {@link #acquireAsync()} until the requested + * capacity is available. Callers must update its internal state by calling {@link #updateRateAfterThrottling()} when getting a + * throttling response or {@link #updateRateAfterSuccess()} when getting successful response. * - *

This class is thread-safe, its internal current state is kept in the inner class {@link PersistentState} which is stored - * using an {@link AtomicReference}. This class is converted to {@link TransientState} when the state needs to be mutated and - * converted back to a {@link PersistentState} and stored using {@link AtomicReference#compareAndSet(Object, Object)}. + *

This class is thread-safe. * *

The algorithm used is adapted from the network congestion avoidance algorithm * CUBIC. */ @SdkInternalApi -public class RateLimiterTokenBucket { +@ThreadSafe +public class RateLimiterTokenBucket implements SdkAutoCloseable { + // Thread used for capacity waiting and notifying. + private final ExecutorService exec; + // the collection of futures returned to threads currently waiting for capacity. Futures are completed in FIFO order. + private final BlockingDeque> waiting = new LinkedBlockingDeque<>(); private final AtomicReference stateReference; private final RateLimiterClock clock; + private final Object lock = new Object(); + private boolean open = true; - RateLimiterTokenBucket(RateLimiterClock clock) { + RateLimiterTokenBucket(RateLimiterClock clock, ExecutorService exec) { this.clock = clock; + this.exec = exec; this.stateReference = new AtomicReference<>(new PersistentState()); + exec.submit(this::doNotify); + } + + @Override + public void close() { + synchronized (lock) { + open = false; + IllegalStateException closedException = new IllegalStateException("Rate limiter bucket is closed"); + while (true) { + CompletableFuture w = waiting.poll(); + if (w == null) { + break; + } + w.completeExceptionally(closedException); + } + } } /** - * Acquire tokens from the bucket. If the bucket contains enough capacity to satisfy the request, this method will return in - * {@link RateLimiterAcquireResponse#delay()} a {@link Duration#ZERO} value, otherwise it will return the amount of time the - * callers need to wait until enough tokens are refilled. + * Acquire a token from the bucket. + * + * @return A future that is completed when the requested amount is acquired from this bucket. */ - public RateLimiterAcquireResponse tryAcquire() { - StateUpdate update = updateState(ts -> ts.tokenBucketAcquire(clock, 1.0)); - return RateLimiterAcquireResponse.create(update.result); + public CompletableFuture acquireAsync() { + synchronized (lock) { + if (!open) { + return CompletableFutureUtils.failedFuture(new IllegalStateException("Rate limiter bucket is closed")); + } + CompletableFuture future = new CompletableFuture<>(); + waiting.add(future); + return future; + } + } + + public void doNotify() { + while (true) { + CompletableFuture w = null; + try { + w = waiting.take(); + + // polling acquire loop. Attempt to acquire capacity from the + // bucket until we're successful or interrupted. + TransientState.AcquireResult acquireResult; + do { + acquireResult = updateState(t -> t.tokenBucketAcquire(clock, 1.0)).result; + + // Not enough capacity. Try again later when enough time + // has passed to refill the bucket at the current rate. + if (!acquireResult.isSuccessful()) { + Thread.sleep(acquireResult.refillWait().toMillis()); + } else { + // Acquire was successful, signal the waiting thread. + w.complete(null); + } + } while (!acquireResult.isSuccessful()); + } catch (InterruptedException ie) { + if (w != null) { + w.completeExceptionally(new RuntimeException("Unable to acquire capacity", ie)); + } + break; + } + } + } + + @SdkTestInternalApi + Deque> waiting() { + return waiting; } /** @@ -95,17 +167,19 @@ private StateUpdate consumeState(Consumer mutator) { * retried until succeeded. */ private StateUpdate updateState(Function mutator) { - PersistentState current; - PersistentState updated; - T result; - do { + synchronized (lock) { + PersistentState current; + PersistentState updated; + T result; + current = stateReference.get(); TransientState transientState = current.toTransient(); result = mutator.apply(transientState); updated = transientState.toPersistent(); - } while (!stateReference.compareAndSet(current, updated)); + stateReference.set(updated); - return new StateUpdate<>(updated, result); + return new StateUpdate<>(updated, result); + } } static class StateUpdate { @@ -163,17 +237,38 @@ PersistentState toPersistent() { * a {@link Duration#ZERO} value, otherwise it will return the amount of time the callers need to wait until enough tokens * are refilled. */ - Duration tokenBucketAcquire(RateLimiterClock clock, double amount) { + AcquireResult tokenBucketAcquire(RateLimiterClock clock, double amount) { if (!this.enabled) { - return Duration.ZERO; + return new AcquireResult(true, Duration.ZERO); } refill(clock); - double waitTime = 0.0; if (this.currentCapacity < amount) { - waitTime = (amount - this.currentCapacity) / this.fillRate; + double diff = amount - currentCapacity; + double waitTime = diff / this.fillRate; + double waitTimeMs = waitTime * 1_000.0; + Duration duration = Duration.ofMillis((long) Math.ceil(waitTimeMs)); + return new AcquireResult(false, duration); } this.currentCapacity -= amount; - return Duration.ofNanos((long) (waitTime * 1_000_000_000.0)); + return new AcquireResult(true, Duration.ZERO); + } + + private static class AcquireResult { + final boolean successful; + final Duration refillWait; + + public AcquireResult(boolean successful, Duration refillWait) { + this.successful = successful; + this.refillWait = refillWait; + } + + public boolean isSuccessful() { + return successful; + } + + public Duration refillWait() { + return refillWait; + } } /** diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java index 303e6d61d7f5..f96035043e1c 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java @@ -15,8 +15,10 @@ package software.amazon.awssdk.retries.internal.ratelimiter; +import java.util.concurrent.ExecutorService; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.annotations.ToBuilderIgnoreField; +import software.amazon.awssdk.utils.SdkAutoCloseable; import software.amazon.awssdk.utils.Validate; import software.amazon.awssdk.utils.builder.CopyableBuilder; import software.amazon.awssdk.utils.builder.ToCopyableBuilder; @@ -27,19 +29,27 @@ */ @SdkInternalApi public final class RateLimiterTokenBucketStore - implements ToCopyableBuilder { + implements ToCopyableBuilder, SdkAutoCloseable { private static final int MAX_ENTRIES = 128; private static final RateLimiterClock DEFAULT_CLOCK = new SystemClock(); private final LruCache scopeToTokenBucket; private final RateLimiterClock clock; + private final ExecutorService executor; private RateLimiterTokenBucketStore(Builder builder) { this.clock = Validate.paramNotNull(builder.clock, "clock"); - this.scopeToTokenBucket = LruCache.builder(x -> new RateLimiterTokenBucket(clock)) + this.executor = Validate.paramNotNull(builder.executor, "executor"); + this.scopeToTokenBucket = LruCache.builder( + x -> new RateLimiterTokenBucket(clock, executor)) .maxSize(MAX_ENTRIES) .build(); } + @Override + public void close() { + executor.shutdownNow(); + } + public RateLimiterTokenBucket tokenBucketForScope(String scope) { return scopeToTokenBucket.get(scope); } @@ -56,6 +66,7 @@ public static RateLimiterTokenBucketStore.Builder builder() { public static class Builder implements CopyableBuilder { private RateLimiterClock clock; + private ExecutorService executor; Builder() { this.clock = DEFAULT_CLOCK; @@ -70,6 +81,11 @@ public Builder clock(RateLimiterClock clock) { return this; } + public Builder executor(ExecutorService executor) { + this.executor = executor; + return this; + } + @Override public RateLimiterTokenBucketStore build() { return new RateLimiterTokenBucketStore(this); diff --git a/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java new file mode 100644 index 000000000000..5c6ad377d3ee --- /dev/null +++ b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java @@ -0,0 +1,37 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file 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 software.amazon.awssdk.retries.internal.ratelimiter; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.concurrent.ScheduledExecutorService; +import org.junit.jupiter.api.Test; + +public class RateLimiterTokenBucketStoreTest { + @Test + void close_closesScheduler() { + ScheduledExecutorService mockScheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder() + .clock(new SystemClock()) + .scheduler(mockScheduler) + .build(); + + store.close(); + + verify(mockScheduler).shutdownNow(); + } +} diff --git a/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java index 14021ebde812..28b689061a11 100644 --- a/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java +++ b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java @@ -16,44 +16,92 @@ package software.amazon.awssdk.retries.internal.ratelimiter; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.AssertionsForClassTypes.within; +import static org.mockito.Mockito.mock; import java.util.Arrays; import java.util.Collection; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.MethodSource; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; class RateLimiterTokenBucketTest { - private static MutableClock clock = null; - private static RateLimiterTokenBucket tokenBucket = null; private static final double EPSILON = 0.0001; + private MutableClock clock = null; + private ScheduledExecutorService scheduler = null; + private RateLimiterTokenBucket tokenBucket = null; - @BeforeAll - static void setup() { + @BeforeEach + void setup() { clock = new MutableClock(); - tokenBucket = new RateLimiterTokenBucket(clock); + scheduler = mock(ScheduledExecutorService.class); + tokenBucket = new RateLimiterTokenBucket(clock, scheduler); } - @ParameterizedTest - @MethodSource("parameters") - void testCase(TestCase testCase) { + @Test + void acquireAsync_bucketClosed_futureCompletedExceptionally() { + tokenBucket.close(); + CompletableFuture f = tokenBucket.acquireAsync(); + assertThatThrownBy(f::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).isExactlyInstanceOf(IllegalStateException.class); + assertThat(cause).hasMessage("Rate limiter bucket is closed"); + }); + } + + @Test + void close_completesAllPendingFutures() { + List> futures = IntStream.range(0, 10) + .mapToObj(i -> tokenBucket.acquireAsync()) + .collect(Collectors.toList()); + + tokenBucket.close(); + + assertThat(futures).allSatisfy(f -> { + assertThatThrownBy(f::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).isExactlyInstanceOf(IllegalStateException.class); + assertThat(cause).hasMessage("Rate limiter bucket is closed"); + }); + }); + + assertThat(tokenBucket.waiting()).isEmpty(); + } + + @Test + void sendingRateEndToEndTest() { + for (TestCase sendingRateTestCase : sendingRateTestCases()) { + assertSendingRateTestCase(sendingRateTestCase); + } + } + + void assertSendingRateTestCase(TestCase testCase) { clock.setCurrent(testCase.givenTimestamp); RateLimiterUpdateResponse res; - tokenBucket.tryAcquire(); + if (testCase.throttleResponse) { res = tokenBucket.updateRateAfterThrottling(); } else { res = tokenBucket.updateRateAfterSuccess(); } double measuredTxRate = res.measuredTxRate(); - assertThat(measuredTxRate).isCloseTo(testCase.expectMeasuredTxRate, within(EPSILON)); + assertThat(measuredTxRate) + .as("%s: Measured TX rate", testCase) + .isCloseTo(testCase.expectMeasuredTxRate, within(EPSILON)); double fillRate = res.fillRate(); - assertThat(fillRate).isCloseTo(testCase.expectFillRate, within(EPSILON)); + assertThat(fillRate) + .as("%s: Fill rate", testCase) + .isCloseTo(testCase.expectFillRate, within(EPSILON)); } - - static Collection parameters() { + static Collection sendingRateTestCases() { + // Note: Test cases are not independent. Each case depends on the state of the bucket being correctly updated from the + // previous test. return Arrays.asList( new TestCase() .givenSuccessResponse() @@ -174,6 +222,15 @@ TestCase expectFillRate(double expectFillRate) { return this; } + @Override + public String toString() { + return "TestCase{" + + "throttleResponse=" + throttleResponse + + ", givenTimestamp=" + givenTimestamp + + ", expectMeasuredTxRate=" + expectMeasuredTxRate + + ", expectFillRate=" + expectFillRate + + '}'; + } } static class MutableClock implements RateLimiterClock {