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.

Author: dagnir

core/retries/pom.xml

@@ -70,5 +70,10 @@
             <artifactId>assertj-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

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

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.
  *
- * <p>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)}.
+ * <p>This class is thread-safe.
  *
  * <p>The algorithm used is adapted from the network congestion avoidance algorithm
  * <a href="https://en.wikipedia.org/wiki/CUBIC_TCP">CUBIC</a>.
  */
 @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<CompletableFuture<Void>> waiting = new ArrayDeque<>();
     private final AtomicReference<PersistentState> 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<Duration> update = updateState(ts -> ts.tokenBucketAcquire(clock, 1.0));
-        return RateLimiterAcquireResponse.create(update.result);
+    public CompletableFuture<Void> acquireAsync() {
+        synchronized (lock) {
+            CompletableFuture<Void> 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<Void> 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<CompletableFuture<Void>> 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<Void> consumeState(Consumer<TransientState> mutator) {
      * retried until succeeded.
      */
     private <T> StateUpdate<T> updateState(Function<TransientState, T> 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<T> {
@@ -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;
+            }
         }
 
         /**

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<RateLimiterTokenBucketStore.Builder, RateLimiterTokenBucketStore> {
+    implements ToCopyableBuilder<RateLimiterTokenBucketStore.Builder, RateLimiterTokenBucketStore>, SdkAutoCloseable {
     private static final int MAX_ENTRIES = 128;
     private static final RateLimiterClock DEFAULT_CLOCK = new SystemClock();
     private final LruCache<String, RateLimiterTokenBucket> scopeToTokenBucket;
     private final RateLimiterClock clock;
+    private final ScheduledExecutorService scheduler;
 
     private RateLimiterTokenBucketStore(Builder builder) {
         this.clock = Validate.paramNotNull(builder.clock, "clock");
-        this.scopeToTokenBucket = LruCache.<String, RateLimiterTokenBucket>builder(x -> new RateLimiterTokenBucket(clock))
+        this.scheduler = Validate.paramNotNull(builder.scheduler, "scheduler");
+        this.scopeToTokenBucket = LruCache.<String, RateLimiterTokenBucket>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<Builder, RateLimiterTokenBucketStore> {
         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);

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();
+    }
+}