aws · dagnir · Jun 30, 2020 · Jun 29, 2020
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java
@@ -18,6 +18,7 @@
 import static com.squareup.javapoet.TypeSpec.Builder;
 import static java.util.Collections.singletonList;
 import static javax.lang.model.element.Modifier.PRIVATE;
+import static javax.lang.model.element.Modifier.STATIC;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applyPaginatorUserAgentMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
 import static software.amazon.awssdk.codegen.poet.client.SyncClientClass.getProtocolSpecs;
@@ -31,8 +32,8 @@
 import com.squareup.javapoet.TypeSpec;
 import java.net.URI;
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.List;
-import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
@@ -55,14 +56,14 @@
 import software.amazon.awssdk.codegen.poet.StaticImport;
 import software.amazon.awssdk.codegen.poet.client.specs.ProtocolSpec;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
+import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
 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;
@@ -71,7 +72,7 @@
 import software.amazon.awssdk.utils.FunctionalUtils;
 
 public final class AsyncClientClass extends AsyncClientInterface {
-    private static final String PUBLISHER_NAME = "metricPublisher";
+    private static final String PUBLISHER_NAME = "metricPublishers";
     private static final String METRIC_COLLECTOR_NAME = "apiCallMetricCollector";
     private final IntermediateModel model;
     private final PoetExtensions poetExtensions;
@@ -108,7 +109,8 @@ public TypeSpec poetSpec() {
                     .addMethods(operations())
                     .addMethod(closeMethod())
                     .addMethods(protocolSpec.additionalMethods())
-                    .addMethod(protocolSpec.initProtocolFactory(model));
+                    .addMethod(protocolSpec.initProtocolFactory(model))
+                    .addMethod(resolveMetricPublishersMethod());
 
         // Kinesis doesn't support CBOR for STS yet so need another protocol factory for JSON
         if (model.getMetadata().isCborProtocol()) {
@@ -240,13 +242,12 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
                                  "() -> $N.exceptionOccurred(t))", paramName);
         }
 
-        builder.addStatement("$T<$T> $N = $T.resolvePublisher(clientConfiguration, $N.overrideConfiguration().orElse(null))",
-                             Optional.class,
+        builder.addStatement("$T<$T> $N = resolveMetricPublishers(clientConfiguration, $N.overrideConfiguration().orElse(null))",
+                             List.class,
                              MetricPublisher.class,
                              PUBLISHER_NAME,
-                             MetricUtils.class,
                              opModel.getInput().getVariableName())
-               .addStatement("$N.ifPresent(p -> p.publish($N.collect()))", PUBLISHER_NAME, "apiCallMetricCollector")
+               .addStatement("$N.forEach(p -> p.publish($N.collect()))", PUBLISHER_NAME, "apiCallMetricCollector")
                .addStatement("return $T.failedFuture(t)", CompletableFutureUtils.class)
                .endControlFlow();
 
@@ -326,4 +327,39 @@ private MethodSpec utilitiesMethod() {
                                        String.join(",", config.getCreateMethodParams()))
                          .build();
     }
+
+    private MethodSpec resolveMetricPublishersMethod() {
+        String clientConfigName = "clientConfiguration";
+        String requestOverrideConfigName = "requestOverrideConfiguration";
+
+        MethodSpec.Builder methodBuilder = MethodSpec.methodBuilder("resolveMetricPublishers")
+                .addModifiers(PRIVATE, STATIC)
+                .returns(ParameterizedTypeName.get(List.class, MetricPublisher.class))
+                .addParameter(SdkClientConfiguration.class, clientConfigName)
+                .addParameter(RequestOverrideConfiguration.class, requestOverrideConfigName);
+
+        String publishersName = "publishers";
+
+        methodBuilder.addStatement("$T $N = null", ParameterizedTypeName.get(List.class, MetricPublisher.class), publishersName);
+
+        methodBuilder.beginControlFlow("if ($N != null)", requestOverrideConfigName)
+                .addStatement("$N = $N.metricPublishers()", publishersName, requestOverrideConfigName)
+                .endControlFlow();
+
+        methodBuilder.beginControlFlow("if ($1N == null || $1N.isEmpty())", publishersName)
+                .addStatement("$N = $N.option($T.$N)",
+                              publishersName,
+                              clientConfigName,
+                              SdkClientOption.class,
+                              "METRIC_PUBLISHERS")
+                .endControlFlow();
+
+        methodBuilder.beginControlFlow("if ($1N == null)", publishersName)
+                .addStatement("$N = $T.emptyList()", publishersName, Collections.class)
+                .endControlFlow();
+
+        methodBuilder.addStatement("return $N", publishersName);
+
+        return methodBuilder.build();
+    }
 }
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java
@@ -17,18 +17,20 @@
 
 import static javax.lang.model.element.Modifier.FINAL;
 import static javax.lang.model.element.Modifier.PRIVATE;
+import static javax.lang.model.element.Modifier.STATIC;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applyPaginatorUserAgentMethod;
 import static software.amazon.awssdk.codegen.poet.client.ClientClassUtils.applySignerOverrideMethod;
 
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeSpec;
 import com.squareup.javapoet.TypeSpec.Builder;
 import java.net.URI;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
-import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -48,12 +50,12 @@
 import software.amazon.awssdk.codegen.poet.client.specs.QueryProtocolSpec;
 import software.amazon.awssdk.codegen.poet.client.specs.XmlProtocolSpec;
 import software.amazon.awssdk.codegen.utils.PaginatorUtils;
+import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
 import software.amazon.awssdk.core.client.handler.SyncClientHandler;
 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;
@@ -89,7 +91,8 @@ public TypeSpec poetSpec() {
                                         .addMethod(constructor())
                                         .addMethod(nameMethod())
                                         .addMethods(protocolSpec.additionalMethods())
-                                        .addMethods(operations());
+                                        .addMethods(operations())
+                                        .addMethod(resolveMetricPublishersMethod());
 
         protocolSpec.createErrorResponseHandler().ifPresent(classBuilder::addMethod);
 
@@ -196,7 +199,7 @@ private List<MethodSpec> operationMethodSpecs(OperationModel opModel) {
         method.addStatement("$1T $2N = $1T.create($3S)",
                 MetricCollector.class, metricCollectorName, "ApiCall");
 
-        String publisherName = "metricPublisher";
+        String publishersName = "metricPublishers";
 
         method.beginControlFlow("try")
                 .addStatement("$N.reportMetric($T.$L, $S)", metricCollectorName, CoreMetric.class, "SERVICE_ID",
@@ -206,13 +209,12 @@ private List<MethodSpec> operationMethodSpecs(OperationModel opModel) {
                 .addCode(protocolSpec.executionHandler(opModel))
                 .endControlFlow()
                 .beginControlFlow("finally")
-                .addStatement("$T<$T> $N = $T.resolvePublisher(clientConfiguration, $N)",
-                              Optional.class,
+                .addStatement("$T<$T> $N = resolveMetricPublishers(clientConfiguration, $N.overrideConfiguration().orElse(null))",
+                              List.class,
                               MetricPublisher.class,
-                              publisherName,
-                              MetricUtils.class,
+                              publishersName,
                               opModel.getInput().getVariableName())
-                .addStatement("$N.ifPresent(p -> p.publish($N.collect()))", publisherName, metricCollectorName)
+                .addStatement("$N.forEach(p -> p.publish($N.collect()))", publishersName, metricCollectorName)
                 .endControlFlow();
 
         methods.add(method.build());
@@ -286,4 +288,39 @@ static ProtocolSpec getProtocolSpecs(PoetExtensions poetExtensions, Intermediate
                 throw new RuntimeException("Unknown protocol: " + protocol.name());
         }
     }
+
+    private MethodSpec resolveMetricPublishersMethod() {
+        String clientConfigName = "clientConfiguration";
+        String requestOverrideConfigName = "requestOverrideConfiguration";
+
+        MethodSpec.Builder methodBuilder = MethodSpec.methodBuilder("resolveMetricPublishers")
+                .addModifiers(PRIVATE, STATIC)
+                .returns(ParameterizedTypeName.get(List.class, MetricPublisher.class))
+                .addParameter(SdkClientConfiguration.class, clientConfigName)
+                .addParameter(RequestOverrideConfiguration.class, requestOverrideConfigName);
+
+        String publishersName = "publishers";
+
+        methodBuilder.addStatement("$T $N = null", ParameterizedTypeName.get(List.class, MetricPublisher.class), publishersName);
+
+        methodBuilder.beginControlFlow("if ($N != null)", requestOverrideConfigName)
+                .addStatement("$N = $N.metricPublishers()", publishersName, requestOverrideConfigName)
+                .endControlFlow();
+
+        methodBuilder.beginControlFlow("if ($1N == null || $1N.isEmpty())", publishersName)
+                .addStatement("$N = $N.option($T.$N)",
+                              publishersName,
+                              clientConfigName,
+                              SdkClientOption.class,
+                              "METRIC_PUBLISHERS")
+                .endControlFlow();
+
+        methodBuilder.beginControlFlow("if ($1N == null)", publishersName)
+                .addStatement("$N = $T.emptyList()", publishersName, Collections.class)
+                .endControlFlow();
+
+        methodBuilder.addStatement("return $N", publishersName);
+
+        return methodBuilder.build();
+    }
 }
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/ProtocolSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/ProtocolSpec.java
@@ -186,8 +186,8 @@ default String publishMetricsWhenComplete() {
     }
 
     default String publishMetrics() {
-        return "Optional<MetricPublisher> metricPublisher = MetricUtils.resolvePublisher(clientConfiguration, "
+        return "List<MetricPublisher> metricPublishers = resolveMetricPublishers(clientConfiguration, "
                + "requestOverrideConfig);\n"
-               + "metricPublisher.ifPresent(p -> p.publish(apiCallMetricCollector.collect()));";
+               + "metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));";
     }
 }