Add full support for AWS chunked content encoding

.changelog/1769635097.md

@@ -0,0 +1,31 @@
+---
+applies_to:
+- aws-sdk-rust
+authors:
+- ysaito1001
+references:
+- smithy-rs#1798
+- smithy-rs#4188
+breaking: false
+new_feature: true
+bug_fix: true
+---
+Add full support for [AWS chunked content encoding](https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-streaming.html). This affects S3's `PutObject` and `UploadPart` operations with streaming payloads. Streaming payloads are now split into 64 KiB chunks by default (the last chunk may be smaller) instead of being sent as a single chunk.
+
+If the default chunk size does not suit the use case, it can be configured via the service config:
+
+```
+// Custom chunk size
+let config = aws_sdk_s3::Config::builder()
+    .chunk_size(Some(10240)) // 10 KiB chunks
+    .build();
+let client = aws_sdk_s3::Client::from_conf(config);
+```
+
+```
+// Disable chunking (buffers entire body in memory and sends it as one chunk)
+let config = aws_sdk_s3::Config::builder()
+    .chunk_size(None)
+    .build();
+let client = aws_sdk_s3::Client::from_conf(config);
+```

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AwsChunkedContentEncodingDecorator.kt

@@ -13,12 +13,15 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection
+import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.Visibility
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.util.getTrait
 import software.amazon.smithy.rust.codegen.core.util.hasStreamingMember
 
@@ -31,34 +34,127 @@ class AwsChunkedContentEncodingDecorator : ClientCodegenDecorator {
     override val order: Byte =
         (minOf(HttpRequestChecksumDecorator.ORDER, HttpRequestCompressionDecorator.ORDER) - 1).toByte()
 
+    override fun configCustomizations(
+        codegenContext: ClientCodegenContext,
+        baseCustomizations: List<ConfigCustomization>,
+    ): List<ConfigCustomization> = baseCustomizations + AwsChunkedConfigCustomization(codegenContext)
+
     override fun operationCustomizations(
         codegenContext: ClientCodegenContext,
         operation: OperationShape,
         baseCustomizations: List<OperationCustomization>,
-    ) = baseCustomizations + AwsChunkedOparationCustomization(codegenContext, operation)
+    ) = baseCustomizations + AwsChunkedOperationCustomization(codegenContext, operation)
+}
+
+// TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Replace this heuristic with a dedicated
+//  Smithy trait once available to determine whether operations require aws-chunked encoding.
+private fun operationRequiresAwsChunked(
+    codegenContext: ClientCodegenContext,
+    operation: OperationShape,
+): Boolean {
+    val checksumTrait = operation.getTrait<HttpChecksumTrait>() ?: return false
+    val requestAlgorithmMember =
+        checksumTrait.requestAlgorithmMemberShape(codegenContext, operation) ?: return false
+    requestAlgorithmMember.getTrait<HttpHeaderTrait>()?.value ?: return false
+    val input = codegenContext.model.expectShape(operation.inputShape, StructureShape::class.java)
+    return input.hasStreamingMember(codegenContext.model)
+}
+
+private fun serviceRequiresAwsChunked(codegenContext: ClientCodegenContext): Boolean =
+    codegenContext.serviceShape.allOperations.any { operationId ->
+        val operation = codegenContext.model.expectShape(operationId, OperationShape::class.java)
+        operationRequiresAwsChunked(codegenContext, operation)
+    }
+
+private class AwsChunkedConfigCustomization(
+    codegenContext: ClientCodegenContext,
+) : ConfigCustomization() {
+    private val runtimeConfig = codegenContext.runtimeConfig
+    private val moduleUseName = codegenContext.moduleUseName()
+    private val serviceRequiresChunking = serviceRequiresAwsChunked(codegenContext)
+
+    override fun section(section: ServiceConfig) =
+        writable {
+            when (section) {
+                ServiceConfig.BuilderImpl -> {
+                    if (!serviceRequiresChunking) return@writable
+                    rustTemplate(
+                        """
+                        /// Sets the chunk size for [`aws-chunked encoding`].
+                        ///
+                        /// Pass `Some(size)` to use a specific chunk size (minimum 8 KiB).
+                        /// Pass `None` to use the content-length as chunk size (no chunking).
+                        ///
+                        /// The minimum chunk size of 8 KiB is validated when the request is sent.
+                        ///
+                        /// **Note:** This setting only applies to operations that support aws-chunked encoding
+                        /// and has no effect on other operations. If this method is not invoked, a default
+                        /// chunk size of 64 KiB is used.
+                        ///
+                        /// [`aws-chunked encoding`]: https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-streaming.html
+                        ///
+                        /// ## Example - Custom chunk size
+                        /// ```no_run
+                        /// ## use $moduleUseName::{Client, Config};
+                        /// ## async fn example(client: Client) -> Result<(), Box<dyn std::error::Error>> {
+                        /// let config = Config::builder()
+                        ///     .chunk_size(Some(10240)) // 10 KiB chunks
+                        ///     .build();
+                        /// let client = Client::from_conf(config);
+                        /// ## Ok(())
+                        /// ## }
+                        /// ```
+                        ///
+                        /// ## Example - No chunking (buffers entire body in memory)
+                        /// ```no_run
+                        /// ## use aws_sdk_s3::{Client, Config};
+                        /// ## async fn example(client: Client) -> Result<(), Box<dyn std::error::Error>> {
+                        /// let config = Config::builder()
+                        ///     .chunk_size(None) // Use entire content as one chunk
+                        ///     .build();
+                        /// let client = Client::from_conf(config);
+                        /// ## Ok(())
+                        /// ## }
+                        /// ```
+                        pub fn chunk_size(mut self, chunk_size: #{Option}<usize>) -> Self {
+                            self.set_chunk_size(#{Some}(chunk_size));
+                            self
+                        }
+
+                        /// Sets the chunk size for aws-chunked encoding.
+                        pub fn set_chunk_size(&mut self, chunk_size: #{Option}<#{Option}<usize>>) -> &mut Self {
+                            if let #{Some}(chunk_size) = chunk_size {
+                                let chunk_size = match chunk_size {
+                                    #{Some}(size) => #{ChunkSize}::Configured(size),
+                                    #{None} => #{ChunkSize}::DisableChunking,
+                                };
+                                self.push_runtime_plugin(#{ChunkSizeRuntimePlugin}::new(chunk_size).into_shared());
+                            }
+                            self
+                        }
+                        """,
+                        *preludeScope,
+                        "ChunkSize" to runtimeConfig.awsChunked().resolve("ChunkSize"),
+                        "ChunkSizeRuntimePlugin" to runtimeConfig.awsChunked().resolve("ChunkSizeRuntimePlugin"),
+                    )
+                }
+
+                else -> emptySection
+            }
+        }
 }
 
-private class AwsChunkedOparationCustomization(
+private class AwsChunkedOperationCustomization(
     private val codegenContext: ClientCodegenContext,
     private val operation: OperationShape,
 ) : OperationCustomization() {
-    private val model = codegenContext.model
     private val runtimeConfig = codegenContext.runtimeConfig
 
     override fun section(section: OperationSection) =
         writable {
             when (section) {
                 is OperationSection.AdditionalInterceptors -> {
-                    // TODO(https://github.com/smithy-lang/smithy-rs/issues/4382): Remove all of these early returns
-                    //  once we have the dedicated trait available in Smithy.
-                    val checksumTrait = operation.getTrait<HttpChecksumTrait>() ?: return@writable
-                    val requestAlgorithmMember =
-                        checksumTrait.requestAlgorithmMemberShape(codegenContext, operation) ?: return@writable
-                    requestAlgorithmMember.getTrait<HttpHeaderTrait>()?.value ?: return@writable
-                    val input = model.expectShape(operation.inputShape, StructureShape::class.java)
-                    if (!input.hasStreamingMember(model)) {
-                        return@writable
-                    }
+                    if (!operationRequiresAwsChunked(codegenContext, operation)) return@writable
 
                     section.registerInterceptor(runtimeConfig, this) {
                         rustTemplate(
@@ -85,10 +181,12 @@ private fun RuntimeConfig.awsChunked() =
             CargoDependency.Http1x,
             CargoDependency.HttpBody1x,
             CargoDependency.Tracing,
+            CargoDependency.HttpBodyUtil01x.toDevDependency(),
             AwsCargoDependency.awsRuntime(this).withFeature("http-02x"),
             CargoDependency.smithyRuntimeApiClient(this),
             CargoDependency.smithyTypes(this),
             AwsCargoDependency.awsSigv4(this),
             CargoDependency.TempFile.toDevDependency(),
+            CargoDependency.Tokio.toDevDependency(),
         ),
     )

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/IntegrationTestDependencies.kt

@@ -22,6 +22,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Compani
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.Http1x
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.HttpBody1x
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.HttpBodyUtil01x
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.PinProjectLite
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.SerdeJson
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.Smol
 import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency.Companion.TempFile
@@ -181,6 +182,7 @@ class S3TestDependencies(private val runtimeConfig: RuntimeConfig) : LibRsCustom
             addDependency(HdrHistogram)
             addDependency(HttpBody1x.toDevDependency().copy(optional = false))
             addDependency(HttpBodyUtil01x.toDevDependency().copy(optional = false))
+            addDependency(PinProjectLite.toDevDependency())
             addDependency(Smol)
             addDependency(TempFile)
             addDependency(TracingAppender)