diff --git a/.changes/next-release/feature-AWSSDKforJavav2-301f836.json b/.changes/next-release/feature-AWSSDKforJavav2-301f836.json new file mode 100644 index 000000000000..f3705d46eb0a --- /dev/null +++ b/.changes/next-release/feature-AWSSDKforJavav2-301f836.json @@ -0,0 +1,6 @@ +{ + "type": "feature", + "category": "AWS SDK for Java v2", + "contributor": "", + "description": "Cache auth scheme resolution results per operation" +} diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java index 379366b4182a..e82e3726e608 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java +++ b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java @@ -162,6 +162,10 @@ public boolean hasSigV4aSupport() { return usesSigV4a() || generateEndpointBasedParams(); } + public boolean hasPerOperationAuthOverrides() { + return AuthSchemeCodegenKnowledgeIndex.of(intermediateModel).hasPerOperationAuthSchemesOverrides(); + } + private static Set setOf(String val1, String val2) { Set result = new HashSet<>(); result.add(val1); diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java index 0525694513f4..1fdd891b7ae8 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java +++ b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java @@ -119,24 +119,24 @@ public AsyncClientClass(GeneratorTaskParams dependencies) { } @Override - protected TypeSpec.Builder createTypeSpec() { + protected Builder createTypeSpec() { return PoetUtils.createClassBuilder(className); } @Override - protected void addInterfaceClass(TypeSpec.Builder type) { + protected void addInterfaceClass(Builder type) { ClassName interfaceClass = poetExtensions.getClientClass(model.getMetadata().getAsyncInterface()); type.addSuperinterface(interfaceClass) .addJavadoc("Internal implementation of {@link $1T}.\n\n@see $1T#builder()", interfaceClass); } @Override - protected void addAnnotations(TypeSpec.Builder type) { + protected void addAnnotations(Builder type) { type.addAnnotation(SdkInternalApi.class); } @Override - protected void addModifiers(TypeSpec.Builder type) { + protected void addModifiers(Builder type) { type.addModifiers(FINAL); } @@ -165,6 +165,8 @@ protected void addFields(Builder type) { model.getEndpointOperation().ifPresent( o -> type.addField(EndpointDiscoveryRefreshCache.class, "endpointDiscoveryCache", PRIVATE)); + + ClientClassUtils.authSchemeCacheField(authSchemeSpecUtils).ifPresent(type::addField); } @Override diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java index 367af1cbf555..09b055c8219e 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java +++ b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java @@ -21,12 +21,15 @@ import com.squareup.javapoet.ClassName; import com.squareup.javapoet.CodeBlock; +import com.squareup.javapoet.FieldSpec; import com.squareup.javapoet.MethodSpec; import com.squareup.javapoet.ParameterSpec; import com.squareup.javapoet.ParameterizedTypeName; import com.squareup.javapoet.TypeName; import com.squareup.javapoet.TypeVariableName; import com.squareup.javapoet.WildcardTypeName; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; @@ -385,6 +388,7 @@ static MethodSpec resolveAuthSchemeOptionsMethod(AuthSchemeSpecUtils authSchemeS + ".orElse(null)", providerInterface, Validate.class, providerInterface, "Expected an instance of " + authSchemeSpecUtils.providerInterfaceName().simpleName()); + builder.addStatement("$T authSchemeProvider = requestAuthSchemeProvider != null " + "? requestAuthSchemeProvider " + ": $T.isInstanceOf($T.class, " @@ -393,13 +397,24 @@ static MethodSpec resolveAuthSchemeOptionsMethod(AuthSchemeSpecUtils authSchemeS SdkInternalExecutionAttribute.class, "Expected an instance of " + authSchemeSpecUtils.providerInterfaceName().simpleName()); + boolean canCache = !authSchemeSpecUtils.useEndpointBasedAuthProvider(); + if (canCache) { + addAuthSchemeCacheLookup(builder, authSchemeSpecUtils); + } + if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) { addEndpointBasedAuthSchemeResolution(builder, authSchemeSpecUtils, endpointRulesSpecUtils); } else { addSimpleAuthSchemeResolution(builder, authSchemeSpecUtils); } - if (endpointRulesSpecUtils.isS3()) { + if (canCache) { + builder.beginControlFlow("if (useCache)"); + builder.addStatement("options = $T.unmodifiableList(options)", Collections.class); + builder.addStatement("authSchemeCache.put(cacheKey, options)"); + builder.endControlFlow(); + builder.addStatement("return options"); + } else if (endpointRulesSpecUtils.isS3()) { ClassName sdkIdentityProperty = ClassName.get("software.amazon.awssdk.core.identity", "SdkIdentityProperty"); builder.addStatement("$T sdkClient = executionAttributes.getAttribute($T.SDK_CLIENT)", SdkClient.class, SdkInternalExecutionAttribute.class); @@ -414,6 +429,58 @@ static MethodSpec resolveAuthSchemeOptionsMethod(AuthSchemeSpecUtils authSchemeS return builder.build(); } + /** + * Returns a field spec for the auth scheme options cache, used when simple (non-endpoint-based) auth is in effect. + */ + static Optional authSchemeCacheField(AuthSchemeSpecUtils authSchemeSpecUtils) { + if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) { + return Optional.empty(); + } + ClassName concurrentHashMap = ClassName.get("java.util.concurrent", "ConcurrentHashMap"); + ParameterizedTypeName mapType = ParameterizedTypeName.get( + concurrentHashMap, + ClassName.get(String.class), + ParameterizedTypeName.get(ClassName.get(List.class), ClassName.get(AuthSchemeOption.class))); + return Optional.of(FieldSpec.builder(mapType, "authSchemeCache", PRIVATE, Modifier.FINAL) + .initializer("new $T<>()", concurrentHashMap) + .build()); + } + + private static void addAuthSchemeCacheLookup(MethodSpec.Builder builder, AuthSchemeSpecUtils authSchemeSpecUtils) { + ClassName defaultProviderClass = authSchemeSpecUtils.defaultAuthSchemeProviderName(); + builder.addStatement("boolean useCache = requestAuthSchemeProvider == null " + + "&& authSchemeProvider instanceof $T", defaultProviderClass); + + ClassName awsExecAttr = ClassName.get("software.amazon.awssdk.awscore", "AwsExecutionAttribute"); + List parts = new ArrayList<>(); + if (authSchemeSpecUtils.hasPerOperationAuthOverrides()) { + parts.add(CodeBlock.of("operationName")); + } + if (authSchemeSpecUtils.usesSigV4()) { + parts.add(CodeBlock.of("executionAttributes.getAttribute($T.AWS_REGION)", awsExecAttr)); + } + if (authSchemeSpecUtils.usesSigV4a()) { + parts.add(CodeBlock.of("executionAttributes.getAttribute($T.AWS_SIGV4A_SIGNING_REGION_SET)", awsExecAttr)); + } + + if (parts.isEmpty()) { + builder.addStatement("$T cacheKey = $S", String.class, "default"); + } else if (parts.size() == 1) { + builder.addStatement("$T cacheKey = $T.valueOf($L)", String.class, String.class, parts.get(0)); + } else { + builder.addStatement("$T cacheKey = $L", String.class, CodeBlock.join(parts, " + \":\" + ")); + } + + builder.beginControlFlow("if (useCache)"); + builder.addStatement("$T<$T> cached = authSchemeCache.get(cacheKey)", + List.class, AuthSchemeOption.class); + builder.beginControlFlow("if (cached != null)"); + builder.addStatement("return cached"); + builder.endControlFlow(); + builder.endControlFlow(); + } + + // Any AwsExecutionAttribute added here must also be added to addAuthSchemeCacheLookup(). Enforced by AuthSchemeCacheKeyTest. private static void addSimpleAuthSchemeResolution(MethodSpec.Builder builder, AuthSchemeSpecUtils authSchemeSpecUtils) { ClassName paramsInterface = authSchemeSpecUtils.parametersInterfaceName(); diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java index af0529666219..346e47fa9805 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java +++ b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java @@ -130,6 +130,7 @@ protected void addFields(TypeSpec.Builder type) { .addField(protocolSpec.protocolFactory(model)) .addField(SdkClientConfiguration.class, "clientConfiguration", PRIVATE, FINAL); protocolSpec.errorResponseMapperField().ifPresent(type::addField); + ClientClassUtils.authSchemeCacheField(authSchemeSpecUtils).ifPresent(type::addField); } @Override diff --git a/codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/AuthSchemeCacheKeyTest.java b/codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/AuthSchemeCacheKeyTest.java new file mode 100644 index 000000000000..eec629b11e58 --- /dev/null +++ b/codegen/src/test/java/software/amazon/awssdk/codegen/poet/client/AuthSchemeCacheKeyTest.java @@ -0,0 +1,95 @@ +/* + * 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.codegen.poet.client; + +import static org.assertj.core.api.Assertions.assertThat; +import static software.amazon.awssdk.codegen.poet.ClientTestModels.customPackageModels; +import static software.amazon.awssdk.codegen.poet.ClientTestModels.opsWithSigv4a; +import static software.amazon.awssdk.codegen.poet.ClientTestModels.restJsonServiceModels; + +import java.util.HashSet; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.junit.Test; +import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel; +import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils; +import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils; + +/** + * Verifies that every AwsExecutionAttribute used in auth scheme params building is also present in the cache key. + * If a new attribute is added to the params without updating the cache key, this test will fail. + */ +public class AuthSchemeCacheKeyTest { + + private static final Pattern AWS_EXEC_ATTR_PATTERN = + Pattern.compile("AwsExecutionAttribute\\.(\\w+)"); + + @Test + public void restJson_cacheKeyCoversAllParamAttributes() { + verifyConsistency(restJsonServiceModels()); + } + + @Test + public void sigv4a_cacheKeyCoversAllParamAttributes() { + verifyConsistency(opsWithSigv4a()); + } + + @Test + public void uniformAuth_cacheKeyCoversAllParamAttributes() { + verifyConsistency(customPackageModels()); + } + + private void verifyConsistency(IntermediateModel model) { + AuthSchemeSpecUtils authSchemeSpecUtils = new AuthSchemeSpecUtils(model); + EndpointRulesSpecUtils endpointRulesSpecUtils = new EndpointRulesSpecUtils(model); + + String source = ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils) + .toString(); + + int resolveCallIndex = source.indexOf(".resolveAuthScheme("); + assertThat(resolveCallIndex).as("resolveAuthScheme call should exist in generated method").isGreaterThan(0); + + int cacheKeyStart = source.indexOf("cacheKey ="); + if (cacheKeyStart < 0) { + // No cache — endpoint-based service, nothing to verify + return; + } + int cacheKeyEnd = source.indexOf(";", cacheKeyStart); + + int paramsStart = source.indexOf("paramsBuilder"); + String paramsSection = source.substring(paramsStart, resolveCallIndex); + + Set paramsAttributes = extractAttributes(paramsSection); + assertThat(paramsAttributes).as("Expected auth params to reference AwsExecutionAttributes").isNotEmpty(); + String cacheKeySection = source.substring(cacheKeyStart, cacheKeyEnd); + Set cacheKeyAttributes = extractAttributes(cacheKeySection); + + assertThat(cacheKeyAttributes) + .as("Cache key must include all AwsExecutionAttributes used in params building. " + + "If you added a new attribute to params, add it to addAuthSchemeCacheLookup() too.") + .containsAll(paramsAttributes); + } + + private Set extractAttributes(String section) { + Set attributes = new HashSet<>(); + Matcher matcher = AWS_EXEC_ATTR_PATTERN.matcher(section); + while (matcher.find()) { + attributes.add(matcher.group(1)); + } + return attributes; + } +} diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-json-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-json-async-client-class.java index 1d4e46dd67fb..3d9a8b62f388 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-json-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-json-async-client-class.java @@ -8,6 +8,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.function.Consumer; import java.util.function.Function; @@ -70,6 +71,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider; +import software.amazon.awssdk.services.json.auth.scheme.internal.DefaultJsonAuthSchemeProvider; import software.amazon.awssdk.services.json.endpoints.JsonEndpointParams; import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider; import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointResolverUtils; @@ -171,8 +173,11 @@ final class DefaultJsonAsyncClient implements JsonAsyncClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + private final Executor executor; + protected DefaultJsonAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -1262,9 +1267,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, JsonAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(JsonAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of JsonAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultJsonAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-async-client-class.java index e050a9abab4e..84728c99d5f1 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-async-client-class.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -51,6 +52,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.querytojsoncompatible.auth.scheme.QueryToJsonCompatibleAuthSchemeParams; import software.amazon.awssdk.services.querytojsoncompatible.auth.scheme.QueryToJsonCompatibleAuthSchemeProvider; +import software.amazon.awssdk.services.querytojsoncompatible.auth.scheme.internal.DefaultQueryToJsonCompatibleAuthSchemeProvider; import software.amazon.awssdk.services.querytojsoncompatible.endpoints.QueryToJsonCompatibleEndpointParams; import software.amazon.awssdk.services.querytojsoncompatible.endpoints.QueryToJsonCompatibleEndpointProvider; import software.amazon.awssdk.services.querytojsoncompatible.endpoints.internal.QueryToJsonCompatibleEndpointResolverUtils; @@ -97,6 +99,8 @@ final class DefaultQueryToJsonCompatibleAsyncClient implements QueryToJsonCompat } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultQueryToJsonCompatibleAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -228,10 +232,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(QueryToJsonCompatibleAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of QueryToJsonCompatibleAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultQueryToJsonCompatibleAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } QueryToJsonCompatibleAuthSchemeParams.Builder paramsBuilder = QueryToJsonCompatibleAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-sync-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-sync-client-class.java index 4bf12d2d6d22..f372a60313c3 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-sync-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-sync-client-class.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -46,6 +47,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.querytojsoncompatible.auth.scheme.QueryToJsonCompatibleAuthSchemeParams; import software.amazon.awssdk.services.querytojsoncompatible.auth.scheme.QueryToJsonCompatibleAuthSchemeProvider; +import software.amazon.awssdk.services.querytojsoncompatible.auth.scheme.internal.DefaultQueryToJsonCompatibleAuthSchemeProvider; import software.amazon.awssdk.services.querytojsoncompatible.endpoints.QueryToJsonCompatibleEndpointParams; import software.amazon.awssdk.services.querytojsoncompatible.endpoints.QueryToJsonCompatibleEndpointProvider; import software.amazon.awssdk.services.querytojsoncompatible.endpoints.internal.QueryToJsonCompatibleEndpointResolverUtils; @@ -92,6 +94,8 @@ final class DefaultQueryToJsonCompatibleClient implements QueryToJsonCompatibleC } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultQueryToJsonCompatibleClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -196,10 +200,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(QueryToJsonCompatibleAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of QueryToJsonCompatibleAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultQueryToJsonCompatibleAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } QueryToJsonCompatibleAuthSchemeParams.Builder paramsBuilder = QueryToJsonCompatibleAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-batchmanager-async.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-batchmanager-async.java index d502d9d8e359..36923ede6e30 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-batchmanager-async.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-batchmanager-async.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ScheduledExecutorService; import java.util.function.Consumer; import java.util.function.Function; @@ -52,6 +53,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.batchmanagertest.auth.scheme.BatchManagerTestAuthSchemeParams; import software.amazon.awssdk.services.batchmanagertest.auth.scheme.BatchManagerTestAuthSchemeProvider; +import software.amazon.awssdk.services.batchmanagertest.auth.scheme.internal.DefaultBatchManagerTestAuthSchemeProvider; import software.amazon.awssdk.services.batchmanagertest.batchmanager.BatchManagerTestAsyncBatchManager; import software.amazon.awssdk.services.batchmanagertest.endpoints.BatchManagerTestEndpointParams; import software.amazon.awssdk.services.batchmanagertest.endpoints.BatchManagerTestEndpointProvider; @@ -96,6 +98,8 @@ final class DefaultBatchManagerTestAsyncClient implements BatchManagerTestAsyncC private final ScheduledExecutorService executorService; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultBatchManagerTestAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -224,10 +228,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(BatchManagerTestAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of BatchManagerTestAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultBatchManagerTestAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } BatchManagerTestAuthSchemeParams.Builder paramsBuilder = BatchManagerTestAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-async-client-class.java index 21731fc695f0..af4884e50e56 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-async-client-class.java @@ -8,6 +8,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.function.Consumer; import java.util.function.Function; @@ -71,6 +72,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider; +import software.amazon.awssdk.services.json.auth.scheme.internal.DefaultJsonAuthSchemeProvider; import software.amazon.awssdk.services.json.endpoints.JsonEndpointParams; import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider; import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointResolverUtils; @@ -174,8 +176,11 @@ final class DefaultJsonAsyncClient implements JsonAsyncClient { private final AwsJsonProtocolFactory jsonProtocolFactory; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + private final Executor executor; + protected DefaultJsonAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -1266,9 +1271,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, JsonAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(JsonAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of JsonAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultJsonAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-client-class.java index 9ff973b70180..dfcd9c933cb1 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-cbor-client-class.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -51,6 +52,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider; +import software.amazon.awssdk.services.json.auth.scheme.internal.DefaultJsonAuthSchemeProvider; import software.amazon.awssdk.services.json.endpoints.JsonEndpointParams; import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider; import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointResolverUtils; @@ -131,6 +133,8 @@ final class DefaultJsonClient implements JsonClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultJsonClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -818,9 +822,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec .isInstanceOf(JsonAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of JsonAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultJsonAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-async-client-class.java index d47ca63513ba..23a669392d17 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-async-client-class.java @@ -8,6 +8,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -52,6 +53,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.foobar.auth.scheme.FooBarAuthSchemeParams; import software.amazon.awssdk.services.foobar.auth.scheme.FooBarAuthSchemeProvider; +import software.amazon.awssdk.services.foobar.auth.scheme.internal.DefaultFooBarAuthSchemeProvider; import software.amazon.awssdk.services.foobar.endpoints.FooBarClientContextParams; import software.amazon.awssdk.services.foobar.endpoints.FooBarEndpointParams; import software.amazon.awssdk.services.foobar.endpoints.FooBarEndpointProvider; @@ -95,6 +97,8 @@ final class DefaultFooBarAsyncClient implements FooBarAsyncClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultFooBarAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -218,9 +222,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec .isInstanceOf(FooBarAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of FooBarAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultFooBarAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } FooBarAuthSchemeParams.Builder paramsBuilder = FooBarAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-sync-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-sync-client-class.java index 5b357c23c042..b041e122f432 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-sync-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custom-context-params-sync-client-class.java @@ -5,6 +5,7 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -47,6 +48,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.foobar.auth.scheme.FooBarAuthSchemeParams; import software.amazon.awssdk.services.foobar.auth.scheme.FooBarAuthSchemeProvider; +import software.amazon.awssdk.services.foobar.auth.scheme.internal.DefaultFooBarAuthSchemeProvider; import software.amazon.awssdk.services.foobar.endpoints.FooBarClientContextParams; import software.amazon.awssdk.services.foobar.endpoints.FooBarEndpointParams; import software.amazon.awssdk.services.foobar.endpoints.FooBarEndpointProvider; @@ -90,6 +92,8 @@ final class DefaultFooBarClient implements FooBarClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultFooBarClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -188,9 +192,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, FooBarAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(FooBarAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of FooBarAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultFooBarAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } FooBarAuthSchemeParams.Builder paramsBuilder = FooBarAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-async.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-async.java index 90abd9e720bc..c2d400d99364 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-async.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-async.java @@ -4,6 +4,7 @@ import foo.bar.helloworld.auth.scheme.ProtocolRestJsonWithCustomPackageAuthSchemeParams; import foo.bar.helloworld.auth.scheme.ProtocolRestJsonWithCustomPackageAuthSchemeProvider; +import foo.bar.helloworld.auth.scheme.internal.DefaultProtocolRestJsonWithCustomPackageAuthSchemeProvider; import foo.bar.helloworld.endpoints.ProtocolRestJsonWithCustomPackageEndpointParams; import foo.bar.helloworld.endpoints.ProtocolRestJsonWithCustomPackageEndpointProvider; import foo.bar.helloworld.endpoints.internal.ProtocolRestJsonWithCustomPackageEndpointResolverUtils; @@ -18,6 +19,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -92,6 +94,8 @@ final class DefaultProtocolRestJsonWithCustomPackageAsyncClient implements Proto } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultProtocolRestJsonWithCustomPackageAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration @@ -218,10 +222,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(ProtocolRestJsonWithCustomPackageAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of ProtocolRestJsonWithCustomPackageAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultProtocolRestJsonWithCustomPackageAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } ProtocolRestJsonWithCustomPackageAuthSchemeParams.Builder paramsBuilder = ProtocolRestJsonWithCustomPackageAuthSchemeParams .builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-sync.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-sync.java index 51b1e35f934c..dfe6b88dfa19 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-sync.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-custompackage-sync.java @@ -2,6 +2,7 @@ import foo.bar.helloworld.auth.scheme.ProtocolRestJsonWithCustomPackageAuthSchemeParams; import foo.bar.helloworld.auth.scheme.ProtocolRestJsonWithCustomPackageAuthSchemeProvider; +import foo.bar.helloworld.auth.scheme.internal.DefaultProtocolRestJsonWithCustomPackageAuthSchemeProvider; import foo.bar.helloworld.endpoints.ProtocolRestJsonWithCustomPackageEndpointParams; import foo.bar.helloworld.endpoints.ProtocolRestJsonWithCustomPackageEndpointProvider; import foo.bar.helloworld.endpoints.internal.ProtocolRestJsonWithCustomPackageEndpointResolverUtils; @@ -15,6 +16,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -87,6 +89,8 @@ final class DefaultProtocolRestJsonWithCustomPackageClient implements ProtocolRe } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultProtocolRestJsonWithCustomPackageClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration @@ -186,10 +190,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(ProtocolRestJsonWithCustomPackageAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of ProtocolRestJsonWithCustomPackageAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultProtocolRestJsonWithCustomPackageAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } ProtocolRestJsonWithCustomPackageAuthSchemeParams.Builder paramsBuilder = ProtocolRestJsonWithCustomPackageAuthSchemeParams .builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-async.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-async.java index 55fb3230f5d4..e6b022bfafd7 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-async.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-async.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -51,6 +52,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.auth.scheme.ProtocolRestJsonWithCustomContentTypeAuthSchemeParams; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.auth.scheme.ProtocolRestJsonWithCustomContentTypeAuthSchemeProvider; +import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.auth.scheme.internal.DefaultProtocolRestJsonWithCustomContentTypeAuthSchemeProvider; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.endpoints.ProtocolRestJsonWithCustomContentTypeEndpointParams; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.endpoints.ProtocolRestJsonWithCustomContentTypeEndpointProvider; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.endpoints.internal.ProtocolRestJsonWithCustomContentTypeEndpointResolverUtils; @@ -92,6 +94,8 @@ final class DefaultProtocolRestJsonWithCustomContentTypeAsyncClient implements P } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultProtocolRestJsonWithCustomContentTypeAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration @@ -218,10 +222,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(ProtocolRestJsonWithCustomContentTypeAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of ProtocolRestJsonWithCustomContentTypeAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultProtocolRestJsonWithCustomContentTypeAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } ProtocolRestJsonWithCustomContentTypeAuthSchemeParams.Builder paramsBuilder = ProtocolRestJsonWithCustomContentTypeAuthSchemeParams .builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-sync.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-sync.java index 2d85e1a94838..f08e1ab354c7 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-sync.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-customservicemetadata-sync.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -46,6 +47,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.auth.scheme.ProtocolRestJsonWithCustomContentTypeAuthSchemeParams; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.auth.scheme.ProtocolRestJsonWithCustomContentTypeAuthSchemeProvider; +import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.auth.scheme.internal.DefaultProtocolRestJsonWithCustomContentTypeAuthSchemeProvider; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.endpoints.ProtocolRestJsonWithCustomContentTypeEndpointParams; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.endpoints.ProtocolRestJsonWithCustomContentTypeEndpointProvider; import software.amazon.awssdk.services.protocolrestjsonwithcustomcontenttype.endpoints.internal.ProtocolRestJsonWithCustomContentTypeEndpointResolverUtils; @@ -87,6 +89,8 @@ final class DefaultProtocolRestJsonWithCustomContentTypeClient implements Protoc } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultProtocolRestJsonWithCustomContentTypeClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration @@ -186,10 +190,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(ProtocolRestJsonWithCustomContentTypeAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of ProtocolRestJsonWithCustomContentTypeAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultProtocolRestJsonWithCustomContentTypeAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } ProtocolRestJsonWithCustomContentTypeAuthSchemeParams.Builder paramsBuilder = ProtocolRestJsonWithCustomContentTypeAuthSchemeParams .builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-async.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-async.java index 83e264696992..4b7c31e64259 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-async.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-async.java @@ -8,6 +8,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -57,6 +58,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.endpointdiscoverytest.auth.scheme.EndpointDiscoveryTestAuthSchemeParams; import software.amazon.awssdk.services.endpointdiscoverytest.auth.scheme.EndpointDiscoveryTestAuthSchemeProvider; +import software.amazon.awssdk.services.endpointdiscoverytest.auth.scheme.internal.DefaultEndpointDiscoveryTestAuthSchemeProvider; import software.amazon.awssdk.services.endpointdiscoverytest.endpoints.EndpointDiscoveryTestEndpointParams; import software.amazon.awssdk.services.endpointdiscoverytest.endpoints.EndpointDiscoveryTestEndpointProvider; import software.amazon.awssdk.services.endpointdiscoverytest.endpoints.internal.EndpointDiscoveryTestEndpointResolverUtils; @@ -109,6 +111,8 @@ final class DefaultEndpointDiscoveryTestAsyncClient implements EndpointDiscovery private EndpointDiscoveryRefreshCache endpointDiscoveryCache; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultEndpointDiscoveryTestAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -479,10 +483,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec : Validate.isInstanceOf(EndpointDiscoveryTestAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of EndpointDiscoveryTestAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultEndpointDiscoveryTestAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } EndpointDiscoveryTestAuthSchemeParams.Builder paramsBuilder = EndpointDiscoveryTestAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-sync.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-sync.java index a4089f26d762..6d6de8fdad2e 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-sync.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-endpoint-discovery-sync.java @@ -6,6 +6,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -53,6 +54,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.endpointdiscoverytest.auth.scheme.EndpointDiscoveryTestAuthSchemeParams; import software.amazon.awssdk.services.endpointdiscoverytest.auth.scheme.EndpointDiscoveryTestAuthSchemeProvider; +import software.amazon.awssdk.services.endpointdiscoverytest.auth.scheme.internal.DefaultEndpointDiscoveryTestAuthSchemeProvider; import software.amazon.awssdk.services.endpointdiscoverytest.endpoints.EndpointDiscoveryTestEndpointParams; import software.amazon.awssdk.services.endpointdiscoverytest.endpoints.EndpointDiscoveryTestEndpointProvider; import software.amazon.awssdk.services.endpointdiscoverytest.endpoints.internal.EndpointDiscoveryTestEndpointResolverUtils; @@ -104,8 +106,11 @@ final class DefaultEndpointDiscoveryTestClient implements EndpointDiscoveryTestC } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + private EndpointDiscoveryRefreshCache endpointDiscoveryCache; + protected DefaultEndpointDiscoveryTestClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -399,10 +404,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec : Validate.isInstanceOf(EndpointDiscoveryTestAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of EndpointDiscoveryTestAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultEndpointDiscoveryTestAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } EndpointDiscoveryTestAuthSchemeParams.Builder paramsBuilder = EndpointDiscoveryTestAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-async-client-class.java index 79e3560406ce..4dc65befa828 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-async-client-class.java @@ -8,6 +8,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; import java.util.function.Consumer; @@ -75,6 +76,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider; +import software.amazon.awssdk.services.json.auth.scheme.internal.DefaultJsonAuthSchemeProvider; import software.amazon.awssdk.services.json.batchmanager.JsonAsyncBatchManager; import software.amazon.awssdk.services.json.endpoints.JsonEndpointParams; import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider; @@ -181,8 +183,11 @@ final class DefaultJsonAsyncClient implements JsonAsyncClient { private final ScheduledExecutorService executorService; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + private final Executor executor; + protected DefaultJsonAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -1458,9 +1463,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec JsonAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(JsonAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of JsonAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultJsonAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-client-class.java index bc8b00d106c4..e3614218d546 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-json-client-class.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -54,6 +55,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider; +import software.amazon.awssdk.services.json.auth.scheme.internal.DefaultJsonAuthSchemeProvider; import software.amazon.awssdk.services.json.endpoints.JsonEndpointParams; import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider; import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointResolverUtils; @@ -136,6 +138,8 @@ final class DefaultJsonClient implements JsonClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultJsonClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -979,9 +983,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec .isInstanceOf(JsonAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of JsonAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultJsonAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-presignedurl-async.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-presignedurl-async.java index 6e67445d26ea..b145e4c312e7 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-presignedurl-async.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-presignedurl-async.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -51,6 +52,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams; import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider; +import software.amazon.awssdk.services.json.auth.scheme.internal.DefaultJsonAuthSchemeProvider; import software.amazon.awssdk.services.json.endpoints.JsonEndpointParams; import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider; import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointResolverUtils; @@ -98,6 +100,8 @@ final class DefaultJsonAsyncClient implements JsonAsyncClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultJsonAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -228,9 +232,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, JsonAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(JsonAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of JsonAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultJsonAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-async-client-class.java index a2cb17b35fc6..1db3165fd3d6 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-async-client-class.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ScheduledExecutorService; import java.util.function.Consumer; import org.slf4j.Logger; @@ -58,6 +59,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeParams; import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeProvider; +import software.amazon.awssdk.services.query.auth.scheme.internal.DefaultQueryAuthSchemeProvider; import software.amazon.awssdk.services.query.endpoints.QueryEndpointParams; import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider; import software.amazon.awssdk.services.query.endpoints.internal.QueryEndpointResolverUtils; @@ -139,6 +141,8 @@ final class DefaultQueryAsyncClient implements QueryAsyncClient { private final ScheduledExecutorService executorService; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultQueryAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -1244,9 +1248,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, QueryAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(QueryAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of QueryAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultQueryAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } QueryAuthSchemeParams.Builder paramsBuilder = QueryAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-client-class.java index c48c63f2afd3..e2805f5e9075 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-query-client-class.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import software.amazon.awssdk.annotations.Generated; import software.amazon.awssdk.annotations.SdkInternalApi; @@ -51,6 +52,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeParams; import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeProvider; +import software.amazon.awssdk.services.query.auth.scheme.internal.DefaultQueryAuthSchemeProvider; import software.amazon.awssdk.services.query.endpoints.QueryEndpointParams; import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider; import software.amazon.awssdk.services.query.endpoints.internal.QueryEndpointResolverUtils; @@ -129,6 +131,8 @@ final class DefaultQueryClient implements QueryClient { private final SdkClientConfiguration clientConfiguration; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultQueryClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -1062,9 +1066,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec .isInstanceOf(QueryAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of QueryAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultQueryAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } QueryAuthSchemeParams.Builder paramsBuilder = QueryAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-async-client-class.java index 6d14da528d9f..61690e57d4ed 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-async-client-class.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -51,6 +52,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.smithyrpcv2protocol.auth.scheme.SmithyRpcV2ProtocolAuthSchemeParams; import software.amazon.awssdk.services.smithyrpcv2protocol.auth.scheme.SmithyRpcV2ProtocolAuthSchemeProvider; +import software.amazon.awssdk.services.smithyrpcv2protocol.auth.scheme.internal.DefaultSmithyRpcV2ProtocolAuthSchemeProvider; import software.amazon.awssdk.services.smithyrpcv2protocol.endpoints.SmithyRpcV2ProtocolEndpointParams; import software.amazon.awssdk.services.smithyrpcv2protocol.endpoints.SmithyRpcV2ProtocolEndpointProvider; import software.amazon.awssdk.services.smithyrpcv2protocol.endpoints.internal.SmithyRpcV2ProtocolEndpointResolverUtils; @@ -140,6 +142,8 @@ final class DefaultSmithyRpcV2ProtocolAsyncClient implements SmithyRpcV2Protocol } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultSmithyRpcV2ProtocolAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -945,10 +949,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(SmithyRpcV2ProtocolAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of SmithyRpcV2ProtocolAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultSmithyRpcV2ProtocolAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } SmithyRpcV2ProtocolAuthSchemeParams.Builder paramsBuilder = SmithyRpcV2ProtocolAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-sync.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-sync.java index 9ad62fa3a88e..a19916101f43 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-sync.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-sync.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -46,6 +47,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.smithyrpcv2protocol.auth.scheme.SmithyRpcV2ProtocolAuthSchemeParams; import software.amazon.awssdk.services.smithyrpcv2protocol.auth.scheme.SmithyRpcV2ProtocolAuthSchemeProvider; +import software.amazon.awssdk.services.smithyrpcv2protocol.auth.scheme.internal.DefaultSmithyRpcV2ProtocolAuthSchemeProvider; import software.amazon.awssdk.services.smithyrpcv2protocol.endpoints.SmithyRpcV2ProtocolEndpointParams; import software.amazon.awssdk.services.smithyrpcv2protocol.endpoints.SmithyRpcV2ProtocolEndpointProvider; import software.amazon.awssdk.services.smithyrpcv2protocol.endpoints.internal.SmithyRpcV2ProtocolEndpointResolverUtils; @@ -135,6 +137,8 @@ final class DefaultSmithyRpcV2ProtocolClient implements SmithyRpcV2ProtocolClien } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultSmithyRpcV2ProtocolClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -824,10 +828,23 @@ private List resolveAuthSchemeOptions(SdkRequest request, : Validate.isInstanceOf(SmithyRpcV2ProtocolAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of SmithyRpcV2ProtocolAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultSmithyRpcV2ProtocolAuthSchemeProvider; + String cacheKey = String.valueOf(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } SmithyRpcV2ProtocolAuthSchemeParams.Builder paramsBuilder = SmithyRpcV2ProtocolAuthSchemeParams.builder().operation( operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-async-client-class.java index f898f9da92fe..ef3bdd024982 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-async-client-class.java @@ -8,6 +8,7 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import org.slf4j.Logger; @@ -57,6 +58,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeParams; import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeProvider; +import software.amazon.awssdk.services.database.auth.scheme.internal.DefaultDatabaseAuthSchemeProvider; import software.amazon.awssdk.services.database.endpoints.DatabaseEndpointParams; import software.amazon.awssdk.services.database.endpoints.DatabaseEndpointProvider; import software.amazon.awssdk.services.database.endpoints.internal.DatabaseEndpointResolverUtils; @@ -133,6 +135,8 @@ final class DefaultDatabaseAsyncClient implements DatabaseAsyncClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultDatabaseAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -905,6 +909,16 @@ private List resolveAuthSchemeOptions(SdkRequest request, DatabaseAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(DatabaseAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of DatabaseAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultDatabaseAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION) + ":" + + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_SIGV4A_SIGNING_REGION_SET); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } DatabaseAuthSchemeParams.Builder paramsBuilder = DatabaseAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); Set sigv4aRegionSet = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_SIGV4A_SIGNING_REGION_SET); @@ -912,6 +926,10 @@ private List resolveAuthSchemeOptions(SdkRequest request, paramsBuilder.regionSet(RegionSet.create(sigv4aRegionSet)); } List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-sync-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-sync-client-class.java index 951dc25368e7..3d79abaea0da 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-sync-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-unsigned-payload-trait-sync-client-class.java @@ -5,6 +5,7 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.Generated; @@ -52,6 +53,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeParams; import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeProvider; +import software.amazon.awssdk.services.database.auth.scheme.internal.DefaultDatabaseAuthSchemeProvider; import software.amazon.awssdk.services.database.endpoints.DatabaseEndpointParams; import software.amazon.awssdk.services.database.endpoints.DatabaseEndpointProvider; import software.amazon.awssdk.services.database.endpoints.internal.DatabaseEndpointResolverUtils; @@ -128,6 +130,8 @@ final class DefaultDatabaseClient implements DatabaseClient { } }; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultDatabaseClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -783,6 +787,16 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec .isInstanceOf(DatabaseAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of DatabaseAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultDatabaseAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION) + ":" + + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_SIGV4A_SIGNING_REGION_SET); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } DatabaseAuthSchemeParams.Builder paramsBuilder = DatabaseAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); Set sigv4aRegionSet = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_SIGV4A_SIGNING_REGION_SET); @@ -790,6 +804,10 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec paramsBuilder.regionSet(RegionSet.create(sigv4aRegionSet)); } List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-async-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-async-client-class.java index 83cc6500a7e7..29a4cf905133 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-async-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-async-client-class.java @@ -7,6 +7,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.function.Consumer; import org.slf4j.Logger; @@ -64,6 +65,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.xml.auth.scheme.XmlAuthSchemeParams; import software.amazon.awssdk.services.xml.auth.scheme.XmlAuthSchemeProvider; +import software.amazon.awssdk.services.xml.auth.scheme.internal.DefaultXmlAuthSchemeProvider; import software.amazon.awssdk.services.xml.endpoints.XmlEndpointParams; import software.amazon.awssdk.services.xml.endpoints.XmlEndpointProvider; import software.amazon.awssdk.services.xml.endpoints.internal.XmlEndpointResolverUtils; @@ -129,8 +131,11 @@ final class DefaultXmlAsyncClient implements XmlAsyncClient { private final SdkClientConfiguration clientConfiguration; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + private final Executor executor; + protected DefaultXmlAsyncClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsAsyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -953,9 +958,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, XmlAuthSchemeProvider authSchemeProvider = requestAuthSchemeProvider != null ? requestAuthSchemeProvider : Validate .isInstanceOf(XmlAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of XmlAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultXmlAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } XmlAuthSchemeParams.Builder paramsBuilder = XmlAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-client-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-client-class.java index e7799d93c4ca..d19470189ed0 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-client-class.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-xml-client-class.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Consumer; import software.amazon.awssdk.annotations.Generated; import software.amazon.awssdk.annotations.SdkInternalApi; @@ -52,6 +53,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.services.xml.auth.scheme.XmlAuthSchemeParams; import software.amazon.awssdk.services.xml.auth.scheme.XmlAuthSchemeProvider; +import software.amazon.awssdk.services.xml.auth.scheme.internal.DefaultXmlAuthSchemeProvider; import software.amazon.awssdk.services.xml.endpoints.XmlEndpointParams; import software.amazon.awssdk.services.xml.endpoints.XmlEndpointProvider; import software.amazon.awssdk.services.xml.endpoints.internal.XmlEndpointResolverUtils; @@ -111,6 +113,8 @@ final class DefaultXmlClient implements XmlClient { private final SdkClientConfiguration clientConfiguration; + private final ConcurrentHashMap> authSchemeCache = new ConcurrentHashMap<>(); + protected DefaultXmlClient(SdkClientConfiguration clientConfiguration) { this.clientHandler = new AwsSyncClientHandler(clientConfiguration); this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this) @@ -722,9 +726,22 @@ private List resolveAuthSchemeOptions(SdkRequest request, Exec .isInstanceOf(XmlAuthSchemeProvider.class, executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER), "Expected an instance of XmlAuthSchemeProvider"); + boolean useCache = requestAuthSchemeProvider == null + && authSchemeProvider instanceof DefaultXmlAuthSchemeProvider; + String cacheKey = operationName + ":" + executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION); + if (useCache) { + List cached = authSchemeCache.get(cacheKey); + if (cached != null) { + return cached; + } + } XmlAuthSchemeParams.Builder paramsBuilder = XmlAuthSchemeParams.builder().operation(operationName); paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION)); List options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build()); + if (useCache) { + options = Collections.unmodifiableList(options); + authSchemeCache.put(cacheKey, options); + } return options; } diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/auth/AuthSchemeCacheTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/auth/AuthSchemeCacheTest.java new file mode 100644 index 000000000000..6474ad1eb101 --- /dev/null +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/auth/AuthSchemeCacheTest.java @@ -0,0 +1,270 @@ +/* + * 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.services.auth; + +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.Mockito.when; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.awscore.AwsExecutionAttribute; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.http.SdkHttpClient; +import software.amazon.awssdk.http.auth.aws.scheme.AwsV4AuthScheme; +import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.multiauth.MultiauthClient; +import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient; +import software.amazon.awssdk.services.protocolrestjson.auth.scheme.ProtocolRestJsonAuthSchemeParams; +import software.amazon.awssdk.services.protocolrestjson.auth.scheme.ProtocolRestJsonAuthSchemeProvider; + +class AuthSchemeCacheTest { + + @Mock + private SdkHttpClient mockHttpClient; + + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + when(mockHttpClient.clientName()).thenReturn("MockHttpClient"); + when(mockHttpClient.prepareRequest(any())) + .thenThrow(new RuntimeException("stop")); + } + + @Test + void customClientProvider_isNotCached() { + AtomicInteger resolveCount = new AtomicInteger(0); + + ProtocolRestJsonAuthSchemeProvider customProvider = params -> { + resolveCount.incrementAndGet(); + return Collections.singletonList( + AuthSchemeOption.builder().schemeId(AwsV4AuthScheme.SCHEME_ID).build()); + }; + + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .authSchemeProvider(customProvider) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(resolveCount.get()).isEqualTo(1); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(resolveCount.get()).isEqualTo(2); + + client.close(); + } + + @Test + void perRequestOverride_usesOverrideProvider() { + AtomicInteger requestProviderCount = new AtomicInteger(0); + + ProtocolRestJsonAuthSchemeProvider requestProvider = params -> { + requestProviderCount.incrementAndGet(); + return Collections.singletonList( + AuthSchemeOption.builder().schemeId(AwsV4AuthScheme.SCHEME_ID).build()); + }; + + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(requestProviderCount.get()).isEqualTo(0); + + assertThatThrownBy(() -> client.allTypes(r -> r.overrideConfiguration( + c -> c.authSchemeProvider(requestProvider) + ))).hasMessageContaining("stop"); + assertThat(requestProviderCount.get()).isEqualTo(1); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(requestProviderCount.get()).isEqualTo(1); + + client.close(); + } + + @Test + void defaultProvider_multipleOperationsSucceed() { + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThatThrownBy(() -> client.deleteOperation(r -> {})).hasMessageContaining("stop"); + assertThatThrownBy(() -> client.deleteOperation(r -> {})).hasMessageContaining("stop"); + + client.close(); + } + + @Test + void defaultProvider_returnsUnmodifiableList() { + List options = ProtocolRestJsonAuthSchemeProvider.defaultProvider().resolveAuthScheme( + ProtocolRestJsonAuthSchemeParams.builder().operation("AllTypes").region(Region.US_WEST_2).build()); + + assertThat(options).isNotEmpty(); + assertThatThrownBy(() -> options.add(AuthSchemeOption.builder().schemeId(AwsV4AuthScheme.SCHEME_ID).build())) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + void defaultProvider_cachesUnmodifiableListAndReusesInstance() throws Exception { + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + + Map> cache = authSchemeCache(client); + assertThat(cache).isNotEmpty(); + List cached = cache.values().iterator().next(); + + assertThatThrownBy(() -> cached.add(AuthSchemeOption.builder().schemeId(AwsV4AuthScheme.SCHEME_ID).build())) + .isInstanceOf(UnsupportedOperationException.class); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(cache.values().iterator().next()).isSameAs(cached); + + client.close(); + } + + @Test + void perOpAuthService_differentOperationsSeparateCacheEntries() throws Exception { + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + Map> cache = authSchemeCache(client); + assertThat(cache).hasSize(1); + + assertThatThrownBy(() -> client.deleteOperation(r -> {})).hasMessageContaining("stop"); + assertThat(cache).hasSize(2); + + client.close(); + } + + + @Test + void regionChange_causesCacheMiss() throws Exception { + AtomicInteger callCount = new AtomicInteger(0); + ExecutionInterceptor regionSwitcher = new ExecutionInterceptor() { + @Override + public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { + if (callCount.incrementAndGet() > 1) { + executionAttributes.putAttribute(AwsExecutionAttribute.AWS_REGION, Region.EU_WEST_1); + } + } + }; + + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .overrideConfiguration(c -> c.addExecutionInterceptor(regionSwitcher)) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + Map> cache = authSchemeCache(client); + assertThat(cache).hasSize(1); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(cache).hasSize(2); + + client.close(); + } + + @Test + void cacheHit_returnsSameInstanceOnRepeatedCalls() throws Exception { + ProtocolRestJsonClient client = ProtocolRestJsonClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .build(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + Map> cache = authSchemeCache(client); + assertThat(cache).hasSize(1); + List firstResult = cache.values().iterator().next(); + + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThatThrownBy(() -> client.allTypes(r -> {})).hasMessageContaining("stop"); + assertThat(cache).hasSize(1); + assertThat(cache.values().iterator().next()).isSameAs(firstResult); + + client.close(); + } + + @Test + void regionSetChange_causesCacheMiss() throws Exception { + AtomicInteger callCount = new AtomicInteger(0); + ExecutionInterceptor regionSetSwitcher = new ExecutionInterceptor() { + @Override + public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) { + if (callCount.incrementAndGet() > 1) { + executionAttributes.putAttribute(AwsExecutionAttribute.AWS_SIGV4A_SIGNING_REGION_SET, + Collections.singleton("eu-west-1")); + } + } + }; + + MultiauthClient client = MultiauthClient.builder() + .httpClient(mockHttpClient) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid"))) + .region(Region.US_WEST_2) + .overrideConfiguration(c -> c.addExecutionInterceptor(regionSetSwitcher)) + .build(); + + assertThatThrownBy(() -> client.multiAuthWithOnlySigv4a(r -> {})).hasMessageContaining("stop"); + Map> cache = authSchemeCache(client); + assertThat(cache).hasSize(1); + + assertThatThrownBy(() -> client.multiAuthWithOnlySigv4a(r -> {})).hasMessageContaining("stop"); + assertThat(cache).hasSize(2); + + client.close(); + } + + @SuppressWarnings("unchecked") + private static Map> authSchemeCache(Object client) throws Exception { + Field field = client.getClass().getDeclaredField("authSchemeCache"); + field.setAccessible(true); + return (Map>) field.get(client); + } +}