Make CustomizableOperation Send and Sync

CHANGELOG.next.toml

@@ -9,4 +9,16 @@
 # message = "Fix typos in module documentation for generated crates"
 # references = ["smithy-rs#920"]
 # meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "client | server | all"}
-# author = "rcoh"
+# author = "rcoh"
+
+[[aws-sdk-rust]]
+message = "`CustomizableOperation`, created as a result of calling the `.customize` method on a fluent builder, ceased to be `Send` and `Sync` in the previous releases. It is now `Send` and `Sync` again."
+references = ["smithy-rs#2944", "smithy-rs#2951"]
+meta = { "breaking" = false, "tada" = false, "bug" = true }
+author = "ysaito1001"
+
+[[smithy-rs]]
+message = "`CustomizableOperation`, created as a result of calling the `.customize` method on a fluent builder, ceased to be `Send` and `Sync` in the previous releases. It is now `Send` and `Sync` again."
+references = ["smithy-rs#2944", "smithy-rs#2951"]
+meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
+author = "ysaito1001"

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt

@@ -155,6 +155,7 @@ class CustomizableOperationGenerator(
                 .resolve("client::interceptors::MapRequestInterceptor"),
             "MutateRequestInterceptor" to RuntimeType.smithyRuntime(runtimeConfig)
                 .resolve("client::interceptors::MutateRequestInterceptor"),
+            "PhantomData" to RuntimeType.Phantom,
             "RuntimePlugin" to RuntimeType.runtimePlugin(runtimeConfig),
             "SharedRuntimePlugin" to RuntimeType.sharedRuntimePlugin(runtimeConfig),
             "SendResult" to ClientRustModule.Client.customize.toType()
@@ -185,14 +186,28 @@ class CustomizableOperationGenerator(
                 rustTemplate(
                     """
                     /// `CustomizableOperation` allows for configuring a single operation invocation before it is sent.
-                    pub struct CustomizableOperation<T, E> {
-                        pub(crate) customizable_send: #{Box}<dyn #{CustomizableSend}<T, E>>,
-                        pub(crate) config_override: #{Option}<crate::config::Builder>,
-                        pub(crate) interceptors: Vec<#{SharedInterceptor}>,
-                        pub(crate) runtime_plugins: Vec<#{SharedRuntimePlugin}>,
+                    pub struct CustomizableOperation<T, E, B> {
+                        customizable_send: B,
+                        config_override: #{Option}<crate::config::Builder>,
+                        interceptors: Vec<#{SharedInterceptor}>,
+                        runtime_plugins: Vec<#{SharedRuntimePlugin}>,
+                        _output: #{PhantomData}<T>,
+                        _error: #{PhantomData}<E>,
                     }
 
-                    impl<T, E> CustomizableOperation<T, E> {
+                    impl<T, E, B> CustomizableOperation<T, E, B> {
+                        /// Creates a new `CustomizableOperation` from `customizable_send`.
+                        pub fn new(customizable_send: B) -> Self {
+                            Self {
+                                customizable_send,
+                                config_override: #{None},
+                                interceptors: vec![],
+                                runtime_plugins: vec![],
+                                _output: #{PhantomData},
+                                _error: #{PhantomData}
+                            }
+                        }
+
                         /// Adds an [`Interceptor`](#{Interceptor}) that runs at specific stages of the request execution pipeline.
                         ///
                         /// Note that interceptors can also be added to `CustomizableOperation` by `config_override`,
@@ -265,6 +280,7 @@ class CustomizableOperationGenerator(
                         ) -> #{SendResult}<T, E>
                         where
                             E: std::error::Error + #{Send} + #{Sync} + 'static,
+                            B: #{CustomizableSend}<T, E>,
                         {
                             let mut config_override = if let Some(config_override) = self.config_override {
                                 config_override
@@ -279,7 +295,7 @@ class CustomizableOperationGenerator(
                                 config_override.push_runtime_plugin(plugin);
                             });
 
-                            (self.customizable_send)(config_override).await
+                            self.customizable_send.send(config_override).await
                         }
 
                         #{additional_methods}
@@ -311,14 +327,12 @@ class CustomizableOperationGenerator(
                     >,
                 >;
 
-                pub trait CustomizableSend<T, E>:
-                    #{FnOnce}(crate::config::Builder) -> BoxFuture<SendResult<T, E>>
-                {}
-
-                impl<F, T, E> CustomizableSend<T, E> for F
-                where
-                    F: #{FnOnce}(crate::config::Builder) -> BoxFuture<SendResult<T, E>>
-                {}
+                pub trait CustomizableSend<T, E>: #{Send} + #{Sync} {
+                    // Takes an owned `self` as the implementation will internally call methods that take `self`.
+                    // If it took `&self`, that would make this trait object safe, but some implementing types do not
+                    // derive `Clone`, unable to yield `self` from `&self`.
+                    fn send(self, config_override: crate::config::Builder) -> BoxFuture<SendResult<T, E>>;
+                }
                 """,
                 *preludeScope,
                 "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig)

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt

@@ -387,6 +387,35 @@ class FluentClientGenerator(
             }
         }
 
+        if (smithyRuntimeMode.generateOrchestrator) {
+            rustTemplate(
+                """
+                impl
+                    crate::client::customize::internal::CustomizableSend<
+                        #{OperationOutput},
+                        #{OperationError},
+                    > for $builderName
+                {
+                    fn send(
+                        self,
+                        config_override: crate::config::Builder,
+                    ) -> crate::client::customize::internal::BoxFuture<
+                        crate::client::customize::internal::SendResult<
+                            #{OperationOutput},
+                            #{OperationError},
+                        >,
+                    > {
+                        #{Box}::pin(async move { self.config_override(config_override).send().await })
+                    }
+                }
+                """,
+                *preludeScope,
+                "OperationError" to errorType,
+                "OperationOutput" to outputType,
+                "SdkError" to RuntimeType.sdkError(runtimeConfig),
+            )
+        }
+
         rustBlockTemplate(
             "impl${generics.inst} $builderName${generics.inst} #{bounds:W}",
             "client" to RuntimeType.smithyClient(runtimeConfig),
@@ -520,22 +549,12 @@ class FluentClientGenerator(
                             #{CustomizableOperation}<
                                 #{OperationOutput},
                                 #{OperationError},
+                                Self,
                             >,
                             #{SdkError}<#{OperationError}>,
                         >
                         {
-                            #{Ok}(#{CustomizableOperation} {
-                                customizable_send: #{Box}::new(move |config_override| {
-                                    #{Box}::pin(async {
-                                        self.config_override(config_override)
-                                            .send()
-                                            .await
-                                    })
-                                }),
-                                config_override: None,
-                                interceptors: vec![],
-                                runtime_plugins: vec![],
-                            })
+                            #{Ok}(#{CustomizableOperation}::new(self))
                         }
                         """,
                         *orchestratorScope,

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGeneratorTest.kt

@@ -0,0 +1,83 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.generators.client
+
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.testutil.TestCodegenSettings
+import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.core.rustlang.rust
+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.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
+import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
+
+class CustomizableOperationGeneratorTest {
+    val model = """
+        namespace com.example
+        use aws.protocols#awsJson1_0
+
+        @awsJson1_0
+        service HelloService {
+            operations: [SayHello],
+            version: "1"
+        }
+
+        @optionalAuth
+        operation SayHello { input: TestInput }
+        structure TestInput {
+           foo: String,
+        }
+    """.asSmithyModel()
+
+    @Test
+    fun `CustomizableOperation is send and sync`() {
+        val test: (ClientCodegenContext, RustCrate) -> Unit = { codegenContext, rustCrate ->
+            rustCrate.integrationTest("customizable_operation_is_send_and_sync") {
+                val moduleName = codegenContext.moduleUseName()
+                rustTemplate(
+                    """
+                    fn check_send_and_sync<T: Send + Sync>(_: T) {}
+
+                    ##[test]
+                    fn test() {
+                        let connector = #{TestConnection}::<#{SdkBody}>::new(Vec::new());
+                        let config = $moduleName::Config::builder()
+                            .endpoint_resolver("http://localhost:1234")
+                            #{set_http_connector}
+                            .build();
+                        let smithy_client = aws_smithy_client::Builder::new()
+                            .connector(connector.clone())
+                            .middleware_fn(|r| r)
+                            .build_dyn();
+                        let client = $moduleName::Client::with_config(smithy_client, config);
+                        check_send_and_sync(client.say_hello().customize());
+                    }
+                    """,
+                    "TestConnection" to CargoDependency.smithyClient(codegenContext.runtimeConfig)
+                        .toDevDependency()
+                        .withFeature("test-util").toType()
+                        .resolve("test_connection::TestConnection"),
+                    "SdkBody" to RuntimeType.sdkBody(codegenContext.runtimeConfig),
+                    "set_http_connector" to writable {
+                        if (codegenContext.smithyRuntimeMode.generateOrchestrator) {
+                            rust(".http_connector(connector.clone())")
+                        }
+                    },
+                )
+            }
+        }
+        clientIntegrationTest(model, TestCodegenSettings.middlewareModeTestParams, test = test)
+        clientIntegrationTest(
+            model,
+            TestCodegenSettings.orchestratorModeTestParams,
+            test = test,
+        )
+    }
+}