Async metrics Part1: Add metrics codegen changes for async clients

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java

@@ -32,6 +32,7 @@
 import java.net.URI;
 import java.nio.ByteBuffer;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
@@ -61,11 +62,17 @@
 import software.amazon.awssdk.core.client.handler.AsyncClientHandler;
 import software.amazon.awssdk.core.endpointdiscovery.EndpointDiscoveryRefreshCache;
 import software.amazon.awssdk.core.endpointdiscovery.EndpointDiscoveryRequest;
+import software.amazon.awssdk.core.internal.util.MetricUtils;
+import software.amazon.awssdk.core.metrics.CoreMetric;
+import software.amazon.awssdk.metrics.MetricCollector;
+import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.protocols.json.AwsJsonProtocolFactory;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 import software.amazon.awssdk.utils.FunctionalUtils;
 
 public final class AsyncClientClass extends AsyncClientInterface {
+    private static final String PUBLISHER_NAME = "metricPublisher";
+    private static final String METRIC_COLLECTOR_NAME = "apiCallMetricCollector";
     private final IntermediateModel model;
     private final PoetExtensions poetExtensions;
     private final ClassName className;
@@ -190,9 +197,18 @@ private MethodSpec closeMethod() {
     protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, OperationModel opModel) {
 
         builder.addModifiers(Modifier.PUBLIC)
-               .addAnnotation(Override.class)
-               .beginControlFlow("try")
-               .addCode(ClientClassUtils.callApplySignerOverrideMethod(opModel))
+               .addAnnotation(Override.class);
+
+        builder.addStatement("$1T $2N = $1T.create($3S)",
+                             MetricCollector.class, METRIC_COLLECTOR_NAME, "ApiCall");
+        builder.beginControlFlow("try");
+
+        builder.addStatement("$N.reportMetric($T.$L, $S)", METRIC_COLLECTOR_NAME, CoreMetric.class, "SERVICE_ID",
+                             model.getMetadata().getServiceId());
+        builder.addStatement("$N.reportMetric($T.$L, $S)", METRIC_COLLECTOR_NAME, CoreMetric.class, "OPERATION_NAME",
+                             opModel.getOperationName());
+
+        builder.addCode(ClientClassUtils.callApplySignerOverrideMethod(opModel))
                .addCode(ClientClassUtils.addEndpointTraitCode(opModel))
                .addCode(protocolSpec.responseHandler(model, opModel));
         protocolSpec.errorResponseHandler(opModel).ifPresent(builder::addCode);
@@ -224,8 +240,17 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
                                  "() -> $N.exceptionOccurred(t))", paramName);
         }
 
-        return builder.addStatement("return $T.failedFuture(t)", CompletableFutureUtils.class)
-                      .endControlFlow();
+        builder.addStatement("$T<$T> $N = $T.resolvePublisher(clientConfiguration, $N.overrideConfiguration().orElse(null))",
+                             Optional.class,
+                             MetricPublisher.class,
+                             PUBLISHER_NAME,
+                             MetricUtils.class,
+                             opModel.getInput().getVariableName())
+               .addStatement("$N.ifPresent(p -> p.publish($N.collect()))", PUBLISHER_NAME, "apiCallMetricCollector")
+               .addStatement("return $T.failedFuture(t)", CompletableFutureUtils.class)
+               .endControlFlow();
+
+        return builder;
     }
 
     @Override

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientInterface.java

@@ -47,7 +47,6 @@
 import software.amazon.awssdk.core.SdkClient;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
-import software.amazon.awssdk.metrics.MetricCollector;
 import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
 
 public class AsyncClientInterface implements ClassSpec {
@@ -285,12 +284,6 @@ private MethodSpec traditionalMethod(OperationModel opModel) {
                 .addParameter(requestType, opModel.getInput().getVariableName())
                 .addJavadoc(opModel.getDocs(model, ClientType.ASYNC));
 
-
-        String metricCollectorName = "apiCallMetricCollector";
-
-        builder.addStatement("$1T $2N = $1T.create($3S)",
-                MetricCollector.class, metricCollectorName, "ApiCall");
-
         if (opModel.hasStreamingInput()) {
             builder.addParameter(ClassName.get(AsyncRequestBody.class), "requestBody");
         } else if (opModel.hasEventStreamInput()) {

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

@@ -28,6 +28,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.eventstream.EventStreamAsyncResponseTransformer;
 import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionPojoSupplier;
 import software.amazon.awssdk.awscore.eventstream.RestEventStreamAsyncResponseTransformer;
@@ -266,6 +267,8 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                     asyncResponseTransformerVariable(isStreaming, isRestJson, opModel));
         String whenComplete = whenCompleteBody(opModel, customerResponseHandler);
         if (!whenComplete.isEmpty()) {
+            builder.addStatement("$T requestOverrideConfig = $L.overrideConfiguration().orElse(null)",
+                                 AwsRequestOverrideConfiguration.class, opModel.getInput().getVariableName());
             builder.add("executeFuture$L;", whenComplete);
         }
         if (opModel.hasEventStreamOutput()) {
@@ -326,7 +329,7 @@ private String whenCompleteBody(OperationModel operationModel, String responseHa
             return streamingOutputWhenComplete(responseHandlerName);
         } else {
             // Non streaming can just return the future as is
-            return "";
+            return publishMetricsWhenComplete();
         }
     }
 
@@ -336,6 +339,7 @@ private String whenCompleteBody(OperationModel operationModel, String responseHa
      * {@link EventStreamAsyncResponseTransformer}. Failure is notified via the normal future (the one returned by the client
      * handler).
      *
+     *
      * @param responseHandlerName Variable name of response handler customer passed in.
      * @return whenComplete to append to future.
      */
@@ -347,12 +351,12 @@ private String eventStreamOutputWhenComplete(String responseHandlerName) {
                              + "         } finally {"
                              + "             future.completeExceptionally(e);"
                              + "         }"
-                             + "     }%n"
-                             + "})", responseHandlerName);
+                             + "     }"
+                             + "%s"
+                             + "})", responseHandlerName, publishMetrics());
     }
 
 
-
     @Override
     public Optional<MethodSpec> createErrorResponseHandler() {
         ClassName httpResponseHandler = ClassName.get(HttpResponseHandler.class);

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/ProtocolSpec.java

@@ -155,7 +155,8 @@ default String streamingOutputWhenComplete(String responseHandlerName) {
                              + "         runAndLogError(log, \"Exception thrown in exceptionOccurred callback, ignoring\", () "
                              + "-> %s.exceptionOccurred(e));%n"
                              + "     }%n"
-                             + "})", responseHandlerName);
+                             + "%s"
+                             + "})", responseHandlerName, publishMetrics());
 
     }
 
@@ -177,4 +178,16 @@ default TypeName executeFutureValueType(OperationModel opModel, PoetExtensions p
     default TypeName getPojoResponseType(OperationModel opModel, PoetExtensions poetExtensions) {
         return poetExtensions.getModelClass(opModel.getReturnType().getReturnType());
     }
+
+    default String publishMetricsWhenComplete() {
+        return String.format(".whenComplete((r, e) -> {%n"
+                             + "%s%n"
+                             + "})", publishMetrics());
+    }
+
+    default String publishMetrics() {
+        return "Optional<MetricPublisher> metricPublisher = MetricUtils.resolvePublisher(clientConfiguration, "
+               + "requestOverrideConfig);\n"
+               + "metricPublisher.ifPresent(p -> p.publish(apiCallMetricCollector.collect()));";
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryProtocolSpec.java

@@ -24,6 +24,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
@@ -154,9 +155,12 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                                                             "errorResponseHandler",
                                                             opModel.getInput().getVariableName(),
                                                             opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "");
-
+        builder.addStatement("$T requestOverrideConfig = $L.overrideConfiguration().orElse(null)",
+                             AwsRequestOverrideConfiguration.class, opModel.getInput().getVariableName());
         if (opModel.hasStreamingOutput()) {
             builder.add("executeFuture$L;", streamingOutputWhenComplete("asyncResponseTransformer"));
+        }  else {
+            builder.add("executeFuture$L;", publishMetricsWhenComplete());
         }
         builder.addStatement("return executeFuture");
         return builder.build();

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/XmlProtocolSpec.java

@@ -20,6 +20,7 @@
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
 import java.util.Optional;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
@@ -169,8 +170,12 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                                     opModel.getInput().getVariableName(),
                                     opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "");
 
+        builder.addStatement("$T requestOverrideConfig = $L.overrideConfiguration().orElse(null)",
+                             AwsRequestOverrideConfiguration.class, opModel.getInput().getVariableName());
         if (opModel.hasStreamingOutput()) {
             builder.add("executeFuture$L;", streamingOutputWhenComplete("asyncResponseTransformer"));
+        } else {
+            builder.add("executeFuture$L;", publishMetricsWhenComplete());
         }
         builder.addStatement("return executeFuture");
         return builder.build();