diff --git a/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml b/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml index 05606dc6d574..33a8b1c77f40 100644 --- a/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml +++ b/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml @@ -547,4 +547,10 @@ + + + + + + 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..31d1eb2574ac 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,128 @@ 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; + private boolean notifierThreadRunning = false; - RateLimiterTokenBucket(RateLimiterClock clock) { + RateLimiterTokenBucket(RateLimiterClock clock, ExecutorService exec) { this.clock = clock; + this.exec = exec; this.stateReference = new AtomicReference<>(new PersistentState()); } + @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")); + } + + if (!stateReference.get().enabled) { + return CompletableFuture.completedFuture(null); + } + + if (!notifierThreadRunning) { + try { + exec.execute(this::doNotify); + notifierThreadRunning = true; + } catch (Throwable t) { + return CompletableFutureUtils.failedFuture(new RuntimeException("Unable to initiate token acquire", t)); + } + } + + 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 +181,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 +251,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; + + 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..633bc848aaa3 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; @@ -63,6 +74,7 @@ public static class Builder implements CopyableBuilder 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 acquireAsync_notEnabled_doesNotSubmitTask() { + CompletableFuture f = tokenBucket.acquireAsync(); + assertThat(f).isCompleted(); + verifyNoInteractions(executor); + } + + @Test + void acquireAsync_enabled_submitsTask() { + // a throttle is the only way to enable rate limiting + tokenBucket.updateRateAfterThrottling(); + + tokenBucket.acquireAsync(); + verify(executor).execute(any(Runnable.class)); + } + + + @Test + void acquireAsync_submitFails_completesFutureExceptionally() { + tokenBucket.updateRateAfterThrottling(); + + doThrow(new RejectedExecutionException("no")).when(executor).execute(any(Runnable.class)); + + CompletableFuture f = tokenBucket.acquireAsync(); + assertThatThrownBy(f::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).hasMessage("Unable to initiate token acquire"); + }); + } + + @Test + void close_completesAllPendingFutures() { + // enable throttling so futures actually get queued instead of being completed immediately + tokenBucket.updateRateAfterThrottling(); + + 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 +261,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 {