Skip to content

Commit 5db9f3d

Browse files
dagnirdagnir
authored
Retry on IDPCommunicationError (#6923)
This commit adds a fix to STS to retry on IDPCommunicationError. This is implemented as a customization for the retry strategy lookup for the client. This is gated behind the AWS_NEW_RETRIES_2026 flag.
1 parent 6e459be commit 5db9f3d

4 files changed

Lines changed: 187 additions & 2 deletions

File tree

‎services/sts/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,12 @@
6363
</dependency>
6464

6565
<!-- Test Dependencies -->
66+
<dependency>
67+
<groupId>software.amazon.awssdk</groupId>
68+
<artifactId>retries</artifactId>
69+
<version>${awsjavasdk.version}</version>
70+
<scope>test</scope>
71+
</dependency>
6672
<dependency>
6773
<artifactId>iam</artifactId>
6874
<groupId>software.amazon.awssdk</groupId>

‎services/sts/src/main/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategy.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services.sts.internal;
17+
18+
import software.amazon.awssdk.annotations.SdkInternalApi;
19+
import software.amazon.awssdk.awscore.retry.AwsRetryStrategy;
20+
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
21+
import software.amazon.awssdk.core.client.config.SdkClientOption;
22+
import software.amazon.awssdk.core.retry.NewRetries2026Resolver;
23+
import software.amazon.awssdk.core.retry.RetryMode;
24+
import software.amazon.awssdk.retries.api.RetryStrategy;
25+
import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException;
26+
27+
/**
28+
* Specialized retry strategy resolution for STS to enable retrying for {@link IdpCommunicationErrorException}.
29+
*/
30+
@SdkInternalApi
31+
public final class StsRetryStrategy {
32+
33+
private StsRetryStrategy() {
34+
}
35+
36+
public static RetryStrategy resolveRetryStrategy(SdkClientConfiguration config) {
37+
RetryStrategy configuredRetryStrategy = config.option(SdkClientOption.RETRY_STRATEGY);
38+
if (configuredRetryStrategy != null) {
39+
return configuredRetryStrategy;
40+
}
41+
42+
// Just return null and let the normal retry strategy resolution occur
43+
if (!isNewRetries2026Enabled(config)) {
44+
return null;
45+
}
46+
47+
RetryMode retryMode = resolveRetryMode(config);
48+
49+
return AwsRetryStrategy.forRetryMode(retryMode, true)
50+
.toBuilder()
51+
.retryOnException(IdpCommunicationErrorException.class)
52+
.build();
53+
}
54+
55+
private static RetryMode resolveRetryMode(SdkClientConfiguration config) {
56+
return RetryMode.resolver()
57+
.profileFile(config.option(SdkClientOption.PROFILE_FILE_SUPPLIER))
58+
.profileName(config.option(SdkClientOption.PROFILE_NAME))
59+
.defaultRetryMode(config.option(SdkClientOption.DEFAULT_RETRY_MODE))
60+
.defaultNewRetries2026(config.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026))
61+
.resolve();
62+
}
63+
64+
private static boolean isNewRetries2026Enabled(SdkClientConfiguration config) {
65+
Boolean defaultNewRetries2026 = config.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026);
66+
return new NewRetries2026Resolver().defaultNewRetries2026(defaultNewRetries2026).resolve();
67+
}
68+
}

‎services/sts/src/main/resources/codegen-resources/customization.config‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,6 @@
2121
"UseGlobalEndpoint with legacy region `us-west-2`": "V2 does not support setting UseGlobalEndpoint. It's regional only/by default"
2222
},
2323

24-
"enableGenerateCompiledEndpointRules": true
25-
24+
"enableGenerateCompiledEndpointRules": true,
25+
"customRetryStrategy" : "software.amazon.awssdk.services.sts.internal.StsRetryStrategy"
2626
}

‎services/sts/src/test/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategyTest.java‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services.sts.internal;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.mockito.Mockito.mock;
20+
21+
import java.time.Duration;
22+
import java.util.stream.Stream;
23+
import org.junit.jupiter.api.Test;
24+
import org.junit.jupiter.params.ParameterizedTest;
25+
import org.junit.jupiter.params.provider.Arguments;
26+
import org.junit.jupiter.params.provider.MethodSource;
27+
import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
28+
import software.amazon.awssdk.core.SdkSystemSetting;
29+
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
30+
import software.amazon.awssdk.core.client.config.SdkClientOption;
31+
import software.amazon.awssdk.retries.AdaptiveRetryStrategy;
32+
import software.amazon.awssdk.retries.LegacyRetryStrategy;
33+
import software.amazon.awssdk.retries.StandardRetryStrategy;
34+
import software.amazon.awssdk.retries.api.AcquireInitialTokenRequest;
35+
import software.amazon.awssdk.retries.api.RefreshRetryTokenRequest;
36+
import software.amazon.awssdk.retries.api.RetryStrategy;
37+
import software.amazon.awssdk.retries.api.RetryToken;
38+
import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException;
39+
40+
public class StsRetryStrategyTest {
41+
@Test
42+
void resolveRetryStrategy_preexistingStrategy_returnsPreexisting() {
43+
RetryStrategy strategy = mock(RetryStrategy.class);
44+
SdkClientConfiguration config = SdkClientConfiguration.builder()
45+
.option(SdkClientOption.RETRY_STRATEGY, strategy)
46+
.build();
47+
48+
assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isSameAs(strategy);
49+
}
50+
51+
@Test
52+
void resolveRetryStrategy_defaultNewRetries2026False_returnsNull() {
53+
SdkClientConfiguration config = SdkClientConfiguration.builder()
54+
.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026, false)
55+
.build();
56+
57+
assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isNull();
58+
}
59+
60+
@Test
61+
void resolveRetryStrategy_newRetries2026False_returnsNull() {
62+
System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "false");
63+
try {
64+
SdkClientConfiguration config = SdkClientConfiguration.builder().build();
65+
assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isNull();
66+
} finally {
67+
System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property());
68+
}
69+
}
70+
71+
@ParameterizedTest
72+
@MethodSource("retryModeResolutionCases")
73+
void resolveRetryStrategy_returnsCorrectStrategyBasedOnMode(String mode, Class<?> expected) {
74+
System.setProperty(SdkSystemSetting.AWS_RETRY_MODE.property(), mode);
75+
System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "true");
76+
try {
77+
SdkClientConfiguration config = SdkClientConfiguration.builder().build();
78+
RetryStrategy resolved = StsRetryStrategy.resolveRetryStrategy(config);
79+
assertThat(resolved).isInstanceOf(expected);
80+
assertRetriesOnIdpCommunicationException(resolved);
81+
} finally {
82+
System.clearProperty(SdkSystemSetting.AWS_RETRY_MODE.property());
83+
System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property());
84+
}
85+
}
86+
87+
void assertRetriesOnIdpCommunicationException(RetryStrategy strategy) {
88+
RetryToken token = strategy.acquireInitialToken(AcquireInitialTokenRequest.create("test")).token();
89+
90+
AwsErrorDetails errorDetails = AwsErrorDetails.builder()
91+
.errorCode("IDPCommunicationError")
92+
.build();
93+
94+
IdpCommunicationErrorException failure = IdpCommunicationErrorException.builder()
95+
.awsErrorDetails(errorDetails)
96+
.build();
97+
RefreshRetryTokenRequest refresh = RefreshRetryTokenRequest.builder()
98+
.token(token)
99+
.failure(failure)
100+
.build();
101+
assertThat(strategy.refreshRetryToken(refresh).delay()).isGreaterThan(Duration.ZERO);
102+
}
103+
104+
private static Stream<Arguments> retryModeResolutionCases() {
105+
return Stream.of(
106+
Arguments.of("standard", StandardRetryStrategy.class),
107+
Arguments.of("legacy", LegacyRetryStrategy.class),
108+
Arguments.of("adaptive", AdaptiveRetryStrategy.class)
109+
);
110+
}
111+
}

0 commit comments

Comments
 (0)