Wire in acquireAsync in adaptive strat

build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml

@@ -547,4 +547,14 @@
         <Method name="marshallFieldViaRegistry"/>
         <Bug pattern="AT_OPERATION_SEQUENCE_ON_CONCURRENT_ABSTRACTION"/>
     </Match>
+
+    <!-- By design. These methods should block. -->
+    <Match>
+        <Class name="software.amazon.awssdk.retries.internal.DefaultAdaptiveRetryStrategy"/>
+        <Or>
+            <Method name="acquireInitialToken"/>
+            <Method name="refreshRetryToken"/>
+        </Or>
+        <Bug pattern="ASYNC_BLOCKING_CALL"/>
+    </Match>
 </FindBugsFilter>

core/retries/src/main/java/software/amazon/awssdk/retries/internal/BaseRetryStrategy.java

@@ -39,6 +39,7 @@
 import software.amazon.awssdk.retries.internal.circuitbreaker.TokenBucket;
 import software.amazon.awssdk.retries.internal.circuitbreaker.TokenBucketStore;
 import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.Pair;
 import software.amazon.awssdk.utils.ToString;
 import software.amazon.awssdk.utils.Validate;
 
@@ -86,7 +87,7 @@ public abstract class BaseRetryStrategy implements DefaultAwareRetryStrategy {
      * @see RetryStrategy#acquireInitialToken(AcquireInitialTokenRequest)
      */
     @Override
-    public final AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) {
+    public AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) {
         logAcquireInitialToken(request);
         DefaultRetryToken token = DefaultRetryToken.builder().scope(request.scope()).build();
         return AcquireInitialTokenResponse.create(token, computeInitialBackoff(request));
@@ -98,7 +99,20 @@ public final AcquireInitialTokenResponse acquireInitialToken(AcquireInitialToken
      * @see RetryStrategy#refreshRetryToken(RefreshRetryTokenRequest)
      */
     @Override
-    public final RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) {
+    public RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) {
+        Pair<DefaultRetryToken, AcquireResponse> refreshedToken = refreshTokenOrThrow(request);
+        Duration backoff = computeBackoff(request, refreshedToken.left());
+
+        logRefreshTokenSuccess(refreshedToken.left(), refreshedToken.right(), backoff);
+        return RefreshRetryTokenResponseImpl.create(refreshedToken.left(), backoff);
+    }
+
+    /**
+     * Attempt to refresh the token for a retry or throws {@link TokenAcquisitionFailedException} if unable to do so.
+     *
+     * @return A pair of the refreshed token and the successful acquire response from the token bucket.
+     */
+    protected Pair<DefaultRetryToken, AcquireResponse> refreshTokenOrThrow(RefreshRetryTokenRequest request) {
         DefaultRetryToken token = asDefaultRetryToken(request.token());
 
         // Check if we meet the preconditions needed for retrying. These will throw if the expected condition is not meet.
@@ -115,12 +129,8 @@ public final RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenReques
         // All the conditions required to retry were meet, update the internal state before retrying.
         updateStateForRetry(request);
 
-        // Refresh the retry token and compute the backoff delay.
-        DefaultRetryToken refreshedToken = refreshToken(request, acquireResponse);
-        Duration backoff = computeBackoff(request, refreshedToken);
-
-        logRefreshTokenSuccess(refreshedToken, acquireResponse, backoff);
-        return RefreshRetryTokenResponseImpl.create(refreshedToken, backoff);
+        // Refresh the retry token
+        return Pair.of(refreshToken(request, acquireResponse), acquireResponse);
     }
 
     /**
@@ -335,15 +345,15 @@ private String acquisitionFailedMessage(AcquireResponse response) {
                              response.maxCapacity());
     }
 
-    private void logAcquireInitialToken(AcquireInitialTokenRequest request) {
+    protected void logAcquireInitialToken(AcquireInitialTokenRequest request) {
         // Request attempt 1 token acquired (backoff: 0ms, cost: 0, capacity: 500/500)
         TokenBucket tokenBucket = tokenBucketStore.tokenBucketForScope(request.scope());
         log.debug(() -> String.format("Request attempt 1 token acquired "
                                       + "(backoff: 0ms, cost: 0, capacity: %d/%d)",
                                       tokenBucket.currentCapacity(), tokenBucket.maxCapacity()));
     }
 
-    private void logRefreshTokenSuccess(DefaultRetryToken token, AcquireResponse acquireResponse, Duration delay) {
+    protected void logRefreshTokenSuccess(DefaultRetryToken token, AcquireResponse acquireResponse, Duration delay) {
         log.debug(() -> String.format("Request attempt %d token acquired "
                                       + "(backoff: %dms, cost: %d, capacity: %d/%d)",
                                       token.attempt(), delay.toMillis(),

core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java

@@ -16,22 +16,27 @@
 package software.amazon.awssdk.retries.internal;
 
 import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.retries.AdaptiveRetryStrategy;
 import software.amazon.awssdk.retries.api.AcquireInitialTokenRequest;
+import software.amazon.awssdk.retries.api.AcquireInitialTokenResponse;
 import software.amazon.awssdk.retries.api.BackoffStrategy;
 import software.amazon.awssdk.retries.api.RefreshRetryTokenRequest;
+import software.amazon.awssdk.retries.api.RefreshRetryTokenResponse;
+import software.amazon.awssdk.retries.internal.circuitbreaker.AcquireResponse;
 import software.amazon.awssdk.retries.internal.circuitbreaker.TokenBucketStore;
 import software.amazon.awssdk.retries.internal.ratelimiter.RateLimiterTokenBucket;
 import software.amazon.awssdk.retries.internal.ratelimiter.RateLimiterTokenBucketStore;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
 import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.Pair;
 import software.amazon.awssdk.utils.Validate;
 
 @SdkInternalApi
 public final class DefaultAdaptiveRetryStrategy
     extends BaseRetryStrategy implements AdaptiveRetryStrategy {
-
     private static final Logger LOG = Logger.loggerFor(DefaultAdaptiveRetryStrategy.class);
     private final RateLimiterTokenBucketStore rateLimiterTokenBucketStore;
 
@@ -42,13 +47,53 @@ public final class DefaultAdaptiveRetryStrategy
     }
 
     @Override
-    protected Duration computeInitialBackoff(AcquireInitialTokenRequest request) {
-        throw new UnsupportedOperationException("TODO");
+    public AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) {
+        return CompletableFutureUtils.joinLikeSync(acquireInitialTokenAsync(request));
+    }
+
+    @Override
+    public RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) {
+        return CompletableFutureUtils.joinLikeSync(refreshRetryTokenAsync(request));
+    }
+
+    @Override
+    public CompletableFuture<AcquireInitialTokenResponse> acquireInitialTokenAsync(AcquireInitialTokenRequest request) {
+        logAcquireInitialToken(request);
+        RateLimiterTokenBucket bucket = rateLimiterTokenBucketStore.tokenBucketForScope(request.scope());
+        CompletableFuture<Void> acquireResult = bucket.acquireAsync();
+
+        return acquireResult.thenApply(r -> {
+            DefaultRetryToken token = DefaultRetryToken.builder().scope(request.scope()).build();
+            return AcquireInitialTokenResponse.create(token, Duration.ZERO);
+        });
+    }
+
+    @Override
+    public CompletableFuture<RefreshRetryTokenResponse> refreshRetryTokenAsync(RefreshRetryTokenRequest request) {
+        DefaultRetryToken token = (DefaultRetryToken) request.token();
+        Pair<DefaultRetryToken, AcquireResponse> refreshResult;
+        try {
+            refreshResult = refreshTokenOrThrow(request);
+        } catch (Throwable t) {
+            return CompletableFutureUtils.failedFuture(t);
+        }
+
+        DefaultRetryToken refreshedToken = refreshResult.left();
+        AcquireResponse acquireResponse = refreshResult.right();
+        RateLimiterTokenBucket bucket = rateLimiterTokenBucketStore.tokenBucketForScope(token.scope());
+        CompletableFuture<Void> acquireResult = bucket.acquireAsync();
+        return acquireResult.thenApply(r -> {
+            // Note: This is the backoff imposed standard retry strategy, *not* the rate limiter. This must still be honored by
+            // the caller before sending the request.
+            Duration backoff = computeBackoff(request, refreshedToken);
+            logRefreshTokenSuccess(refreshedToken, acquireResponse, backoff);
+            return RefreshRetryTokenResponse.create(refreshedToken, backoff);
+        });
     }
 
     @Override
-    protected Duration computeBackoff(RefreshRetryTokenRequest request, DefaultRetryToken token) {
-        throw new UnsupportedOperationException("TODO");
+    public void close() {
+        rateLimiterTokenBucketStore.close();
     }
 
     @Override

core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java

@@ -15,10 +15,13 @@
 
 package software.amazon.awssdk.retries.internal.ratelimiter;
 
+import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.annotations.ToBuilderIgnoreField;
 import software.amazon.awssdk.utils.SdkAutoCloseable;
+import software.amazon.awssdk.utils.ThreadFactoryBuilder;
 import software.amazon.awssdk.utils.Validate;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
 import software.amazon.awssdk.utils.builder.ToCopyableBuilder;
@@ -31,30 +34,57 @@
 public final class RateLimiterTokenBucketStore
     implements ToCopyableBuilder<RateLimiterTokenBucketStore.Builder, RateLimiterTokenBucketStore>, SdkAutoCloseable {
     private static final int MAX_ENTRIES = 128;
+    private static final String THREAD_NAME_PREFIX = "sdk-adaptive-rate-limiter-";
+
     private static final RateLimiterClock DEFAULT_CLOCK = new SystemClock();
     private final LruCache<String, RateLimiterTokenBucket> scopeToTokenBucket;
     private final RateLimiterClock clock;
     private final ScheduledExecutorService scheduler;
+    private final boolean closeScheduler;
 
     private RateLimiterTokenBucketStore(Builder builder) {
-        this.clock = Validate.paramNotNull(builder.clock, "clock");
-        this.scheduler = Validate.paramNotNull(builder.scheduler, "scheduler");
+        this(builder.clock,
+             resolveScheduler(builder),
+             builder.scheduler == null);
+    }
+
+    private RateLimiterTokenBucketStore(RateLimiterClock clock, ScheduledExecutorService scheduler, boolean closeScheduler) {
+        this.clock = Validate.paramNotNull(clock, "clock");
+        this.scheduler = Validate.paramNotNull(scheduler, "scheduler");
+        this.closeScheduler = closeScheduler;
         this.scopeToTokenBucket = LruCache.<String, RateLimiterTokenBucket>builder(
-            x -> new RateLimiterTokenBucket(clock, scheduler))
+                                              x -> new RateLimiterTokenBucket(clock, scheduler))
                                           .maxSize(MAX_ENTRIES)
                                           .build();
     }
 
     @Override
     public void close() {
         scopeToTokenBucket.evictAll();
-        scheduler.shutdownNow();
+        if (closeScheduler) {
+            scheduler.shutdownNow();
+        }
     }
 
     public RateLimiterTokenBucket tokenBucketForScope(String scope) {
         return scopeToTokenBucket.get(scope);
     }
 
+    @SdkTestInternalApi
+    ScheduledExecutorService scheduler() {
+        return scheduler;
+    }
+
+    private static ScheduledExecutorService resolveScheduler(Builder b) {
+        if (b.scheduler != null) {
+            return b.scheduler;
+        }
+        return Executors.newSingleThreadScheduledExecutor(new ThreadFactoryBuilder()
+                                                              .daemonThreads(true)
+                                                              .threadNamePrefix(THREAD_NAME_PREFIX)
+                                                              .build());
+    }
+
     @Override
     @ToBuilderIgnoreField("scopeToTokenBucket")
     public Builder toBuilder() {
@@ -83,7 +113,13 @@ public Builder clock(RateLimiterClock clock) {
             return this;
         }
 
-        public Builder executor(ScheduledExecutorService scheduler) {
+        /**
+         * The scheduler used by the {@link RateLimiterTokenBucket rate limter buckets} to perform async notifications.
+         * The configured scheduler <strong>will not</strong> be closed when {@link #close() closing} this bucket store.
+         *
+         * @return This object for method chaining.
+         */
+        public Builder scheduler(ScheduledExecutorService scheduler) {
             this.scheduler = scheduler;
             return this;
         }

core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java

@@ -18,27 +18,39 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import org.junit.jupiter.api.Test;
 
 public class RateLimiterTokenBucketStoreTest {
     @Test
-    void close_closesScheduler() {
+    void close_schedulerProvided_schedulerNotClosed() {
         ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
         RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder()
                                                                        .clock(new SystemClock())
-                                                                       .executor(scheduler)
+                                                                       .scheduler(scheduler)
                                                                        .build();
 
         store.close();
 
-        verify(scheduler).shutdownNow();
+        verify(scheduler, never()).shutdownNow();
+        verify(scheduler, never()).shutdown();
+    }
+
+    @Test
+    void close_schedulerNotProvidedOnBuilder_schedulerClosed() {
+        RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder()
+                                                                       .clock(new SystemClock())
+                                                                       .build();
+
+        store.close();
+
+        assertThat(store.scheduler().isShutdown()).isTrue();
     }
 
     @Test
@@ -49,7 +61,7 @@ void close_closesAllCacheEntries() {
         ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
         RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder()
                                                                        .clock(new SystemClock())
-                                                                       .executor(scheduler)
+                                                                       .scheduler(scheduler)
                                                                        .build();
 
         List<CompletableFuture<Void>> futures = new ArrayList<>(entries);

test/codegen-generated-classes-test/src/main/resources/codegen-resources/json/service-2.json

@@ -0,0 +1,38 @@
+{
+  "version":"2.0",
+  "metadata":{
+    "apiVersion":"2010-05-08",
+    "endpointPrefix":"json-service-endpoint",
+    "globalEndpoint": "json-service.amazonaws.com",
+    "jsonVersion":"1.0",
+    "protocol":"json",
+    "serviceAbbreviation":"Aws Json Service",
+    "serviceFullName":"Some Service That Uses AWS JSON",
+    "serviceId":"Aws Json Service",
+    "signingName": "aws-json-service",
+    "signatureVersion":"v4",
+    "uid":"aws-json-service-2010-05-08",
+    "awsQueryCompatible":{}
+  },
+  "operations":{
+    "AllType": {
+      "name": "APostOperation",
+      "http": {
+        "method": "POST",
+        "requestUri": "/"
+      },
+      "httpChecksumRequired": true
+    }
+  },
+  "shapes": {
+    "OneShape": {
+      "type": "structure",
+      "members": {
+        "StringMember": {
+          "shape": "String"
+        }
+      }
+    },
+    "String":{"type":"string"}
+  }
+}