Implementing Protocol Rpc V2 Cbor FeatureID

.changes/next-release/feature-AWSSDKforJavav2-ee91960.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Add business metrics support for RPC v2 CBOR protocol to track smithy rpcv2 cbor protocol usage."
+}

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

@@ -25,10 +25,14 @@
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeVariableName;
 import com.squareup.javapoet.WildcardTypeName;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 import java.util.function.Function;
 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;
@@ -48,6 +52,7 @@
 import software.amazon.awssdk.codegen.poet.client.traits.RequestCompressionTrait;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
 import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
+import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.SdkPojoBuilder;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
@@ -56,6 +61,7 @@
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.HttpResponseHandler;
 import software.amazon.awssdk.core.protocol.VoidSdkResponse;
+import software.amazon.awssdk.core.useragent.BusinessMetricFeatureId;
 import software.amazon.awssdk.protocols.cbor.AwsCborProtocolFactory;
 import software.amazon.awssdk.protocols.core.ExceptionMetadata;
 import software.amazon.awssdk.protocols.json.AwsJsonProtocol;
@@ -224,7 +230,10 @@ public CodeBlock executionHandler(OperationModel opModel) {
                      .add(discoveredEndpoint(opModel))
                      .add(credentialType(opModel, model))
                      .add(".withRequestConfiguration(clientConfiguration)")
-                     .add(".withInput($L)\n", opModel.getInput().getVariableName())
+                     .add(".withInput($L)\n", 
+                          model.getMetadata().isRpcV2CborProtocol() ? 
+                          "applyRpcV2CborUserAgent(" + opModel.getInput().getVariableName() + ")" : 
+                          opModel.getInput().getVariableName())
                      .add(".withMetricCollector(apiCallMetricCollector)")
                      .add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
                      .add(HttpChecksumTrait.create(opModel));
@@ -320,7 +329,10 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
 
         builder.add(RequestCompressionTrait.create(opModel, model))
                .add(".withInput($L)$L)",
-                    opModel.getInput().getVariableName(), asyncResponseTransformerVariable(isStreaming, isRestJson, opModel))
+                    intermediateModel.getMetadata().isRpcV2CborProtocol() ? 
+                    "applyRpcV2CborUserAgent(" + opModel.getInput().getVariableName() + ")" : 
+                    opModel.getInput().getVariableName(), 
+                    asyncResponseTransformerVariable(isStreaming, isRestJson, opModel))
                .add(opModel.getEndpointDiscovery() != null ? ");" : ";");
 
         if (opModel.hasStreamingOutput()) {
@@ -568,4 +580,49 @@ private String protocolFactoryLiteral(IntermediateModel model, OperationModel op
     private boolean isRestJson(IntermediateModel model) {
         return model.getMetadata().getProtocol() == Protocol.REST_JSON;
     }
+
+    @Override
+    public List<MethodSpec> additionalMethods() {
+        List<MethodSpec> methods = new ArrayList<>();
+
+        applyRpcV2CborUserAgentMethod().ifPresent(methods::add);
+
+        return methods;
+    }
+
+    private Optional<MethodSpec> applyRpcV2CborUserAgentMethod() {
+        if (!model.getMetadata().isRpcV2CborProtocol()) {
+            return Optional.empty();
+        }
+
+        TypeVariableName typeVariableName =
+            TypeVariableName.get("T", poetExtensions.getModelClass(model.getSdkRequestBaseClassName()));
+
+        ParameterizedTypeName parameterizedTypeName = ParameterizedTypeName
+            .get(ClassName.get(Consumer.class), ClassName.get(AwsRequestOverrideConfiguration.Builder.class));
+
+        CodeBlock codeBlock = CodeBlock.builder()
+                                       .addStatement("$T userAgentApplier = b -> "
+                                                     + "b.addApiName($T.builder().name($S).version($S).build())",
+                                                     parameterizedTypeName, ApiName.class,
+                                                     "sdk-metrics",
+                                                     BusinessMetricFeatureId.PROTOCOL_RPC_V2_CBOR.value())
+                                       .addStatement("$T overrideConfiguration =\n"
+                                                     + "            request.overrideConfiguration().map(c -> c.toBuilder()"
+                                                     + ".applyMutation(userAgentApplier).build())\n"
+                                                     + "            .orElse((AwsRequestOverrideConfiguration.builder()"
+                                                     + ".applyMutation(userAgentApplier).build()))",
+                                                     AwsRequestOverrideConfiguration.class)
+                                       .addStatement("return (T) request.toBuilder().overrideConfiguration(overrideConfiguration)"
+                                                     + ".build()")
+                                       .build();
+
+        return Optional.of(MethodSpec.methodBuilder("applyRpcV2CborUserAgent")
+                                    .addModifiers(Modifier.PRIVATE, Modifier.STATIC)
+                                    .addParameter(typeVariableName, "request")
+                                    .addTypeVariable(typeVariableName)
+                                    .addCode(codeBlock)
+                                    .returns(typeVariableName)
+                                    .build());
+    }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-rpcv2-async-client-class.java

@@ -12,11 +12,13 @@
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.annotations.Generated;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.client.handler.AwsAsyncClientHandler;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.awscore.internal.AwsProtocolMetadata;
 import software.amazon.awssdk.awscore.internal.AwsServiceProtocol;
 import software.amazon.awssdk.awscore.retry.AwsRetryStrategy;
+import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.SdkPlugin;
 import software.amazon.awssdk.core.SdkRequest;
@@ -66,6 +68,7 @@
 import software.amazon.awssdk.services.smithyrpcv2protocol.model.SimpleScalarPropertiesRequest;
 import software.amazon.awssdk.services.smithyrpcv2protocol.model.SimpleScalarPropertiesResponse;
 import software.amazon.awssdk.services.smithyrpcv2protocol.model.SmithyRpcV2ProtocolException;
+import software.amazon.awssdk.services.smithyrpcv2protocol.model.SmithyRpcV2ProtocolRequest;
 import software.amazon.awssdk.services.smithyrpcv2protocol.model.SparseNullsOperationRequest;
 import software.amazon.awssdk.services.smithyrpcv2protocol.model.SparseNullsOperationResponse;
 import software.amazon.awssdk.services.smithyrpcv2protocol.model.ValidationException;
@@ -173,7 +176,7 @@ public CompletableFuture<EmptyInputOutputResponse> emptyInputOutput(EmptyInputOu
                              .withMarshaller(new EmptyInputOutputRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(emptyInputOutputRequest));
+                             .withInput(applyRpcV2CborUserAgent(emptyInputOutputRequest)));
             CompletableFuture<EmptyInputOutputResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -246,7 +249,7 @@ public CompletableFuture<Float16Response> float16(Float16Request float16Request)
                                                                                      .withProtocolMetadata(protocolMetadata).withMarshaller(new Float16RequestMarshaller(protocolFactory))
                                                                                      .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                                                                                      .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                                                                                     .withInput(float16Request));
+                                                                                     .withInput(applyRpcV2CborUserAgent(float16Request)));
             CompletableFuture<Float16Response> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -321,7 +324,7 @@ public CompletableFuture<FractionalSecondsResponse> fractionalSeconds(Fractional
                              .withMarshaller(new FractionalSecondsRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(fractionalSecondsRequest));
+                             .withInput(applyRpcV2CborUserAgent(fractionalSecondsRequest)));
             CompletableFuture<FractionalSecondsResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -398,7 +401,7 @@ public CompletableFuture<GreetingWithErrorsResponse> greetingWithErrors(Greeting
                              .withMarshaller(new GreetingWithErrorsRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(greetingWithErrorsRequest));
+                             .withInput(applyRpcV2CborUserAgent(greetingWithErrorsRequest)));
             CompletableFuture<GreetingWithErrorsResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -472,7 +475,7 @@ public CompletableFuture<NoInputOutputResponse> noInputOutput(NoInputOutputReque
                              .withMarshaller(new NoInputOutputRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(noInputOutputRequest));
+                             .withInput(applyRpcV2CborUserAgent(noInputOutputRequest)));
             CompletableFuture<NoInputOutputResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -549,7 +552,7 @@ public CompletableFuture<OperationWithDefaultsResponse> operationWithDefaults(
                              .withMarshaller(new OperationWithDefaultsRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(operationWithDefaultsRequest));
+                             .withInput(applyRpcV2CborUserAgent(operationWithDefaultsRequest)));
             CompletableFuture<OperationWithDefaultsResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -625,7 +628,7 @@ public CompletableFuture<OptionalInputOutputResponse> optionalInputOutput(
                              .withMarshaller(new OptionalInputOutputRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(optionalInputOutputRequest));
+                             .withInput(applyRpcV2CborUserAgent(optionalInputOutputRequest)));
             CompletableFuture<OptionalInputOutputResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -700,7 +703,7 @@ public CompletableFuture<RecursiveShapesResponse> recursiveShapes(RecursiveShape
                              .withMarshaller(new RecursiveShapesRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(recursiveShapesRequest));
+                             .withInput(applyRpcV2CborUserAgent(recursiveShapesRequest)));
             CompletableFuture<RecursiveShapesResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -776,7 +779,7 @@ public CompletableFuture<RpcV2CborDenseMapsResponse> rpcV2CborDenseMaps(RpcV2Cbo
                              .withMarshaller(new RpcV2CborDenseMapsRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(rpcV2CborDenseMapsRequest));
+                             .withInput(applyRpcV2CborUserAgent(rpcV2CborDenseMapsRequest)));
             CompletableFuture<RpcV2CborDenseMapsResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -851,7 +854,7 @@ public CompletableFuture<RpcV2CborListsResponse> rpcV2CborLists(RpcV2CborListsRe
                              .withMarshaller(new RpcV2CborListsRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(rpcV2CborListsRequest));
+                             .withInput(applyRpcV2CborUserAgent(rpcV2CborListsRequest)));
             CompletableFuture<RpcV2CborListsResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -928,7 +931,7 @@ public CompletableFuture<RpcV2CborSparseMapsResponse> rpcV2CborSparseMaps(
                              .withMarshaller(new RpcV2CborSparseMapsRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(rpcV2CborSparseMapsRequest));
+                             .withInput(applyRpcV2CborUserAgent(rpcV2CborSparseMapsRequest)));
             CompletableFuture<RpcV2CborSparseMapsResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -1004,7 +1007,7 @@ public CompletableFuture<SimpleScalarPropertiesResponse> simpleScalarProperties(
                              .withMarshaller(new SimpleScalarPropertiesRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(simpleScalarPropertiesRequest));
+                             .withInput(applyRpcV2CborUserAgent(simpleScalarPropertiesRequest)));
             CompletableFuture<SimpleScalarPropertiesResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -1080,7 +1083,7 @@ public CompletableFuture<SparseNullsOperationResponse> sparseNullsOperation(
                              .withMarshaller(new SparseNullsOperationRequestMarshaller(protocolFactory))
                              .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                              .withRequestConfiguration(clientConfiguration).withMetricCollector(apiCallMetricCollector)
-                             .withInput(sparseNullsOperationRequest));
+                             .withInput(applyRpcV2CborUserAgent(sparseNullsOperationRequest)));
             CompletableFuture<SparseNullsOperationResponse> whenCompleted = executeFuture.whenComplete((r, e) -> {
                 metricPublishers.forEach(p -> p.publish(apiCallMetricCollector.collect()));
             });
@@ -1102,6 +1105,15 @@ public final String serviceName() {
         return SERVICE_NAME;
     }
 
+    private static <T extends SmithyRpcV2ProtocolRequest> T applyRpcV2CborUserAgent(T request) {
+        Consumer<AwsRequestOverrideConfiguration.Builder> userAgentApplier = b -> b.addApiName(ApiName.builder()
+                                                                                                      .name("sdk-metrics").version("M").build());
+        AwsRequestOverrideConfiguration overrideConfiguration = request.overrideConfiguration()
+                                                                       .map(c -> c.toBuilder().applyMutation(userAgentApplier).build())
+                                                                       .orElse((AwsRequestOverrideConfiguration.builder().applyMutation(userAgentApplier).build()));
+        return (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
+    }
+
     private <T extends BaseAwsJsonProtocolFactory.Builder<T>> T init(T builder) {
         return builder.clientConfiguration(clientConfiguration)
                       .defaultServiceExceptionSupplier(SmithyRpcV2ProtocolException::builder)
@@ -1170,4 +1182,4 @@ private HttpResponseHandler<AwsServiceException> createErrorResponseHandler(Base
     public void close() {
         clientHandler.close();
     }
-}
+}