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..3834906f5057 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,127 @@ package software.amazon.awssdk.retries.internal.ratelimiter; import java.time.Duration; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; 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.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 ScheduledExecutorService scheduler; + // the collection of futures returned to threads currently waiting for capacity. Futures are completed in FIFO order. + private final Deque> waiting = new ArrayDeque<>(); private final AtomicReference stateReference; private final RateLimiterClock clock; + private final Object lock = new Object(); - RateLimiterTokenBucket(RateLimiterClock clock) { + private boolean notifierRunning = false; + + RateLimiterTokenBucket(RateLimiterClock clock, ScheduledExecutorService scheduler) { this.clock = clock; + this.scheduler = scheduler; this.stateReference = new AtomicReference<>(new PersistentState()); } + @Override + public void close() { + scheduler.shutdownNow(); + } + /** - * 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) { + CompletableFuture future = new CompletableFuture<>(); + if (!notifierRunning) { + waiting.add(future); + boolean scheduled = scheduleOrFail(this::doNotify, Duration.ZERO, future); + if (scheduled) { + notifierRunning = true; + } else { + waiting.pop(); + } + } + return future; + } + } + + public void doNotify() { + synchronized (lock) { + while (true) { + CompletableFuture w = waiting.poll(); + if (w == null) { + notifierRunning = false; + return; + } + + PersistentState persistentState = stateReference.get(); + TransientState ts = persistentState.toTransient(); + TransientState.AcquireResult acquireResult = ts.tokenBucketAcquire(clock, 1.0); + stateReference.set(ts.toPersistent()); + + // Not enough capacity. Try again later when enough time has passed to refill the bucket at the current rate. + if (!acquireResult.isSuccessful()) { + waiting.push(w); + + if (!scheduleOrFail(this::doNotify, acquireResult.refillWait(), w)) { + waiting.pop(); + return; + } + } + + // Acquire was successful, signal the waiting thread. + w.complete(null); + } + } + } + + @SdkTestInternalApi + Deque> waiting() { + return waiting; + } + + private void schedule(Runnable command, Duration d) { + scheduler.schedule(command, d.toMillis(), TimeUnit.MILLISECONDS); + } + + /** + * @return true if schedule was successful, false otherwise. + */ + private boolean scheduleOrFail(Runnable command, Duration d, CompletableFuture future) { + try { + schedule(command, d); + return true; + } catch (Throwable t) { + RuntimeException e = new RuntimeException("Unable to initiate token acquire", t); + future.completeExceptionally(e); + } + return false; } /** @@ -95,17 +180,20 @@ private StateUpdate consumeState(Consumer mutator) { * retried until succeeded. */ private StateUpdate updateState(Function mutator) { - PersistentState current; - PersistentState updated; - T result; - do { - current = stateReference.get(); - TransientState transientState = current.toTransient(); - result = mutator.apply(transientState); - updated = transientState.toPersistent(); - } while (!stateReference.compareAndSet(current, updated)); - - return new StateUpdate<>(updated, result); + synchronized (lock) { + PersistentState current; + PersistentState updated; + T result; + do { + current = stateReference.get(); + TransientState transientState = current.toTransient(); + result = mutator.apply(transientState); + updated = transientState.toPersistent(); + } while (!stateReference.compareAndSet(current, updated)); + + + return new StateUpdate<>(updated, result); + } } static class StateUpdate { @@ -163,17 +251,39 @@ 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; + this.currentCapacity = 0; + 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..680c3b51f31e 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.ScheduledExecutorService; 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 ScheduledExecutorService scheduler; private RateLimiterTokenBucketStore(Builder builder) { this.clock = Validate.paramNotNull(builder.clock, "clock"); - this.scopeToTokenBucket = LruCache.builder(x -> new RateLimiterTokenBucket(clock)) + this.scheduler = Validate.paramNotNull(builder.scheduler, "scheduler"); + this.scopeToTokenBucket = LruCache.builder( + x -> new RateLimiterTokenBucket(clock, scheduler)) .maxSize(MAX_ENTRIES) .build(); } + @Override + public void close() { + scheduler.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 ScheduledExecutorService scheduler; Builder() { this.clock = DEFAULT_CLOCK; @@ -70,6 +81,11 @@ public Builder clock(RateLimiterClock clock) { return this; } + public Builder scheduler(ScheduledExecutorService scheduler) { + this.scheduler = scheduler; + 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..9bea5a4dc787 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,132 @@ 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.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; import java.util.Arrays; import java.util.Collection; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.MethodSource; +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 static ScheduledExecutorService scheduler = null; + private static RateLimiterTokenBucket tokenBucket = null; @BeforeAll static 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_notifierNotRunning_schedules() { + ScheduledExecutorService testScheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucket testBucket = new RateLimiterTokenBucket(new MutableClock(), testScheduler); + + testBucket.acquireAsync(); + verify(testScheduler).schedule(any(Runnable.class), eq(0L), any(TimeUnit.class)); + } + + @Test + void acquireAsync_notifierRunning_doesNotSchedule() { + ScheduledExecutorService testScheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucket testBucket = new RateLimiterTokenBucket(new MutableClock(), testScheduler); + testBucket.acquireAsync(); + verify(testScheduler).schedule(any(Runnable.class), eq(0L), any(TimeUnit.class)); + + testBucket.acquireAsync(); + verifyNoMoreInteractions(testScheduler); + } + + @Test + void acquireAsync_notifierNotRunning_schedulerThrows_completesFuture() { + ScheduledExecutorService testScheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucket testBucket = new RateLimiterTokenBucket(new MutableClock(), testScheduler); + when(testScheduler.schedule(any(Runnable.class), eq(0L), any(TimeUnit.class))) + .thenThrow(new RejectedExecutionException("rejected")); + + assertThatThrownBy(testBucket.acquireAsync()::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).isExactlyInstanceOf(RuntimeException.class); + assertThat(cause).hasMessage("Unable to initiate token acquire"); + assertThat(cause.getCause()).hasMessage("rejected"); + }); + + assertThat(testBucket.waiting()).isEmpty(); + } + + @Test + void acquireAsync_notifierRunning_schedulerThrows_completesFuture() { + ScheduledExecutorService testScheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucket testBucket = new RateLimiterTokenBucket(new MutableClock(), testScheduler); + + // Enable the bucket + testBucket.updateRateAfterThrottling(); + + when(testScheduler.schedule(any(Runnable.class), anyLong(), any())) + .thenAnswer(i -> { + Runnable command = i.getArgument(0); + Long delay = i.getArgument(1); + if (delay == 0L) { + command.run(); + return null; + } + throw new RejectedExecutionException("rejected"); + }); + + assertThatThrownBy(testBucket.acquireAsync()::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).isExactlyInstanceOf(RuntimeException.class); + assertThat(cause).hasMessage("Unable to initiate token acquire"); + assertThat(cause.getCause()).hasMessage("rejected"); + }); + + assertThat(testBucket.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 +262,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 {