Fix adaptive retries implementation

.changes/next-release/bugfix-AWSSDKforJavav2-5417813.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Fixed the client-side rate limiting behavior of `AdaptiveRetryStrategy` under throttling conditions. Previously, the computed rate limit was overly aggressive, potentially allowing a request rate much lower than what the service allows, and was unstable, varying over time even when the service's allowed throughput was constant."
+}

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-spi/pom.xml

@@ -64,5 +64,11 @@
             <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-spi/src/main/java/software/amazon/awssdk/retries/api/RetryStrategy.java

@@ -15,9 +15,12 @@
 
 package software.amazon.awssdk.retries.api;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
 import software.amazon.awssdk.utils.builder.SdkBuilder;
 
 /**
@@ -40,9 +43,9 @@
  */
 @ThreadSafe
 @SdkPublicApi
-public interface RetryStrategy {
+public interface RetryStrategy extends SdkAutoCloseable {
     /**
-     * Invoked before the first request attempt.
+     * Acquire a retry token synchronously. Invoked before the first request attempt.
      *
      * <p>Callers MUST wait for the {@code delay} returned by this call before making the first attempt. Callers that wish to
      * retry a failed attempt MUST call {@link #refreshRetryToken} before doing so.
@@ -55,7 +58,27 @@ public interface RetryStrategy {
     AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request);
 
     /**
-     * Invoked before each subsequent (non-first) request attempt.
+     * Acquire a retry token asynchronously. Invoked before the first request attempt.
+     *
+     * <p>Callers MUST wait for the {@code delay} returned by this call before making the first attempt. Callers that wish to
+     * retry a failed attempt MUST call {@link #refreshRetryToken} before doing so.
+     *
+     * <p>If the attempt was successful, callers MUST call {@link #recordSuccess}.
+     *
+     * <p>The future is completed exceptionally with {@link NullPointerException} if a required parameter is {@code null}.
+     * <p>The future is completed exceptionally with {@link TokenAcquisitionFailedException} if a token cannot be acquired.
+     */
+    default CompletableFuture<AcquireInitialTokenResponse> acquireInitialTokenAsync(AcquireInitialTokenRequest request) {
+        try {
+            AcquireInitialTokenResponse response = acquireInitialToken(request);
+            return CompletableFuture.completedFuture(response);
+        } catch (Throwable t) {
+            return CompletableFutureUtils.failedFuture(t);
+        }
+    }
+
+    /**
+     * Refresh a retry token synchronously. Invoked before each subsequent (non-first) request attempt.
      *
      * <p>Callers MUST wait for the {@code delay} returned by this call before making the next attempt. If the next attempt
      * fails, callers MUST re-call {@link #refreshRetryToken} before attempting another retry. This call invalidates the provided
@@ -70,6 +93,29 @@ public interface RetryStrategy {
      */
     RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request);
 
+    /**
+     * Refresh a retry token asynchronously. Invoked before each subsequent (non-first) request attempt.
+     *
+     * <p>Callers MUST wait for the {@code delay} returned by this call before making the next attempt. If the next attempt
+     * fails, callers MUST re-call {@link #refreshRetryToken} before attempting another retry. This call invalidates the provided
+     * token, and returns a new one. Callers MUST use the new token.
+     *
+     * <p>If the attempt was successful, callers MUST call {@link #recordSuccess}.
+     *
+     * <p>The future is completed exceptionally with {@link NullPointerException} if a required parameter is not specified.
+     * <p>The future is completed exceptionally with {@link IllegalArgumentException} if the provided token was not issued by
+     * this strategy or the provided token was already used for a previous refresh or success call.
+     * <p>The future is completed exceptionally with {@link TokenAcquisitionFailedException} if a token cannot be acquired.
+     */
+    default CompletableFuture<RefreshRetryTokenResponse> refreshRetryTokenAsync(RefreshRetryTokenRequest request) {
+        try {
+            RefreshRetryTokenResponse response = refreshRetryToken(request);
+            return CompletableFuture.completedFuture(response);
+        } catch (Throwable t) {
+            return CompletableFutureUtils.failedFuture(t);
+        }
+    }
+
     /**
      * Invoked after an attempt succeeds.
      *
@@ -102,6 +148,10 @@ default boolean useClientDefaults() {
      */
     Builder<?, ?> toBuilder();
 
+    @Override
+    default void close() {
+    }
+
     /**
      * Builder to create immutable instances of {@link RetryStrategy}.
      */

core/retries-spi/src/test/java/software/amazon/awssdk/retries/api/RetryStrategyTest.java

@@ -0,0 +1,61 @@
+/*
+ * 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.api;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.junit.jupiter.api.Test;
+
+public class RetryStrategyTest {
+    @Test
+    void acquireInitialTokenAsync_syncThrows_reportedThroughFuture() {
+        RetryStrategy retryStrategy = mock(RetryStrategy.class);
+
+        when(retryStrategy.acquireInitialTokenAsync(any(AcquireInitialTokenRequest.class)))
+            .thenCallRealMethod();
+
+        RuntimeException t = new RuntimeException("oops");
+        when(retryStrategy.acquireInitialToken(any(AcquireInitialTokenRequest.class)))
+            .thenThrow(t);
+
+        assertThatThrownBy(() -> retryStrategy.acquireInitialTokenAsync(AcquireInitialTokenRequest.create("test")).join())
+            .hasRootCause(t);
+    }
+
+    @Test
+    void refreshRetryTokenAsync_syncThrows_reportedThroughFuture() {
+        RetryStrategy retryStrategy = mock(RetryStrategy.class);
+
+        when(retryStrategy.refreshRetryTokenAsync(any(RefreshRetryTokenRequest.class)))
+            .thenCallRealMethod();
+
+        RuntimeException t = new RuntimeException("oops");
+        when(retryStrategy.refreshRetryToken(any(RefreshRetryTokenRequest.class)))
+            .thenThrow(t);
+
+        RetryToken token = mock(RetryToken.class);
+        RefreshRetryTokenRequest request = RefreshRetryTokenRequest.builder()
+                                                                   .token(token)
+                                                                   .failure(new RuntimeException("failure"))
+                                                                   .build();
+
+        assertThatThrownBy(() -> retryStrategy.refreshRetryTokenAsync(request).join())
+            .hasRootCause(t);
+    }
+}

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/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,16 +47,53 @@ public final class DefaultAdaptiveRetryStrategy
     }
 
     @Override
-    protected Duration computeInitialBackoff(AcquireInitialTokenRequest request) {
+    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());
-        return bucket.tryAcquire().delay();
+        CompletableFuture<Void> acquireResult = bucket.acquireAsync();
+
+        return acquireResult.thenApply(r -> {
+            DefaultRetryToken token = DefaultRetryToken.builder().scope(request.scope()).build();
+            return AcquireInitialTokenResponse.create(token, Duration.ZERO);
+        });
     }
 
     @Override
-    protected Duration computeBackoff(RefreshRetryTokenRequest request, DefaultRetryToken token) {
-        Duration backoff = super.computeBackoff(request, token);
+    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());
-        return backoff.plus(bucket.tryAcquire().delay());
+        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
+    public void close() {
+        rateLimiterTokenBucketStore.close();
     }
 
     @Override