Skip to content

Commit fd683d8

Browse files
gosargosar
authored
Use joinLikeSync when resolving identity in sync code (#3953)
1 parent 53e2ffe commit fd683d8

File tree

10 files changed

+118
-30
lines changed

10 files changed

+118
-30
lines changed

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,12 @@
2929
import com.squareup.javapoet.MethodSpec;
3030
import com.squareup.javapoet.ParameterizedTypeName;
3131
import com.squareup.javapoet.TypeSpec;
32+
import com.squareup.javapoet.WildcardTypeName;
3233
import java.net.URI;
3334
import java.util.ArrayList;
3435
import java.util.Collections;
3536
import java.util.List;
37+
import java.util.concurrent.CompletableFuture;
3638
import java.util.stream.Collectors;
3739
import software.amazon.awssdk.annotations.SdkInternalApi;
3840
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
@@ -58,9 +60,11 @@
5860
import software.amazon.awssdk.core.endpointdiscovery.EndpointDiscoveryRefreshCache;
5961
import software.amazon.awssdk.core.endpointdiscovery.EndpointDiscoveryRequest;
6062
import software.amazon.awssdk.core.metrics.CoreMetric;
63+
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
6164
import software.amazon.awssdk.metrics.MetricCollector;
6265
import software.amazon.awssdk.metrics.MetricPublisher;
6366
import software.amazon.awssdk.metrics.NoOpMetricCollector;
67+
import software.amazon.awssdk.utils.CompletableFutureUtils;
6468
import software.amazon.awssdk.utils.Logger;
6569

6670
public class SyncClientClass extends SyncClientInterface {
@@ -258,11 +262,18 @@ private List<MethodSpec> operationMethodSpecs(OperationModel opModel) {
258262
method.addStatement("$T cachedEndpoint = null", URI.class);
259263
method.beginControlFlow("if (endpointDiscoveryEnabled)");
260264

261-
method.addCode("$T key = $N.overrideConfiguration()", String.class, opModel.getInput().getVariableName())
265+
ParameterizedTypeName identityFutureTypeName =
266+
ParameterizedTypeName.get(ClassName.get(CompletableFuture.class),
267+
WildcardTypeName.subtypeOf(AwsCredentialsIdentity.class));
268+
method.addCode("$T identityFuture = $N.overrideConfiguration()",
269+
identityFutureTypeName,
270+
opModel.getInput().getVariableName())
262271
.addCode(" .flatMap($T::credentialsIdentityProvider)", AwsRequestOverrideConfiguration.class)
263272
.addCode(" .orElseGet(() -> clientConfiguration.option($T.CREDENTIALS_IDENTITY_PROVIDER))",
264273
AwsClientOption.class)
265-
.addCode(" .resolveIdentity().join().accessKeyId();");
274+
.addCode(" .resolveIdentity();");
275+
276+
method.addCode("$T key = $T.joinLikeSync(identityFuture).accessKeyId();", String.class, CompletableFutureUtils.class);
266277

267278
method.addCode("$1T endpointDiscoveryRequest = $1T.builder()", EndpointDiscoveryRequest.class)
268279
.addCode(" .required($L)", opModel.getInputShape().getEndpointDiscovery().isRequired())

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-sync.java

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import java.net.URI;
44
import java.util.Collections;
55
import java.util.List;
6+
import java.util.concurrent.CompletableFuture;
67
import software.amazon.awssdk.annotations.Generated;
78
import software.amazon.awssdk.annotations.SdkInternalApi;
89
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
@@ -19,6 +20,7 @@
1920
import software.amazon.awssdk.core.exception.SdkClientException;
2021
import software.amazon.awssdk.core.http.HttpResponseHandler;
2122
import software.amazon.awssdk.core.metrics.CoreMetric;
23+
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
2224
import software.amazon.awssdk.metrics.MetricCollector;
2325
import software.amazon.awssdk.metrics.MetricPublisher;
2426
import software.amazon.awssdk.metrics.NoOpMetricCollector;
@@ -39,6 +41,7 @@
3941
import software.amazon.awssdk.services.endpointdiscoverytest.transform.TestDiscoveryIdentifiersRequiredRequestMarshaller;
4042
import software.amazon.awssdk.services.endpointdiscoverytest.transform.TestDiscoveryOptionalRequestMarshaller;
4143
import software.amazon.awssdk.services.endpointdiscoverytest.transform.TestDiscoveryRequiredRequestMarshaller;
44+
import software.amazon.awssdk.utils.CompletableFutureUtils;
4245
import software.amazon.awssdk.utils.Logger;
4346

4447
/**
@@ -154,10 +157,12 @@ public TestDiscoveryIdentifiersRequiredResponse testDiscoveryIdentifiersRequired
154157
}
155158
URI cachedEndpoint = null;
156159
if (endpointDiscoveryEnabled) {
157-
String key = testDiscoveryIdentifiersRequiredRequest.overrideConfiguration()
158-
.flatMap(AwsRequestOverrideConfiguration::credentialsIdentityProvider)
159-
.orElseGet(() -> clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER))
160-
.resolveIdentity().join().accessKeyId();
160+
CompletableFuture<? extends AwsCredentialsIdentity> identityFuture =
161+
testDiscoveryIdentifiersRequiredRequest.overrideConfiguration()
162+
.flatMap(AwsRequestOverrideConfiguration::credentialsIdentityProvider)
163+
.orElseGet(() -> clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER))
164+
.resolveIdentity();
165+
String key = CompletableFutureUtils.joinLikeSync(identityFuture).accessKeyId();
161166
EndpointDiscoveryRequest endpointDiscoveryRequest = EndpointDiscoveryRequest.builder().required(true)
162167
.defaultEndpoint(clientConfiguration.option(SdkClientOption.ENDPOINT))
163168
.overrideConfiguration(testDiscoveryIdentifiersRequiredRequest.overrideConfiguration().orElse(null)).build();
@@ -211,10 +216,12 @@ public TestDiscoveryOptionalResponse testDiscoveryOptional(TestDiscoveryOptional
211216
boolean endpointOverridden = clientConfiguration.option(SdkClientOption.ENDPOINT_OVERRIDDEN) == Boolean.TRUE;
212217
URI cachedEndpoint = null;
213218
if (endpointDiscoveryEnabled) {
214-
String key = testDiscoveryOptionalRequest.overrideConfiguration()
215-
.flatMap(AwsRequestOverrideConfiguration::credentialsIdentityProvider)
216-
.orElseGet(() -> clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER))
217-
.resolveIdentity().join().accessKeyId();
219+
CompletableFuture<? extends AwsCredentialsIdentity> identityFuture =
220+
testDiscoveryOptionalRequest.overrideConfiguration()
221+
.flatMap(AwsRequestOverrideConfiguration::credentialsIdentityProvider)
222+
.orElseGet(() -> clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER))
223+
.resolveIdentity();
224+
String key = CompletableFutureUtils.joinLikeSync(identityFuture).accessKeyId();
218225
EndpointDiscoveryRequest endpointDiscoveryRequest = EndpointDiscoveryRequest.builder().required(false)
219226
.defaultEndpoint(clientConfiguration.option(SdkClientOption.ENDPOINT))
220227
.overrideConfiguration(testDiscoveryOptionalRequest.overrideConfiguration().orElse(null)).build();
@@ -275,10 +282,12 @@ public TestDiscoveryRequiredResponse testDiscoveryRequired(TestDiscoveryRequired
275282
}
276283
URI cachedEndpoint = null;
277284
if (endpointDiscoveryEnabled) {
278-
String key = testDiscoveryRequiredRequest.overrideConfiguration()
279-
.flatMap(AwsRequestOverrideConfiguration::credentialsIdentityProvider)
280-
.orElseGet(() -> clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER))
281-
.resolveIdentity().join().accessKeyId();
285+
CompletableFuture<? extends AwsCredentialsIdentity> identityFuture =
286+
testDiscoveryRequiredRequest.overrideConfiguration()
287+
.flatMap(AwsRequestOverrideConfiguration::credentialsIdentityProvider)
288+
.orElseGet(() -> clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER))
289+
.resolveIdentity();
290+
String key = CompletableFutureUtils.joinLikeSync(identityFuture).accessKeyId();
282291
EndpointDiscoveryRequest endpointDiscoveryRequest = EndpointDiscoveryRequest.builder().required(true)
283292
.defaultEndpoint(clientConfiguration.option(SdkClientOption.ENDPOINT))
284293
.overrideConfiguration(testDiscoveryRequiredRequest.overrideConfiguration().orElse(null)).build();

‎core/auth/src/main/java/software/amazon/awssdk/auth/credentials/AwsCredentialsProviderChain.java

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import software.amazon.awssdk.core.exception.SdkClientException;
2525
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
2626
import software.amazon.awssdk.identity.spi.IdentityProvider;
27+
import software.amazon.awssdk.utils.CompletableFutureUtils;
2728
import software.amazon.awssdk.utils.IoUtils;
2829
import software.amazon.awssdk.utils.Logger;
2930
import software.amazon.awssdk.utils.SdkAutoCloseable;
@@ -99,15 +100,13 @@ public static AwsCredentialsProviderChain of(IdentityProvider<? extends AwsCrede
99100
@Override
100101
public AwsCredentials resolveCredentials() {
101102
if (reuseLastProviderEnabled && lastUsedProvider != null) {
102-
// TODO: Exception handling for join?
103-
return CredentialUtils.toCredentials(lastUsedProvider.resolveIdentity().join());
103+
return CredentialUtils.toCredentials(CompletableFutureUtils.joinLikeSync(lastUsedProvider.resolveIdentity()));
104104
}
105105

106106
List<String> exceptionMessages = null;
107107
for (IdentityProvider<? extends AwsCredentialsIdentity> provider : credentialsProviders) {
108108
try {
109-
// TODO: Exception handling for join?
110-
AwsCredentialsIdentity credentials = provider.resolveIdentity().join();
109+
AwsCredentialsIdentity credentials = CompletableFutureUtils.joinLikeSync(provider.resolveIdentity());
111110

112111
log.debug(() -> "Loading credentials from " + provider);
113112

‎core/auth/src/main/java/software/amazon/awssdk/auth/credentials/CredentialUtils.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
2020
import software.amazon.awssdk.identity.spi.AwsSessionCredentialsIdentity;
2121
import software.amazon.awssdk.identity.spi.IdentityProvider;
22+
import software.amazon.awssdk.utils.CompletableFutureUtils;
2223

2324
@SdkProtectedApi
2425
public final class CredentialUtils {
@@ -101,8 +102,8 @@ public static AwsCredentialsProvider toCredentialsProvider(
101102
return (AwsCredentialsProvider) identityProvider;
102103
}
103104
return () -> {
104-
// TODO: Exception handling for CompletionException thrown from join?
105-
AwsCredentialsIdentity awsCredentialsIdentity = identityProvider.resolveIdentity().join();
105+
AwsCredentialsIdentity awsCredentialsIdentity =
106+
CompletableFutureUtils.joinLikeSync(identityProvider.resolveIdentity());
106107
return toCredentials(awsCredentialsIdentity);
107108
};
108109
}

‎core/auth/src/main/java/software/amazon/awssdk/auth/token/credentials/SdkTokenProviderChain.java

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import software.amazon.awssdk.core.exception.SdkClientException;
2626
import software.amazon.awssdk.identity.spi.IdentityProvider;
2727
import software.amazon.awssdk.identity.spi.TokenIdentity;
28+
import software.amazon.awssdk.utils.CompletableFutureUtils;
2829
import software.amazon.awssdk.utils.IoUtils;
2930
import software.amazon.awssdk.utils.Logger;
3031
import software.amazon.awssdk.utils.SdkAutoCloseable;
@@ -95,15 +96,13 @@ public static SdkTokenProviderChain of(IdentityProvider<? extends TokenIdentity>
9596
@Override
9697
public SdkToken resolveToken() {
9798
if (reuseLastProviderEnabled && lastUsedProvider != null) {
98-
// TODO: Exception handling for join?
99-
return TokenUtils.toSdkToken(lastUsedProvider.resolveIdentity().join());
99+
return TokenUtils.toSdkToken(CompletableFutureUtils.joinLikeSync(lastUsedProvider.resolveIdentity()));
100100
}
101101

102102
List<String> exceptionMessages = null;
103103
for (IdentityProvider<? extends TokenIdentity> provider : sdkTokenProviders) {
104104
try {
105-
// TODO: Exception handling for join?
106-
TokenIdentity token = provider.resolveIdentity().join();
105+
TokenIdentity token = CompletableFutureUtils.joinLikeSync(provider.resolveIdentity());
107106

108107
log.debug(() -> "Loading token from " + provider);
109108

‎services/polly/src/main/java/software/amazon/awssdk/services/polly/internal/presigner/DefaultPollyPresigner.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
import software.amazon.awssdk.services.polly.presigner.PollyPresigner;
6060
import software.amazon.awssdk.services.polly.presigner.model.PresignedSynthesizeSpeechRequest;
6161
import software.amazon.awssdk.services.polly.presigner.model.SynthesizeSpeechPresignRequest;
62+
import software.amazon.awssdk.utils.CompletableFutureUtils;
6263
import software.amazon.awssdk.utils.IoUtils;
6364
import software.amazon.awssdk.utils.Validate;
6465

@@ -212,7 +213,7 @@ private IdentityProvider<? extends AwsCredentialsIdentity> resolveCredentialsPro
212213
}
213214

214215
private AwsCredentialsIdentity resolveCredentials(IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider) {
215-
return credentialsProvider.resolveIdentity().join();
216+
return CompletableFutureUtils.joinLikeSync(credentialsProvider.resolveIdentity());
216217
}
217218

218219
private Presigner resolvePresigner(PollyRequest request) {

‎services/rds/src/main/java/software/amazon/awssdk/services/rds/DefaultRdsUtilities.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import software.amazon.awssdk.identity.spi.IdentityProvider;
3333
import software.amazon.awssdk.regions.Region;
3434
import software.amazon.awssdk.services.rds.model.GenerateAuthenticationTokenRequest;
35+
import software.amazon.awssdk.utils.CompletableFutureUtils;
3536
import software.amazon.awssdk.utils.StringUtils;
3637

3738
@Immutable
@@ -111,11 +112,12 @@ private Region resolveRegion(GenerateAuthenticationTokenRequest request) {
111112
// TODO: update this to use AwsCredentialsIdentity when we migrate Signers to accept the new type.
112113
private AwsCredentials resolveCredentials(GenerateAuthenticationTokenRequest request) {
113114
if (request.credentialsIdentityProvider() != null) {
114-
return CredentialUtils.toCredentials(request.credentialsIdentityProvider().resolveIdentity().join());
115+
return CredentialUtils.toCredentials(
116+
CompletableFutureUtils.joinLikeSync(request.credentialsIdentityProvider().resolveIdentity()));
115117
}
116118

117119
if (this.credentialsProvider != null) {
118-
return CredentialUtils.toCredentials(this.credentialsProvider.resolveIdentity().join());
120+
return CredentialUtils.toCredentials(CompletableFutureUtils.joinLikeSync(this.credentialsProvider.resolveIdentity()));
119121
}
120122

121123
throw new IllegalArgumentException("CredentialProvider should be provided either in GenerateAuthenticationTokenRequest " +

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crt/CrtCredentialsProviderAdapter.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
2727
import software.amazon.awssdk.identity.spi.AwsSessionCredentialsIdentity;
2828
import software.amazon.awssdk.identity.spi.IdentityProvider;
29+
import software.amazon.awssdk.utils.CompletableFutureUtils;
2930
import software.amazon.awssdk.utils.SdkAutoCloseable;
3031

3132
/**
@@ -45,8 +46,8 @@ public CrtCredentialsProviderAdapter(IdentityProvider<? extends AwsCredentialsId
4546
return Credentials.createAnonymousCredentials();
4647
}
4748

48-
// TODO: Exception handling for join?
49-
AwsCredentialsIdentity sdkCredentials = credentialsProvider.resolveIdentity().join();
49+
AwsCredentialsIdentity sdkCredentials =
50+
CompletableFutureUtils.joinLikeSync(credentialsProvider.resolveIdentity());
5051
byte[] accessKey = sdkCredentials.accessKeyId().getBytes(StandardCharsets.UTF_8);
5152
byte[] secreteKey = sdkCredentials.secretAccessKey().getBytes(StandardCharsets.UTF_8);
5253

‎utils/src/main/java/software/amazon/awssdk/utils/CompletableFutureUtils.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,4 +239,22 @@ public static void joinInterruptiblyIgnoringFailures(CompletableFuture<?> future
239239
// Ignore
240240
}
241241
}
242+
243+
/**
244+
* Joins (interruptibly) on the future, and re-throws any RuntimeExceptions or Errors just like the async task would have
245+
* thrown if it was executed synchronously.
246+
*/
247+
public static <T> T joinLikeSync(CompletableFuture<T> future) {
248+
try {
249+
return joinInterruptibly(future);
250+
} catch (CompletionException e) {
251+
Throwable cause = e.getCause();
252+
if (cause instanceof RuntimeException) {
253+
// Make sure we don't lose the context of where the join is in the stack...
254+
cause.addSuppressed(new RuntimeException("Task failed."));
255+
throw (RuntimeException) cause;
256+
}
257+
throw e;
258+
}
259+
}
242260
}

‎utils/src/test/java/software/amazon/awssdk/utils/CompletableFutureUtilsTest.java

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@
1919
import static org.assertj.core.api.Assertions.assertThatThrownBy;
2020
import static org.junit.Assert.fail;
2121

22+
import java.util.concurrent.CancellationException;
2223
import java.util.concurrent.CompletableFuture;
24+
import java.util.concurrent.CompletionException;
2325
import java.util.concurrent.ExecutorService;
2426
import java.util.concurrent.Executors;
2527
import org.junit.AfterClass;
@@ -159,4 +161,49 @@ public void allOfExceptionForwarded_allFutureSucceed_shouldComplete() {
159161
assertThat(resultFuture.isDone()).isTrue();
160162
assertThat(resultFuture.isCompletedExceptionally()).isFalse();
161163
}
162-
}
164+
165+
@Test(timeout = 1000)
166+
public void joinLikeSync_completesExceptionally_throwsUnderlyingException() {
167+
Exception e = new RuntimeException("BOOM");
168+
CompletableFuture future = new CompletableFuture();
169+
future.completeExceptionally(e);
170+
171+
assertThatThrownBy(() -> CompletableFutureUtils.joinLikeSync(future))
172+
.hasSuppressedException(new RuntimeException("Task failed."))
173+
.isEqualTo(e);
174+
}
175+
176+
@Test(timeout = 1000)
177+
public void joinLikeSync_completesExceptionallyChecked_throwsCompletionException() {
178+
Exception e = new Exception("BOOM");
179+
CompletableFuture future = new CompletableFuture();
180+
future.completeExceptionally(e);
181+
182+
assertThatThrownBy(() -> CompletableFutureUtils.joinLikeSync(future))
183+
.hasNoSuppressedExceptions()
184+
.hasCause(e)
185+
.isInstanceOf(CompletionException.class);
186+
}
187+
188+
@Test(timeout = 1000)
189+
public void joinLikeSync_completesExceptionallyWithError_throwsError() {
190+
Error e = new Error("BOOM");
191+
CompletableFuture future = new CompletableFuture();
192+
future.completeExceptionally(e);
193+
194+
assertThatThrownBy(() -> CompletableFutureUtils.joinLikeSync(future))
195+
.hasNoSuppressedExceptions()
196+
.isEqualTo(e);
197+
}
198+
199+
@Test(timeout = 1000)
200+
public void joinLikeSync_canceled_throwsCancellationException() {
201+
Exception e = new Exception("BOOM");
202+
CompletableFuture future = new CompletableFuture();
203+
future.cancel(false);
204+
205+
assertThatThrownBy(() -> CompletableFutureUtils.joinLikeSync(future))
206+
.hasNoSuppressedExceptions()
207+
.hasNoCause()
208+
.isInstanceOf(CancellationException.class);
209+
}}

0 commit comments

Comments
 (0)