Add tracking of RequestBody/ResponseTransfromer implementations used in UserAgent

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

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "\"AWS SDK for Java v2\"",
+    "contributor": "",
+    "description": "Add tracking of RequestBody/ResponseTransfromer implementations used in UserAgent."
+}

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

@@ -235,6 +235,10 @@ public CodeBlock executionHandler(OperationModel opModel) {
 
         codeBlock.add(RequestCompressionTrait.create(opModel, model));
 
+        if (opModel.hasStreamingOutput()) {
+            codeBlock.add(".withResponseTransformer(responseTransformer)");
+        }
+
         if (opModel.hasStreamingInput()) {
             codeBlock.add(".withRequestBody(requestBody)")
                      .add(".withMarshaller($L)", syncStreamingMarshaller(model, opModel, marshaller));
@@ -310,6 +314,10 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
             builder.add(NoneAuthTypeRequestTrait.create(opModel));
         }
 
+        if (opModel.hasStreamingOutput()) {
+            builder.add(".withAsyncResponseTransformer(asyncResponseTransformer)");
+        }
+
         builder.add(RequestCompressionTrait.create(opModel, model))
                .add(".withInput($L)$L)",
                     opModel.getInput().getVariableName(), asyncResponseTransformerVariable(isStreaming, isRestJson, opModel))

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

@@ -128,7 +128,9 @@ public CodeBlock executionHandler(OperationModel opModel) {
         }
 
         codeBlock.add(RequestCompressionTrait.create(opModel, intermediateModel));
-
+        if (opModel.hasStreamingOutput()) {
+            codeBlock.add(".withResponseTransformer(responseTransformer)");
+        }
         if (opModel.hasStreamingInput()) {
             return codeBlock.add(".withRequestBody(requestBody)")
                             .add(".withMarshaller($L));", syncStreamingMarshaller(intermediateModel, opModel, marshaller))
@@ -170,6 +172,10 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
 
         builder.add(RequestCompressionTrait.create(opModel, intermediateModel));
 
+        if (opModel.hasStreamingOutput()) {
+            builder.add(".withAsyncResponseTransformer(asyncResponseTransformer)");
+        }
+
         builder.add(hostPrefixExpression(opModel) + asyncRequestBody + ".withInput($L)$L);",
                     opModel.getInput().getVariableName(),
                     opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "");

core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/AwsExecutionContextBuilder.java

@@ -21,6 +21,8 @@
 import static software.amazon.awssdk.core.interceptor.SdkExecutionAttribute.RESOLVED_CHECKSUM_SPECS;
 import static software.amazon.awssdk.core.internal.useragent.BusinessMetricsUtils.resolveRetryMode;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -36,6 +38,8 @@
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.SelectedAuthScheme;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
@@ -49,8 +53,11 @@
 import software.amazon.awssdk.core.internal.InternalCoreExecutionAttribute;
 import software.amazon.awssdk.core.internal.util.HttpChecksumResolver;
 import software.amazon.awssdk.core.signer.Signer;
+import software.amazon.awssdk.core.sync.ResponseTransformer;
+import software.amazon.awssdk.core.useragent.AdditionalMetadata;
 import software.amazon.awssdk.core.useragent.BusinessMetricCollection;
 import software.amazon.awssdk.endpoints.EndpointProvider;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.auth.scheme.NoAuthAuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
@@ -159,6 +166,8 @@ private AwsExecutionContextBuilder() {
                                              signer, executionAttributes, executionAttributes.getOptionalAttribute(
                                                  AwsSignerExecutionAttribute.AWS_CREDENTIALS).orElse(null)));
 
+        putStreamingInputOutputTypesMetadata(executionAttributes, executionParams);
+
         return ExecutionContext.builder()
                                .interceptorChain(executionInterceptorChain)
                                .interceptorContext(interceptorContext)
@@ -168,6 +177,57 @@ private AwsExecutionContextBuilder() {
                                .build();
     }
 
+    private static <InputT extends SdkRequest, OutputT extends SdkResponse> void putStreamingInputOutputTypesMetadata(
+        ExecutionAttributes executionAttributes, ClientExecutionParams<InputT, OutputT> executionParams) {
+        List<AdditionalMetadata> userAgentMetadata = new ArrayList<>();
+
+        if (executionParams.getRequestBody() != null) {
+            userAgentMetadata.add(
+                AdditionalMetadata
+                    .builder()
+                    .name("rb")
+                    .value(ContentStreamProvider.ProviderType.shortValueFromName(
+                        executionParams.getRequestBody().contentStreamProvider().name())
+                    )
+                    .build());
+        }
+
+        if (executionParams.getAsyncRequestBody() != null) {
+            userAgentMetadata.add(
+                AdditionalMetadata
+                    .builder()
+                    .name("rb")
+                    .value(AsyncRequestBody.BodyType.shortValueFromName(
+                        executionParams.getAsyncRequestBody().body())
+                    )
+                    .build());
+        }
+
+        if (executionParams.getResponseTransformer() != null) {
+            userAgentMetadata.add(
+                AdditionalMetadata
+                    .builder()
+                    .name("rt")
+                    .value(ResponseTransformer.TransformerType.shortValueFromName(
+                        executionParams.getResponseTransformer().name())
+                    )
+                    .build());
+        }
+
+        if (executionParams.getAsyncResponseTransformer() != null) {
+            userAgentMetadata.add(
+                AdditionalMetadata
+                    .builder()
+                    .name("rt")
+                    .value(AsyncResponseTransformer.TransformerType.shortValueFromName(
+                        executionParams.getAsyncResponseTransformer().name())
+                    )
+                    .build());
+        }
+
+        executionAttributes.putAttribute(SdkInternalExecutionAttribute.USER_AGENT_METADATA, userAgentMetadata);
+    }
+
     /**
      * We will load the old (non-SRA) signer if this client seems like an old version or the customer has provided a signer
      * override. We assume that if there's no auth schemes defined, we're on the old code path.

core/aws-core/src/test/java/software/amazon/awssdk/awscore/internal/AwsExecutionContextBuilderTest.java

@@ -22,11 +22,13 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.io.File;
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
 import org.junit.Before;
@@ -44,6 +46,8 @@
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.SelectedAuthScheme;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.checksums.ChecksumSpecs;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
@@ -59,18 +63,19 @@
 import software.amazon.awssdk.core.internal.util.HttpChecksumUtils;
 import software.amazon.awssdk.core.signer.NoOpSigner;
 import software.amazon.awssdk.core.signer.Signer;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.core.sync.ResponseTransformer;
+import software.amazon.awssdk.core.useragent.AdditionalMetadata;
 import software.amazon.awssdk.http.auth.aws.scheme.AwsV4AuthScheme;
 import software.amazon.awssdk.http.auth.scheme.NoAuthAuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.http.auth.spi.signer.HttpSigner;
-import software.amazon.awssdk.http.auth.spi.signer.SignerProperty;
 import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.IdentityProviders;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.profiles.ProfileFile;
-import software.amazon.awssdk.regions.RegionScope;
 
 @RunWith(MockitoJUnitRunner.class)
 public class AwsExecutionContextBuilderTest {
@@ -437,6 +442,74 @@ public void invokeInterceptorsAndCreateExecutionContext_requestOverrideForIdenti
         assertThat(actualTokenProvider).isSameAs(requestTokenProvider);
     }
 
+    @Test
+    public void invokeInterceptorsAndCreateExecutionContext_withRequestBody_addsUserAgentMetadata() throws IOException {
+        ClientExecutionParams<SdkRequest, SdkResponse> executionParams = clientExecutionParams();
+        File testFile = File.createTempFile("testFile", UUID.randomUUID().toString());
+        testFile.deleteOnExit();
+        executionParams.withRequestBody(RequestBody.fromFile(testFile));
+
+        ExecutionContext executionContext =
+            AwsExecutionContextBuilder.invokeInterceptorsAndCreateExecutionContext(executionParams,
+                                                                                   testClientConfiguration().build());
+
+        ExecutionAttributes executionAttributes = executionContext.executionAttributes();
+        assertThat(executionAttributes.getAttribute(SdkInternalExecutionAttribute.USER_AGENT_METADATA)).isEqualTo(
+            Collections.singletonList(AdditionalMetadata.builder().name("rb").value("f").build())
+        );
+    }
+
+    @Test
+    public void invokeInterceptorsAndCreateExecutionContext_withResponseTransformer_addsUserAgentMetadata() throws IOException {
+        ClientExecutionParams<SdkRequest, SdkResponse> executionParams = clientExecutionParams();
+        File testFile = File.createTempFile("testFile", UUID.randomUUID().toString());
+        testFile.deleteOnExit();
+        executionParams.withResponseTransformer(ResponseTransformer.toFile(testFile));
+
+        ExecutionContext executionContext =
+            AwsExecutionContextBuilder.invokeInterceptorsAndCreateExecutionContext(executionParams,
+                                                                                   testClientConfiguration().build());
+
+        ExecutionAttributes executionAttributes = executionContext.executionAttributes();
+        assertThat(executionAttributes.getAttribute(SdkInternalExecutionAttribute.USER_AGENT_METADATA)).isEqualTo(
+            Collections.singletonList(AdditionalMetadata.builder().name("rt").value("f").build())
+        );
+    }
+
+    @Test
+    public void invokeInterceptorsAndCreateExecutionContext_withAsyncRequestBody_addsUserAgentMetadata() throws IOException {
+        ClientExecutionParams<SdkRequest, SdkResponse> executionParams = clientExecutionParams();
+        File testFile = File.createTempFile("testFile", UUID.randomUUID().toString());
+        testFile.deleteOnExit();
+        executionParams.withAsyncRequestBody(AsyncRequestBody.fromFile(testFile));
+
+        ExecutionContext executionContext =
+            AwsExecutionContextBuilder.invokeInterceptorsAndCreateExecutionContext(executionParams,
+                                                                                   testClientConfiguration().build());
+
+        ExecutionAttributes executionAttributes = executionContext.executionAttributes();
+        assertThat(executionAttributes.getAttribute(SdkInternalExecutionAttribute.USER_AGENT_METADATA)).isEqualTo(
+            Collections.singletonList(AdditionalMetadata.builder().name("rb").value("f").build())
+        );
+    }
+
+    @Test
+    public void invokeInterceptorsAndCreateExecutionContext_withAsyncResponseTransformer_addsUserAgentMetadata() throws IOException {
+        ClientExecutionParams<SdkRequest, SdkResponse> executionParams = clientExecutionParams();
+        File testFile = File.createTempFile("testFile", UUID.randomUUID().toString());
+        testFile.deleteOnExit();
+        executionParams.withAsyncResponseTransformer(AsyncResponseTransformer.toFile(testFile));
+
+        ExecutionContext executionContext =
+            AwsExecutionContextBuilder.invokeInterceptorsAndCreateExecutionContext(executionParams,
+                                                                                   testClientConfiguration().build());
+
+        ExecutionAttributes executionAttributes = executionContext.executionAttributes();
+        assertThat(executionAttributes.getAttribute(SdkInternalExecutionAttribute.USER_AGENT_METADATA)).isEqualTo(
+            Collections.singletonList(AdditionalMetadata.builder().name("rt").value("f").build())
+        );
+    }
+
     private ClientExecutionParams<SdkRequest, SdkResponse> clientExecutionParams() {
         return clientExecutionParams(sdkRequest);
     }

core/sdk-core/src/main/java/software/amazon/awssdk/core/async/AsyncRequestBody.java

@@ -23,11 +23,13 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.function.Consumer;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
+import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.core.FileRequestBodyConfiguration;
 import software.amazon.awssdk.core.internal.async.ByteBuffersAsyncRequestBody;
@@ -37,6 +39,7 @@
 import software.amazon.awssdk.core.internal.util.Mimetype;
 import software.amazon.awssdk.utils.BinaryUtils;
 import software.amazon.awssdk.utils.Validate;
+import software.amazon.awssdk.utils.internal.EnumUtils;
 
 /**
  * Interface to allow non-blocking streaming of request content. This follows the reactive streams pattern where this interface is
@@ -74,6 +77,16 @@ default String contentType() {
         return Mimetype.MIMETYPE_OCTET_STREAM;
     }
 
+    /**
+     * Each AsyncRequestBody should return a well-formed name that can be used to identify the implementation.
+     * The body name should only include alphanumeric characters.
+     *
+     * @return String containing the identifying name of this AsyncRequestBody implementation.
+     */
+    default String body() {
+        return BodyType.UNKNOWN.getName();
+    }
+
     /**
      * Creates an {@link AsyncRequestBody} the produces data from the input ByteBuffer publisher. The data is delivered when the
      * publisher publishes the data.
@@ -96,6 +109,11 @@ public Optional<Long> contentLength() {
             public void subscribe(Subscriber<? super ByteBuffer> s) {
                 publisher.subscribe(s);
             }
+
+            @Override
+            public String body() {
+                return BodyType.PUBLISHER.getName();
+            }
         };
     }
 
@@ -513,4 +531,36 @@ default SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfi
         Validate.notNull(splitConfiguration, "splitConfiguration");
         return split(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
     }
+
+    @SdkProtectedApi
+    enum BodyType {
+        FILE("File", "f"),
+        BYTES("Bytes", "b"),
+        STREAM("Stream", "s"),
+        PUBLISHER("Publisher", "p"),
+        UNKNOWN("Unknown", "u");
+
+        private static final Map<String, BodyType> VALUE_MAP =
+            EnumUtils.uniqueIndex(BodyType.class, BodyType::getName);
+
+        private final String name;
+        private final String shortValue;
+
+        BodyType(String name, String shortValue) {
+            this.name = name;
+            this.shortValue = shortValue;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public String getShortValue() {
+            return shortValue;
+        }
+
+        public static String shortValueFromName(String name) {
+            return VALUE_MAP.getOrDefault(name, UNKNOWN).getShortValue();
+        }
+    }
 }