Add support for tracking business metrics from resolved endpoints

Author: alextwoods

.changes/next-release/feature-AWSSDKforJavaV2-6e40ae3.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java V2",
+    "contributor": "",
+    "description": "Add support for tracking business metrics from resolved endpoints."
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointResolverInterceptorSpec.java

@@ -75,6 +75,7 @@
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.endpoints.Endpoint;
+import software.amazon.awssdk.endpoints.EndpointAttributeKey;
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.http.auth.aws.scheme.AwsV4AuthScheme;
 import software.amazon.awssdk.http.auth.aws.scheme.AwsV4aAuthScheme;
@@ -160,6 +161,8 @@ public TypeSpec poetSpec() {
         }
 
         endpointParamsKnowledgeIndex.addAccountIdMethodsIfPresent(b);
+
+        b.addMethod(setMetricValuesMethod());
         return b.build();
     }
 
@@ -255,6 +258,7 @@ private MethodSpec modifyRequestMethod(String endpointAuthSchemeStrategyFieldNam
         }
 
         b.addStatement("executionAttributes.putAttribute(SdkInternalExecutionAttribute.RESOLVED_ENDPOINT, endpoint)");
+        b.addStatement("setMetricValues(endpoint, executionAttributes)");
         b.addStatement("return result");
         b.endControlFlow();
         b.beginControlFlow("catch ($T e)", CompletionException.class);
@@ -905,4 +909,19 @@ private MethodSpec constructorMethodSpec(String endpointAuthSchemeFieldName) {
         b.addStatement("this.$N = $N.endpointAuthSchemeStrategy()", endpointAuthSchemeFieldName, factoryLocalVarName);
         return b.build();
     }
+
+    private MethodSpec setMetricValuesMethod() {
+        MethodSpec.Builder b = MethodSpec.methodBuilder("setMetricValues")
+                                         .addModifiers(Modifier.PRIVATE)
+                                         .addParameter(Endpoint.class, "endpoint")
+                                         .addParameter(ExecutionAttributes.class, "executionAttributes")
+                                         .returns(void.class);
+
+        b.beginControlFlow("if (endpoint.attribute($T.METRIC_VALUES) != null)", AwsEndpointAttribute.class);
+        b.addStatement("executionAttributes.getOptionalAttribute($T.BUSINESS_METRICS).ifPresent("
+                       + "metrics -> endpoint.attribute($T.METRIC_VALUES).forEach(v -> metrics.addMetric(v))))",
+                       SdkInternalExecutionAttribute.class, AwsEndpointAttribute.class);
+        b.endControlFlow();
+        return b.build();
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/TestGeneratorUtils.java

@@ -115,6 +115,8 @@ private static void addEndpointAttributeBlock(CodeBlock.Builder builder, String
                                                   Map<String, KeyTypePair> knownEndpointAttributes) {
         if ("authSchemes".equals(attrName)) {
             addAuthSchemesBlock(builder, attrValue);
+        } else if ("metricValues".equals(attrName)) {
+            addMetricValuesBlock(builder, attrValue);
         } else if (knownEndpointAttributes.containsKey(attrName)) {
             addAttributeBlock(builder, attrValue, knownEndpointAttributes.get(attrName));
         } else {
@@ -124,6 +126,29 @@ private static void addEndpointAttributeBlock(CodeBlock.Builder builder, String
         }
     }
 
+    private static void addMetricValuesBlock(CodeBlock.Builder builder, TreeNode attrValue) {
+        CodeBlock keyExpr = CodeBlock.builder()
+                                     .add("$T.METRIC_VALUES", AwsEndpointAttribute.class)
+                                     .build();
+
+        CodeBlock.Builder schemesListExpr = CodeBlock.builder()
+                                                     .add("$T.asList(", Arrays.class);
+
+        JrsArray schemesArray = (JrsArray) attrValue;
+
+        Iterator<JrsValue> elementsIter = schemesArray.elements();
+        while (elementsIter.hasNext()) {
+            schemesListExpr.add("$S", elementsIter.next().asText());
+
+            if (elementsIter.hasNext()) {
+                schemesListExpr.add(",");
+            }
+        }
+        schemesListExpr.add(")");
+
+        builder.add(".putAttribute($L, $L)", keyExpr, schemesListExpr.build());
+    }
+
     private static void addAttributeBlock(CodeBlock.Builder builder, TreeNode attrValue, KeyTypePair keyType) {
         CodeBlock keyExpr = CodeBlock.builder()
                                      .add("$L", keyType.getKey())

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules2/CodeGeneratorVisitor.java

@@ -353,6 +353,8 @@ public Void visitPropertiesExpression(PropertiesExpression e) {
         properties.forEach((k, v) -> {
             if ("authSchemes".equals(k)) {
                 addAuthSchemesBlock(v);
+            } else if ("metricValues".equals(k)) {
+                addMetricValuesBlock(v);
             } else if (knownEndpointAttributes.containsKey(k)) {
                 addAttributeBlock(k, v);
             } else {
@@ -436,6 +438,12 @@ private ClassName authSchemeClass(String name) {
         }
     }
 
+    private void addMetricValuesBlock(RuleExpression v) {
+        builder.add(".putAttribute($T.METRIC_VALUES, ", AwsEndpointAttribute.class);
+        v.accept(this);
+        builder.add(")");
+    }
+
     private void addAttributeBlock(String k, RuleExpression v) {
         KeyTypePair keyType = knownEndpointAttributes.get(k);
         ClassConstant classConstant = parseClassConstant(keyType.getKey());

codegen/src/test/java/software/amazon/awssdk/codegen/poet/ClientTestModels.java

@@ -206,6 +206,26 @@ public static IntermediateModel queryServiceModelsWithUnknownEndpointProperties(
         return new IntermediateModelBuilder(models).build();
     }
 
+    public static IntermediateModel queryServiceModelsWithUnknownEndpointMetricValues() {
+        File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/query/service-2.json").getFile());
+        File waitersModel = new File(ClientTestModels.class.getResource("client/c2j/query/waiters-2.json").getFile());
+        File endpointRuleSetModel =
+            new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-rule-set-metric-values.json").getFile());
+        File endpointTestsModel =
+            new File(ClientTestModels.class.getResource("client/c2j/query/endpoint-tests-metric-values.json").getFile());
+
+        C2jModels models = C2jModels
+            .builder()
+            .serviceModel(getServiceModel(serviceModel))
+            .waitersModel(getWaiters(waitersModel))
+            .customizationConfig(CustomizationConfig.create())
+            .endpointRuleSetModel(getEndpointRuleSet(endpointRuleSetModel))
+            .endpointTestSuiteModel(getEndpointTestSuite(endpointTestsModel))
+            .build();
+
+        return new IntermediateModelBuilder(models).build();
+    }
+
     public static IntermediateModel queryServiceModelsWithUriCache() {
         File serviceModel = new File(ClientTestModels.class.getResource("client/c2j/query/service-2.json").getFile());
         File customizationModel =

codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules/EndpointProviderCompiledRulesClassSpecTest.java

@@ -51,4 +51,11 @@ void endpointProviderClassWithUriCache() {
             new EndpointProviderSpec2(ClientTestModels.queryServiceModelsWithUriCache());
         assertThat(endpointProviderSpec, generatesTo("endpoint-provider-uri-cache-class.java"));
     }
+
+    @Test
+    void endpointProviderClassWithMetricValues() {
+        ClassSpec endpointProviderSpec =
+            new EndpointProviderSpec2(ClientTestModels.queryServiceModelsWithUnknownEndpointMetricValues());
+        assertThat(endpointProviderSpec, generatesTo("endpoint-provider-metric-values-class.java"));
+    }
 }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules/EndpointRulesClientTestSpecTest.java

@@ -40,4 +40,10 @@ public void endpointProviderTestClassWithUnknownProperties() {
         ClassSpec endpointProviderSpec = new EndpointProviderTestSpec(ClientTestModels.queryServiceModelsWithUnknownEndpointProperties());
         assertThat(endpointProviderSpec, generatesTo("endpoint-rules-unknownproperty-test-class.java"));
     }
+
+    @Test
+    public void endpointProviderTestClassWithMetricValues() {
+        ClassSpec endpointProviderSpec = new EndpointProviderTestSpec(ClientTestModels.queryServiceModelsWithUnknownEndpointMetricValues());
+        assertThat(endpointProviderSpec, generatesTo("endpoint-rules-metric-values-test-class.java"));
+    }
 }