Skip to content

perf(endpoints): Add peephole optimizations to endpoint resolution codegen #7305

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

Draft
alextwoods wants to merge 1 commit into from
Draft

perf(endpoints): Add peephole optimizations to endpoint resolution codegen #7305

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
61 changes: 55 additions & 6 deletions codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules2/CodeGeneratorVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,15 +328,64 @@ private String callParams(String ruleId) {

@Override
public Void visitEndpointExpression(EndpointExpression e) {
builder.add("return $T.builder().endpointUrl(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
builder.add(")");
e.headers().accept(this);
e.properties().accept(this);
builder.addStatement(".build()");
Map<String, RuleExpression> properties = e.properties().properties();
boolean hasHeaders = !e.headers().headers().isEmpty();
boolean hasAuthSchemesOnly = !hasHeaders && properties.size() == 1 && properties.containsKey("authSchemes");
boolean hasTwoAttrs = !hasHeaders && properties.size() == 2 && properties.containsKey("authSchemes");

if (hasAuthSchemesOnly) {
// Optimized: Endpoint.ofAttribute(url, AUTH_SCHEMES, list)
builder.add("return $T.ofAttribute(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
builder.add(", $T.AUTH_SCHEMES, ", AwsEndpointAttribute.class);
addAuthSchemesInlineValue(properties.get("authSchemes"));
builder.addStatement(")");
} else if (hasTwoAttrs) {
// Optimized: Endpoint.ofAttributes(url, key1, val1, key2, val2)
builder.add("return $T.ofAttributes(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
for (Map.Entry<String, RuleExpression> entry : properties.entrySet()) {
builder.add(", ");
if ("authSchemes".equals(entry.getKey())) {
builder.add("$T.AUTH_SCHEMES, ", AwsEndpointAttribute.class);
addAuthSchemesInlineValue(entry.getValue());
} else if (knownEndpointAttributes.containsKey(entry.getKey())) {
KeyTypePair keyType = knownEndpointAttributes.get(entry.getKey());
ClassConstant classConstant = parseClassConstant(keyType.getKey());
builder.add("$T.$L, ", classConstant.className(), classConstant.fieldName());
entry.getValue().accept(this);
} else {
builder.add("$T.AUTH_SCHEMES, ", AwsEndpointAttribute.class);
entry.getValue().accept(this);
}
}
builder.addStatement(")");
} else {
// General case: use builder pattern
builder.add("return $T.builder().endpointUrl(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
builder.add(")");
e.headers().accept(this);
e.properties().accept(this);
builder.addStatement(".build()");
}
return null;
}

private void addAuthSchemesInlineValue(RuleExpression authSchemesExpr) {
ListExpression expr = (ListExpression) authSchemesExpr;
builder.add("$T.asList(", Arrays.class);
boolean isFirst = true;
for (RuleExpression authSchemeExpr : expr.expressions()) {
if (!isFirst) {
builder.add(", ");
}
addAuthSchemesBody(authSchemeExpr);
isFirst = false;
}
builder.add(")");
}

@Override
public Void visitPropertiesExpression(PropertiesExpression e) {
Map<String, RuleExpression> properties = e.properties();
Expand Down
63 changes: 56 additions & 7 deletions ...in/java/software/amazon/awssdk/codegen/poet/rules2/bdd/BddResultCodeGeneratorVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,16 +252,65 @@ public RuleType visitListExpression(ListExpression e) {

@Override
public RuleType visitEndpointExpression(EndpointExpression e) {
// Return Endpoint directly — no RuleResult wrapper
builder.add("return $T.builder().endpointUrl(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
builder.add(")");
e.headers().accept(this);
e.properties().accept(this);
builder.addStatement(".build()");
Map<String, RuleExpression> properties = e.properties().properties();
boolean hasHeaders = !e.headers().headers().isEmpty();
boolean hasAuthSchemesOnly = !hasHeaders && properties.size() == 1 && properties.containsKey("authSchemes");
boolean hasTwoAttrs = !hasHeaders && properties.size() == 2 && properties.containsKey("authSchemes");

if (hasAuthSchemesOnly) {
// Optimized: Endpoint.ofAttribute(url, AUTH_SCHEMES, list)
builder.add("return $T.ofAttribute(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
builder.add(", $T.AUTH_SCHEMES, ", AwsEndpointAttribute.class);
addAuthSchemesValue(properties.get("authSchemes"));
builder.addStatement(")");
} else if (hasTwoAttrs) {
// Optimized: Endpoint.ofAttributes(url, key1, val1, key2, val2)
builder.add("return $T.ofAttributes(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
for (Map.Entry<String, RuleExpression> entry : properties.entrySet()) {
builder.add(", ");
if ("authSchemes".equals(entry.getKey())) {
builder.add("$T.AUTH_SCHEMES, ", AwsEndpointAttribute.class);
addAuthSchemesValue(entry.getValue());
} else if (knownEndpointAttributes.containsKey(entry.getKey())) {
KeyTypePair keyType = knownEndpointAttributes.get(entry.getKey());
ClassConstant classConstant = parseClassConstant(keyType.getKey());
builder.add("$T.$L, ", classConstant.className(), classConstant.fieldName());
entry.getValue().accept(this);
} else {
// Unknown attribute — use AwsEndpointAttribute for safety
builder.add("$T.AUTH_SCHEMES, ", AwsEndpointAttribute.class);
entry.getValue().accept(this);
}
}
builder.addStatement(")");
} else {
// General case: use builder pattern
builder.add("return $T.builder().endpointUrl(", Endpoint.class);
EndpointUrlCodeEmitter.emit(e.url(), builder, this);
builder.add(")");
e.headers().accept(this);
e.properties().accept(this);
builder.addStatement(".build()");
}
return null;
}

private void addAuthSchemesValue(RuleExpression authSchemesExpr) {
ListExpression expr = (ListExpression) authSchemesExpr;
builder.add("$T.asList(", Arrays.class);
boolean isFirst = true;
for (RuleExpression authSchemeExpr : expr.expressions()) {
if (!isFirst) {
builder.add(", ");
}
addAuthSchemesBody(authSchemeExpr);
isFirst = false;
}
builder.add(")");
}

@Override
public RuleType visitPropertiesExpression(PropertiesExpression e) {
Map<String, RuleExpression> properties = e.properties();
Expand Down
25 changes: 24 additions & 1 deletion ...gen/src/main/resources/software/amazon/awssdk/codegen/rules2/RulesFunctions.java.resource
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,21 @@ public class RulesFunctions {
private static final int MAX_HOST_LABEL_SIZE = 63;
private static final int MIN_BUCKET_SIZE = 3;

// Per-thread caches to avoid recomputing results for repeated calls with the same inputs.
// These are common in endpoint resolution where the same region/bucket is resolved many times.
private static final ThreadLocal<FunctionCache> CACHE = new ThreadLocal<FunctionCache>() {
@Override
protected FunctionCache initialValue() {
return new FunctionCache();
}
};

private static final class FunctionCache {
// awsPartition cache: (lastRegion -> lastPartition)
String lastPartitionRegion;
RulePartition lastPartitionResult;
}

@SafeVarargs
public static <T> T coalesce(T... args) {
if (args == null || args.length < 2) {
Expand Down Expand Up @@ -191,6 +206,11 @@ public class RulesFunctions {
// AWS related functions

public static RulePartition awsPartition(String regionName) {
FunctionCache cache = CACHE.get();
if (regionName != null && regionName.equals(cache.lastPartitionRegion)) {
return cache.lastPartitionResult;
}

PartitionData data = PARTITION_DATA.value();
Partition matchedPartition;

Expand All @@ -211,7 +231,10 @@ public class RulesFunctions {
matchedPartition = AWS_PARTITION.value();
}

return RulePartition.from(matchedPartition.id(), matchedPartition.outputs());
RulePartition result = RulePartition.from(matchedPartition.id(), matchedPartition.outputs());
cache.lastPartitionRegion = regionName;
cache.lastPartitionResult = result;
return result;
}

public static RuleArn awsParseArn(String value) {
Expand Down
12 changes: 6 additions & 6 deletions ...rc/test/resources/software/amazon/awssdk/codegen/poet/rules2/endpoint-provider-class.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ private static Endpoint endpointRule1(QueryEndpointParams params, String region)
if (parsedArn != null) {
String arnResourceId = RulesFunctions.listAccess(parsedArn.resourceId(), 0);
if (arnResourceId != null) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", arnResourceId + "." + params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", arnResourceId + "." + params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
}
}
Expand All @@ -79,9 +79,9 @@ private static Endpoint endpointRule2(QueryEndpointParams params, RulePartition
throw SdkClientException.create("FIPS endpoints not supported with multi-region endpoints");
}
if (params.useFipsEndpoint() == null && params.useDualStackEndpoint() != null && params.useDualStackEndpoint()) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
return null;
}
Expand All @@ -90,13 +90,13 @@ private static Endpoint endpointRule6(QueryEndpointParams params, String region,
RulePartition partitionResult) {
if (RulesFunctions.isValidHostLabel(region, false)) {
if (params.useFipsEndpoint() != null && params.useFipsEndpoint() && params.useDualStackEndpoint() == null) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
if (params.useDualStackEndpoint() != null && params.useDualStackEndpoint() && params.useFipsEndpoint() == null) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query." + region + "." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build(), SigV4AuthScheme.builder().signingName("query").signingRegion(region).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", "query." + region + "." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build(), SigV4AuthScheme.builder().signingName("query").signingRegion(region).build()));
}
if (params.useDualStackEndpoint() != null && params.useFipsEndpoint() != null && params.useDualStackEndpoint() && params.useFipsEndpoint()) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query." + region + "." + partitionResult.dnsSuffix(), -1, "")).build();
}
Expand Down
12 changes: 6 additions & 6 deletions ...oftware/amazon/awssdk/codegen/poet/rules2/endpoint-provider-know-prop-override-class.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ private static Endpoint endpointRule1(QueryEndpointParams params, String region)
if (parsedArn != null) {
String arnResourceId = RulesFunctions.listAccess(parsedArn.resourceId(), 0);
if (arnResourceId != null) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", arnResourceId + "." + params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", arnResourceId + "." + params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
}
}
Expand All @@ -79,9 +79,9 @@ private static Endpoint endpointRule2(QueryEndpointParams params, RulePartition
throw SdkClientException.create("FIPS endpoints not supported with multi-region endpoints");
}
if (params.useFipsEndpoint() == null && params.useDualStackEndpoint() != null && params.useDualStackEndpoint()) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", params.endpointId() + ".query." + partitionResult.dnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
return null;
}
Expand All @@ -90,13 +90,13 @@ private static Endpoint endpointRule6(QueryEndpointParams params, String region,
RulePartition partitionResult) {
if (RulesFunctions.isValidHostLabel(region, false)) {
if (params.useFipsEndpoint() != null && params.useFipsEndpoint() && params.useDualStackEndpoint() == null) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
if (params.useDualStackEndpoint() != null && params.useDualStackEndpoint() && params.useFipsEndpoint() == null) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query." + region + "." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build(), SigV4AuthScheme.builder().signingName("query").signingRegion(region).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", "query." + region + "." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build(), SigV4AuthScheme.builder().signingName("query").signingRegion(region).build()));
}
if (params.useDualStackEndpoint() != null && params.useFipsEndpoint() != null && params.useDualStackEndpoint() && params.useFipsEndpoint()) {
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dualStackDnsSuffix(), -1, "")).putAttribute(AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build())).build();
return Endpoint.ofAttribute(EndpointUrl.fromComponents("https", "query-fips." + region + "." + partitionResult.dualStackDnsSuffix(), -1, ""), AwsEndpointAttribute.AUTH_SCHEMES, Arrays.asList(SigV4aAuthScheme.builder().signingName("query").signingRegionSet(Arrays.asList("*")).build()));
}
return Endpoint.builder().endpointUrl(EndpointUrl.fromComponents("https", "query." + region + "." + partitionResult.dnsSuffix(), -1, "")).build();
}
Expand Down
Loading