Cache auth scheme resolution

.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"
+}

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<String> setOf(String val1, String val2) {
         Set<String> result = new HashSet<>();
         result.add(val1);

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

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<FieldSpec> 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<CodeBlock> 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();

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

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<String> paramsAttributes = extractAttributes(paramsSection);
+        assertThat(paramsAttributes).as("Expected auth params to reference AwsExecutionAttributes").isNotEmpty();
+        String cacheKeySection = source.substring(cacheKeyStart, cacheKeyEnd);
+        Set<String> 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<String> extractAttributes(String section) {
+        Set<String> attributes = new HashSet<>();
+        Matcher matcher = AWS_EXEC_ATTR_PATTERN.matcher(section);
+        while (matcher.find()) {
+            attributes.add(matcher.group(1));
+        }
+        return attributes;
+    }
+}

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<String, List<AuthSchemeOption>> 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<AuthSchemeOption> 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<AuthSchemeOption> cached = authSchemeCache.get(cacheKey);
+            if (cached != null) {
+                return cached;
+            }
+        }
         JsonAuthSchemeParams.Builder paramsBuilder = JsonAuthSchemeParams.builder().operation(operationName);
         paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION));
         List<AuthSchemeOption> options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build());
+        if (useCache) {
+            options = Collections.unmodifiableList(options);
+            authSchemeCache.put(cacheKey, options);
+        }
         return options;
     }
 

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<String, List<AuthSchemeOption>> 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<AuthSchemeOption> 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<AuthSchemeOption> cached = authSchemeCache.get(cacheKey);
+            if (cached != null) {
+                return cached;
+            }
+        }
         QueryToJsonCompatibleAuthSchemeParams.Builder paramsBuilder = QueryToJsonCompatibleAuthSchemeParams.builder().operation(
             operationName);
         paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION));
         List<AuthSchemeOption> options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build());
+        if (useCache) {
+            options = Collections.unmodifiableList(options);
+            authSchemeCache.put(cacheKey, options);
+        }
         return options;
     }