Skip to content

Update AsyncRetryableStage + tests #7192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dagnir wants to merge 2 commits into
base: feature/master/adaptive-retries-fixes
Choose a base branch
from
Open

Update AsyncRetryableStage + tests #7192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f4ab985
Update AsyncRetryableStage + tests
dagnir Jul 21, 2026
70480b3
Handle race condition in CRT tx pause
dagnir Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 37 additions & 33 deletions ...n/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@
import software.amazon.awssdk.http.SdkHttpFullRequest;
import software.amazon.awssdk.http.SdkHttpResponse;
import software.amazon.awssdk.utils.CompletableFutureUtils;
import software.amazon.awssdk.utils.Either;
import software.amazon.awssdk.utils.Logger;

/**
Expand Down Expand Up @@ -90,14 +89,14 @@ public CompletableFuture<Response<OutputT>> execute() {
}

public void attemptFirstExecute(CompletableFuture<Response<OutputT>> future) {
Duration backoffDelay = retryableStageHelper.acquireInitialToken();
if (backoffDelay.isZero()) {
attemptExecute(future);
} else {
retryableStageHelper.logBackingOff(backoffDelay);
long totalDelayMillis = backoffDelay.toMillis();
scheduledExecutor.schedule(() -> attemptExecute(future), totalDelayMillis, MILLISECONDS);
}
retryableStageHelper.acquireInitialTokenAsync().whenComplete((r, t) -> {
if (t != null) {
future.completeExceptionally(t);
return;
}
scheduledExecutor.schedule(() -> attemptExecute(future), r.toMillis(), MILLISECONDS);
});

}

private void attemptExecute(CompletableFuture<Response<OutputT>> future) {
Expand Down Expand Up @@ -140,33 +139,38 @@ private void attemptExecute(CompletableFuture<Response<OutputT>> future) {
}

public void maybeAttemptExecute(CompletableFuture<Response<OutputT>> future) {
Either<Duration, Duration> backoffDelay = retryableStageHelper.tryRefreshToken(suggestedDelay());

Optional<Duration> acquireFailureDelay = backoffDelay.right();
if (acquireFailureDelay.isPresent()) {
Duration delay = acquireFailureDelay.get();
retryableStageHelper.logAcquireFailureBackingOff(delay);
SdkException disallowedException = retryableStageHelper.retryPolicyDisallowedRetryException();
// Avoid needless scheduling if we won't wait
if (delay.isZero()) {
future.completeExceptionally(disallowedException);
} else {
scheduledExecutor.schedule(() -> future.completeExceptionally(disallowedException),
delay.toMillis(), MILLISECONDS);
retryableStageHelper.tryRefreshTokenAsync(suggestedDelay()).whenComplete((backoffDelay, t) -> {
if (t != null) {
future.completeExceptionally(t);
return;
}
return;
}
// We failed the last attempt, but will retry. The response handler wants to know when that happens.
responseHandler.onError(retryableStageHelper.getLastException());

// Reset the request provider to the original one before retries, in case it was modified downstream.
context.requestProvider(originalRequestBody);
Optional<Duration> acquireFailureDelay = backoffDelay.right();
if (acquireFailureDelay.isPresent()) {
Duration delay = acquireFailureDelay.get();
retryableStageHelper.logAcquireFailureBackingOff(delay);
SdkException disallowedException = retryableStageHelper.retryPolicyDisallowedRetryException();
// Avoid needless scheduling if we won't wait
if (delay.isZero()) {
future.completeExceptionally(disallowedException);
} else {
scheduledExecutor.schedule(() -> future.completeExceptionally(disallowedException),
delay.toMillis(), MILLISECONDS);
}
return;
}
// We failed the last attempt, but will retry. The response handler wants to know when that happens.
responseHandler.onError(retryableStageHelper.getLastException());

// Reset the request provider to the original one before retries, in case it was modified downstream.
context.requestProvider(originalRequestBody);

// get() is safe, Either requires left OR right to be present
Duration successDelay = backoffDelay.left().get();
retryableStageHelper.logBackingOff(successDelay);
long totalDelayMillis = successDelay.toMillis();
scheduledExecutor.schedule(() -> attemptExecute(future), totalDelayMillis, MILLISECONDS);
// get() is safe, Either requires left OR right to be present
Duration successDelay = backoffDelay.left().get();
retryableStageHelper.logBackingOff(successDelay);
long totalDelayMillis = successDelay.toMillis();
scheduledExecutor.schedule(() -> attemptExecute(future), totalDelayMillis, MILLISECONDS);
});
}

private void maybeRetryExecute(CompletableFuture<Response<OutputT>> future, Exception exception) {
Expand Down
106 changes: 87 additions & 19 deletions ...software/amazon/awssdk/core/internal/http/pipeline/stages/utils/RetryableStageHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.core.Response;
Expand Down Expand Up @@ -60,6 +61,7 @@ public final class RetryableStageHelper {
public static final String SDK_RETRY_INFO_HEADER = "amz-sdk-request";
public static final ExecutionAttribute<Duration> LAST_BACKOFF_DELAY_DURATION =
new ExecutionAttribute<>("LastBackoffDuration");
private static final String RETRY_TOKEN_SCOPE = "GLOBAL";

private final SdkHttpFullRequest request;
private final boolean isLongPollingOperation;
Expand All @@ -72,6 +74,7 @@ public final class RetryableStageHelper {
private SdkHttpResponse lastResponse;
private SdkException lastException;


public RetryableStageHelper(SdkHttpFullRequest request,
RequestExecutionContext context,
HttpClientDependencies dependencies) {
Expand Down Expand Up @@ -106,14 +109,31 @@ public void startingAttempt() {
* value is {@link AdaptiveRetryStrategy}.
*/
public Duration acquireInitialToken() {
String scope = "GLOBAL";
AcquireInitialTokenRequest acquireRequest = AcquireInitialTokenRequest.create(scope);
AcquireInitialTokenResponse acquireResponse = retryStrategy().acquireInitialToken(acquireRequest);
RetryToken retryToken = acquireResponse.token();
Duration delay = acquireResponse.delay();
context.executionAttributes().putAttribute(RETRY_TOKEN, retryToken);
context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, delay);
return delay;
AcquireInitialTokenResponse result = doBlockAcquireInitialToken();
context.executionAttributes().putAttribute(RETRY_TOKEN, result.token());
context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, result.delay());
return result.delay();
}

private AcquireInitialTokenResponse doBlockAcquireInitialToken() {
AcquireInitialTokenRequest acquireRequest = AcquireInitialTokenRequest.create(RETRY_TOKEN_SCOPE);

return retryStrategy().acquireInitialToken(acquireRequest);
}

public CompletableFuture<Duration> acquireInitialTokenAsync() {
AcquireInitialTokenRequest acquireRequest = AcquireInitialTokenRequest.create(RETRY_TOKEN_SCOPE);

return retryStrategy().acquireInitialTokenAsync(acquireRequest).whenComplete(
(result, t) -> {
if (t != null) {
return;
}

context.executionAttributes().putAttribute(RETRY_TOKEN, result.token());
context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, result.delay());
})
.thenApply(AcquireInitialTokenResponse::delay);
}

/**
Expand All @@ -138,16 +158,12 @@ public void recordAttemptSucceeded() {
* code should not retry.
*/
public Either<Duration, Duration> tryRefreshToken(Duration suggestedDelay) {
RetryToken retryToken = context.executionAttributes().getAttribute(RETRY_TOKEN);
RefreshRetryTokenResponse refreshResponse;
RetryToken retryToken;
Duration attemptDelay;
try {
RefreshRetryTokenRequest refreshRequest = RefreshRetryTokenRequest.builder()
.failure(this.lastException)
.token(retryToken)
.isLongPolling(isLongPollingOperation)
.suggestedDelay(suggestedDelay)
.build();
refreshResponse = retryStrategy().refreshRetryToken(refreshRequest);
RefreshRetryTokenResponse tryRefreshResult = doBlockingRefreshRetryToken(suggestedDelay);
attemptDelay = tryRefreshResult.delay();
retryToken = tryRefreshResult.token();
} catch (TokenAcquisitionFailedException e) {
context.executionAttributes().putAttribute(RETRY_TOKEN, e.token());
Optional<Duration> acquireFailureDelay = e.delay();
Expand All @@ -158,12 +174,64 @@ public Either<Duration, Duration> tryRefreshToken(Duration suggestedDelay) {
}
return Either.right(Duration.ZERO);
}
Duration acquireSuccessDelay = refreshResponse.delay();
context.executionAttributes().putAttribute(RETRY_TOKEN, refreshResponse.token());
Duration acquireSuccessDelay = attemptDelay;
context.executionAttributes().putAttribute(RETRY_TOKEN, retryToken);
context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, acquireSuccessDelay);
return Either.left(acquireSuccessDelay);
}

private RefreshRetryTokenResponse doBlockingRefreshRetryToken(Duration suggestedDelay) {
RetryToken retryToken = context.executionAttributes().getAttribute(RETRY_TOKEN);

RefreshRetryTokenRequest refreshRequest = RefreshRetryTokenRequest.builder()
.failure(this.lastException)
.token(retryToken)
.isLongPolling(isLongPollingOperation)
.suggestedDelay(suggestedDelay)
.build();

return retryStrategy().refreshRetryToken(refreshRequest);
}

public CompletableFuture<Either<Duration, Duration>> tryRefreshTokenAsync(Duration suggestedDelay) {
CompletableFuture<Either<Duration, Duration>> cf = new CompletableFuture<>();

RetryToken retryToken = context.executionAttributes().getAttribute(RETRY_TOKEN);

RefreshRetryTokenRequest refreshRequest = RefreshRetryTokenRequest.builder()
.failure(this.lastException)
.token(retryToken)
.isLongPolling(isLongPollingOperation)
.suggestedDelay(suggestedDelay)
.build();

retryStrategy().refreshRetryTokenAsync(refreshRequest).whenComplete((r, t) -> {
if (t != null) {
if (t instanceof TokenAcquisitionFailedException) {
TokenAcquisitionFailedException e = (TokenAcquisitionFailedException) t;
context.executionAttributes().putAttribute(RETRY_TOKEN, e.token());
Optional<Duration> acquireFailureDelay = e.delay();
if (acquireFailureDelay.isPresent()) {
Duration acquireDelay = acquireFailureDelay.get();
context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, acquireDelay);
cf.complete(Either.right(acquireDelay));

@davidh44 davidh44 Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need a return here in this branch after complete(). Maybe a test case for this?

@dagnir dagnir Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah good catch. I don't think it actually matters though, because if you fall through to the second complete, it becomes a no-op.

}
cf.complete(Either.right(Duration.ZERO));
} else {
cf.completeExceptionally(t);
}
return;
}

Duration acquireSuccessDelay = r.delay();
context.executionAttributes().putAttribute(RETRY_TOKEN, r.token());
context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, acquireSuccessDelay);
cf.complete(Either.left(acquireSuccessDelay));
});

return cf;
}

/**
* Return the exception that should be thrown, because the retry strategy did not allow the request to be retried.
*/
Expand Down
4 changes: 4 additions & 0 deletions ...c/test/java/software/amazon/awssdk/core/client/AsyncClientRetryStrategyExceptionTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ public class AsyncClientRetryStrategyExceptionTest {
public void exceptionInInitialTokenReportedInFuture() {
Exception exception = new RuntimeException(MESSAGE);
when(retryStrategy.acquireInitialToken(any())).thenThrow(exception);
when(retryStrategy.acquireInitialTokenAsync(any())).thenCallRealMethod();

CompletableFuture<SdkResponse> responseFuture = makeRequest();

Expand All @@ -73,8 +74,11 @@ public void exceptionInRefreshTokenReportedInFuture() {
AcquireInitialTokenResponse.create(new RetryToken() {
}, Duration.ZERO)
);
when(retryStrategy.acquireInitialTokenAsync(any())).thenCallRealMethod();

Exception exception = new RuntimeException(MESSAGE);
when(retryStrategy.refreshRetryToken(any())).thenThrow(exception);
when(retryStrategy.refreshRetryTokenAsync(any())).thenCallRealMethod();

CompletableFuture<SdkResponse> responseFuture = makeRequest();

Expand Down
13 changes: 12 additions & 1 deletion ...va/software/amazon/awssdk/core/client/handler/AsyncClientHandlerSignerValidationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import static org.assertj.core.api.Assertions.assertThatNoException;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.util.Arrays;
Expand All @@ -26,6 +28,8 @@
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
Expand Down Expand Up @@ -179,7 +183,6 @@ public void execute_requestHasHttpsEndpoint_usesBearerAuth_succeeds() throws Exc

SdkAsyncClientHandler sdkAsyncClientHandler = new SdkAsyncClientHandler(testClientConfiguration());
CompletableFuture<SdkResponse> execute = sdkAsyncClientHandler.execute(executionParams());

SdkAsyncHttpResponseHandler capturedHandler = executeRequestCaptor.getValue().responseHandler();
Map<String, List<String>> headers = new HashMap<>();
headers.put("foo", Arrays.asList("bar"));
Expand Down Expand Up @@ -227,10 +230,18 @@ private ClientExecutionParams<SdkRequest, SdkResponse> executionParams() {
}

private SdkClientConfiguration testClientConfiguration() {
ScheduledExecutorService exec = mock(ScheduledExecutorService.class);
when(exec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> {
Runnable r = i.getArgument(0);
r.run();
return null;
});

return HttpTestUtils.testClientConfiguration()
.toBuilder()
.option(SdkClientOption.ASYNC_HTTP_CLIENT, httpClient)
.option(SdkAdvancedClientOption.SIGNER, signer)
.option(SdkClientOption.SCHEDULED_EXECUTOR_SERVICE, exec)
.build();
}
}
10 changes: 10 additions & 0 deletions ...core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

Expand All @@ -27,6 +29,7 @@
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import org.junit.Before;
import org.junit.Test;
Expand Down Expand Up @@ -143,10 +146,17 @@ private ClientExecutionParams<SdkRequest, SdkResponse> clientExecutionParams() {
}

public SdkClientConfiguration clientConfiguration() {
ScheduledExecutorService exec = mock(ScheduledExecutorService.class);
when(exec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> {
Runnable r = i.getArgument(0);
r.run();
return null;
});
return HttpTestUtils.testClientConfiguration().toBuilder()
.option(SdkClientOption.ASYNC_HTTP_CLIENT, httpClient)
.option(SdkClientOption.RETRY_POLICY, RetryPolicy.none())
.option(SdkClientOption.RETRY_STRATEGY, DefaultRetryStrategy.doNotRetry())
.option(SdkClientOption.SCHEDULED_EXECUTOR_SERVICE, exec)
.build();
}
}
Loading
Loading