Skip to content

Cache auth scheme resolution results per operation #7282

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
davidh44 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Cache auth scheme resolution results per operation #7282

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6dfd73c
Cache auth scheme resolution results per operation
davidh44 Aug 14, 2026
421e679
Update codegenerated test files
davidh44 Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .changes/next-release/feature-AWSSDKforJavav2-301f836.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
{
"type": "feature",
"category": "AWS SDK for Java v2",
"contributor": "",
"description": "Cache auth scheme resolution results per operation"
}
10 changes: 6 additions & 4 deletions codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down Expand Up @@ -165,6 +165,8 @@ protected void addFields(Builder type) {

model.getEndpointOperation().ifPresent(
o -> type.addField(EndpointDiscoveryRefreshCache.class, "endpointDiscoveryCache", PRIVATE));

ClientClassUtils.authSchemeCacheField(authSchemeSpecUtils, endpointRulesSpecUtils).ifPresent(type::addField);
}

@Override
Expand Down
42 changes: 41 additions & 1 deletion codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

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;
Expand Down Expand Up @@ -385,6 +386,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, "
Expand All @@ -393,13 +395,33 @@ static MethodSpec resolveAuthSchemeOptionsMethod(AuthSchemeSpecUtils authSchemeS
SdkInternalExecutionAttribute.class,
"Expected an instance of " + authSchemeSpecUtils.providerInterfaceName().simpleName());

// Use cache for simple (non-endpoint-based, non-S3) auth scheme resolution when using the default provider
boolean canCache = !authSchemeSpecUtils.useEndpointBasedAuthProvider() && !endpointRulesSpecUtils.isS3();
if (canCache) {
ClassName defaultProviderClass = authSchemeSpecUtils.defaultAuthSchemeProviderName();
builder.addStatement("boolean useCache = requestAuthSchemeProvider == null "
+ "&& authSchemeProvider instanceof $T", defaultProviderClass);
builder.beginControlFlow("if (useCache)");
builder.addStatement("$T<$T> cached = authSchemeCache.get(operationName)",
List.class, AuthSchemeOption.class);
builder.beginControlFlow("if (cached != null)");
builder.addStatement("return cached");
builder.endControlFlow();
builder.endControlFlow();
}

if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
addEndpointBasedAuthSchemeResolution(builder, authSchemeSpecUtils, endpointRulesSpecUtils);
} else {
addSimpleAuthSchemeResolution(builder, authSchemeSpecUtils);
}

if (endpointRulesSpecUtils.isS3()) {
if (canCache) {
builder.beginControlFlow("if (useCache)");
builder.addStatement("authSchemeCache.put(operationName, 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);
Expand All @@ -414,6 +436,24 @@ 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,
EndpointRulesSpecUtils endpointRulesSpecUtils) {
if (authSchemeSpecUtils.useEndpointBasedAuthProvider() || endpointRulesSpecUtils.isS3()) {
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 addSimpleAuthSchemeResolution(MethodSpec.Builder builder,
AuthSchemeSpecUtils authSchemeSpecUtils) {
ClassName paramsInterface = authSchemeSpecUtils.parametersInterfaceName();
Expand Down
1 change: 1 addition & 0 deletions codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, endpointRulesSpecUtils).ifPresent(type::addField);
}

@Override
Expand Down
16 changes: 16 additions & 0 deletions ...esources/software/amazon/awssdk/codegen/poet/client/test-aws-json-async-client-class.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -1262,9 +1267,20 @@ 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;
if (useCache) {
List<AuthSchemeOption> cached = authSchemeCache.get(operationName);
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) {
authSchemeCache.put(operationName, options);
}
return options;
}

Expand Down
15 changes: 15 additions & 0 deletions .../amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-async-client-class.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -228,10 +232,21 @@ 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;
if (useCache) {
List<AuthSchemeOption> cached = authSchemeCache.get(operationName);
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) {
authSchemeCache.put(operationName, options);
}
return options;
}

Expand Down
15 changes: 15 additions & 0 deletions ...e/amazon/awssdk/codegen/poet/client/test-aws-query-compatible-json-sync-client-class.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -92,6 +94,8 @@ final class DefaultQueryToJsonCompatibleClient implements QueryToJsonCompatibleC
}
};

private final ConcurrentHashMap<String, List<AuthSchemeOption>> authSchemeCache = new ConcurrentHashMap<>();

protected DefaultQueryToJsonCompatibleClient(SdkClientConfiguration clientConfiguration) {
this.clientHandler = new AwsSyncClientHandler(clientConfiguration);
this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this)
Expand Down Expand Up @@ -196,10 +200,21 @@ 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;
if (useCache) {
List<AuthSchemeOption> cached = authSchemeCache.get(operationName);
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) {
authSchemeCache.put(operationName, options);
}
return options;
}

Expand Down
15 changes: 15 additions & 0 deletions ...rc/test/resources/software/amazon/awssdk/codegen/poet/client/test-batchmanager-async.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -96,6 +98,8 @@ final class DefaultBatchManagerTestAsyncClient implements BatchManagerTestAsyncC

private final ScheduledExecutorService executorService;

private final ConcurrentHashMap<String, List<AuthSchemeOption>> authSchemeCache = new ConcurrentHashMap<>();

protected DefaultBatchManagerTestAsyncClient(SdkClientConfiguration clientConfiguration) {
this.clientHandler = new AwsAsyncClientHandler(clientConfiguration);
this.clientConfiguration = clientConfiguration.toBuilder().option(SdkClientOption.SDK_CLIENT, this)
Expand Down Expand Up @@ -224,10 +228,21 @@ private List<AuthSchemeOption> 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;
if (useCache) {
List<AuthSchemeOption> cached = authSchemeCache.get(operationName);
if (cached != null) {
return cached;
}
}
BatchManagerTestAuthSchemeParams.Builder paramsBuilder = BatchManagerTestAuthSchemeParams.builder().operation(
operationName);
paramsBuilder.region(executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION));
List<AuthSchemeOption> options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build());
if (useCache) {
authSchemeCache.put(operationName, options);
}
return options;
}

Expand Down
Loading
Loading