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 {