Unverified Commit 0beb8905 authored by ysaito1001's avatar ysaito1001 Committed by GitHub
Browse files

Render `CustomizableOperation` per crate in orchestrator (#2726) 

## Motivation and Context
Render `CustomizableOperation` once in the `customizable` module instead
of rendering it in every operation's `builders` module.

## Description
This PR is a follow-up on #2706. The previous PR ported
`CustomizableOperation` in the middleware to the orchestrator but it had
a flaw in that the type was rendered per every operation. This was
clearly a codegen regression because the code for
`CustomizableOperation` is repeated many times and is rendered even if
no customization is requested for an operation.

This PR attempts to fix that problem. Just like `CustomizableOperation`
in the middleware, we render the new `CustomizableOperation` once in the
`customize` module per crate. To accomplish that, the new
`CustomizableOperation` uses a closure, called `customizable_send`, to
encapsulate a call to a fluent builder's `send` method and invokes it
when its own `send` method is called.

~~As a side note, this PR will not take care of the following. We have
[a separate PR](https://github.com/awslabs/smithy-rs/pull/2723

) to fix
them and once we've merged it to this PR, they should be addressed
automatically (there will be merge conflicts, though):~~
- ~~Removing `send_orchestrator_with_plugin` (which is why a type
parameter `R` is present in the code changes, but it'll soon go away)~~
- ~~Incorrect signature of a closure in orchestrator's
`CustomizableOperaton::map_request`~~

## Testing
No new tests added as it's purely refactoring. Confirmed `sra-test`
passed (except for the one that's known to fail).

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Co-authored-by: default avatarYuki Saito <awsaito@amazon.com>
parent 8773a704

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

+21 −26
Original line number Diff line number Diff line
@@ -37,32 +37,7 @@ class CustomizableOperationTestHelpers(runtimeConfig: RuntimeConfig) : 
    override fun section(section: CustomizableOperationSection): Writable =

        writable {

            if (section is CustomizableOperationSection.CustomizableOperationImpl) {

                if (section.operationShape == null) {

                    // TODO(enableNewSmithyRuntime): Delete this branch when middleware is no longer used

                    // This branch customizes CustomizableOperation in the middleware. section.operationShape being

                    // null means that this customization is rendered in a place where we don't need to figure out

                    // the module for an operation (which is the case for CustomizableOperation in the middleware

                    // that is rendered in the customize module).

                    rustTemplate(

                        """

                        ##[doc(hidden)]

                        // This is a temporary method for testing. NEVER use it in production

                        pub fn request_time_for_tests(mut self, request_time: ::std::time::SystemTime) -> Self {

                            self.operation.properties_mut().insert(request_time);

                            self

                        }



                        ##[doc(hidden)]

                        // This is a temporary method for testing. NEVER use it in production

                        pub fn user_agent_for_tests(mut self) -> Self {

                            self.operation.properties_mut().insert(#{AwsUserAgent}::for_tests());

                            self

                        }

                        """,

                        *codegenScope,

                    )

                } else {

                    // The else branch is for rendering customization for the orchestrator.

                if (section.isRuntimeModeOrchestrator) {

                    rustTemplate(

                        """

                        ##[doc(hidden)]
@@ -97,6 +72,26 @@ class CustomizableOperationTestHelpers(runtimeConfig: RuntimeConfig) : 
                        """,

                        *codegenScope,

                    )

                } else {

                    // TODO(enableNewSmithyRuntime): Delete this branch when middleware is no longer used

                    rustTemplate(

                        """

                        ##[doc(hidden)]

                        // This is a temporary method for testing. NEVER use it in production

                        pub fn request_time_for_tests(mut self, request_time: ::std::time::SystemTime) -> Self {

                            self.operation.properties_mut().insert(request_time);

                            self

                        }



                        ##[doc(hidden)]

                        // This is a temporary method for testing. NEVER use it in production

                        pub fn user_agent_for_tests(mut self) -> Self {

                            self.operation.properties_mut().insert(#{AwsUserAgent}::for_tests());

                            self

                        }

                        """,

                        *codegenScope,

                    )

                }

            }

        }

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

+1 −2
Original line number Diff line number Diff line
@@ -5,14 +5,13 @@ 


package software.amazon.smithy.rust.codegen.client.smithy.generators.client



import software.amazon.smithy.model.shapes.OperationShape

import software.amazon.smithy.rust.codegen.core.smithy.customize.NamedCustomization

import software.amazon.smithy.rust.codegen.core.smithy.customize.Section



sealed class CustomizableOperationSection(name: String) : Section(name) {

    /** Write custom code into a customizable operation's impl block */

    data class CustomizableOperationImpl(

        val operationShape: OperationShape?,

        val isRuntimeModeOrchestrator: Boolean,

    ) : CustomizableOperationSection("CustomizableOperationImpl")

}

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

+153 −111
Original line number Diff line number Diff line
@@ -5,20 +5,20 @@ 


package software.amazon.smithy.rust.codegen.client.smithy.generators.client



import software.amazon.smithy.model.shapes.OperationShape

import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext

import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule

import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency

import software.amazon.smithy.rust.codegen.core.rustlang.GenericTypeArg

import software.amazon.smithy.rust.codegen.core.rustlang.RustGenerics

import software.amazon.smithy.rust.codegen.core.rustlang.RustModule

import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter

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.RuntimeType

import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope

import software.amazon.smithy.rust.codegen.core.smithy.RustCrate

import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations

import software.amazon.smithy.rust.codegen.core.util.outputShape



/**

 * Generates the code required to add the `.customize()` function to the
@@ -135,21 +135,18 @@ class CustomizableOperationGenerator( 
            """,

            *codegenScope,

            "additional_methods" to writable {

                writeCustomizations(customizations, CustomizableOperationSection.CustomizableOperationImpl(null))

                writeCustomizations(customizations, CustomizableOperationSection.CustomizableOperationImpl(false))

            },

        )

    }



    fun renderForOrchestrator(writer: RustWriter, operation: OperationShape) {

        val symbolProvider = codegenContext.symbolProvider

        val model = codegenContext.model



        val builderName = operation.fluentBuilderType(symbolProvider).name

        val outputType = symbolProvider.toSymbol(operation.outputShape(model))

        val errorType = symbolProvider.symbolForOperationError(operation)



    fun renderForOrchestrator(crate: RustCrate) {

        val codegenScope = arrayOf(

            *preludeScope,

            "CustomizableOperation" to ClientRustModule.Client.customize.toType()

                .resolve("orchestrator::CustomizableOperation"),

            "CustomizableSend" to ClientRustModule.Client.customize.toType()

                .resolve("internal::CustomizableSend"),

            "HttpRequest" to RuntimeType.smithyRuntimeApi(runtimeConfig)

                .resolve("client::orchestrator::HttpRequest"),

            "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig)
@@ -160,27 +157,42 @@ class CustomizableOperationGenerator( 
                .resolve("client::interceptor::MapRequestInterceptor"),

            "MutateRequestInterceptor" to RuntimeType.smithyRuntime(runtimeConfig)

                .resolve("client::interceptor::MutateRequestInterceptor"),

            "OperationError" to errorType,

            "OperationOutput" to outputType,

            "RuntimePlugin" to RuntimeType.runtimePlugin(runtimeConfig),

            "SendResult" to ClientRustModule.Client.customize.toType()

                .resolve("internal::SendResult"),

            "SdkBody" to RuntimeType.sdkBody(runtimeConfig),

            "SdkError" to RuntimeType.sdkError(runtimeConfig),

            "SharedInterceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig)

                .resolve("client::interceptors::SharedInterceptor"),

        )



        // TODO(enableNewSmithyRuntime): Centralize this struct rather than generating it once per operation

        writer.rustTemplate(

        val customizeModule = ClientRustModule.Client.customize



        crate.withModule(customizeModule) {

            renderConvenienceAliases(customizeModule, this)



            // TODO(enableNewSmithyRuntime): Render it directly under the customize module when CustomizableOperation

            //  in the middleware has been removed.

            withInlineModule(

                RustModule.new(

                    "orchestrator",

                    Visibility.PUBLIC,

                    true,

                    customizeModule,

                    documentationOverride = "Module for defining types for `CustomizableOperation` in the orchestrator",

                ),

                null,

            ) {

                rustTemplate(

                    """

            /// A wrapper type for [`$builderName`]($builderName) that allows for configuring a single

            /// operation invocation.

            pub struct CustomizableOperation {

                pub(crate) fluent_builder: $builderName,

                    /// `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}>,

                    }



            impl CustomizableOperation {

                    impl<T, E> CustomizableOperation<T, E> {

                        /// 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`,
@@ -193,13 +205,13 @@ class CustomizableOperationGenerator( 
                        }



                        /// Allows for customizing the operation's request.

                pub fn map_request<F, E>(mut self, f: F) -> Self

                        pub fn map_request<F, MapE>(mut self, f: F) -> Self

                        where

                    F: #{Fn}(#{HttpRequest}) -> #{Result}<#{HttpRequest}, E>

                            F: #{Fn}(#{HttpRequest}) -> #{Result}<#{HttpRequest}, MapE>

                                + #{Send}

                                + #{Sync}

                                + 'static,

                    E: ::std::error::Error + #{Send} + #{Sync} + 'static,

                            MapE: ::std::error::Error + #{Send} + #{Sync} + 'static,

                        {

                            self.interceptors.push(

                                #{SharedInterceptor}::new(
@@ -242,14 +254,11 @@ class CustomizableOperationGenerator( 


                        /// Sends the request and returns the response.

                        pub async fn send(

                    self

                ) -> #{Result}<

                    #{OperationOutput},

                    #{SdkError}<

                        #{OperationError},

                        #{HttpResponse}

                    >

                > {

                            self,

                        ) -> #{SendResult}<T, E>

                        where

                            E: std::error::Error + #{Send} + #{Sync} + 'static,

                        {

                            let mut config_override = if let Some(config_override) = self.config_override {

                                config_override

                            } else {
@@ -260,10 +269,7 @@ class CustomizableOperationGenerator( 
                                config_override.add_interceptor(interceptor);

                            });



                    self.fluent_builder

                        .config_override(config_override)

                        .send_orchestrator()

                        .await

                            (self.customizable_send)(config_override).await

                        }



                        #{additional_methods}
@@ -271,11 +277,47 @@ class CustomizableOperationGenerator( 
                    """,

                    *codegenScope,

                    "additional_methods" to writable {

                writeCustomizations(customizations, CustomizableOperationSection.CustomizableOperationImpl(operation))

                        writeCustomizations(

                            customizations,

                            CustomizableOperationSection.CustomizableOperationImpl(true),

                        )

                    },

                )

            }

        }

    }



    private fun renderConvenienceAliases(parentModule: RustModule, writer: RustWriter) {

        writer.withInlineModule(RustModule.new("internal", Visibility.PUBCRATE, true, parentModule), null) {

            rustTemplate(

                """

                pub type BoxFuture<T> = ::std::pin::Pin<#{Box}<dyn ::std::future::Future<Output = T> + #{Send}>>;



                pub type SendResult<T, E> = #{Result}<

                    T,

                    #{SdkError}<

                        E,

                        #{HttpResponse},

                    >,

                >;



                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>>

                {}

                """,

                *preludeScope,

                "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig)

                    .resolve("client::orchestrator::HttpResponse"),

                "SdkError" to RuntimeType.sdkError(runtimeConfig),

            )

        }

    }

}



fun renderCustomizableOperationSend(codegenContext: ClientCodegenContext, generics: FluentClientGenerics, writer: RustWriter) {

    val runtimeConfig = codegenContext.runtimeConfig

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

+35 −11
Original line number Diff line number Diff line
@@ -88,13 +88,13 @@ class FluentClientGenerator( 
        operations.forEach { operation ->

            crate.withModule(symbolProvider.moduleForBuilder(operation)) {

                renderFluentBuilder(operation)

                if (codegenContext.smithyRuntimeMode.generateOrchestrator) {

                    customizableOperationGenerator.renderForOrchestrator(this, operation)

                }

            }

        }



        customizableOperationGenerator.render(crate)

        if (codegenContext.smithyRuntimeMode.generateOrchestrator) {

            customizableOperationGenerator.renderForOrchestrator(crate)

        }

    }



    private fun renderFluentClient(crate: RustCrate) {
@@ -450,13 +450,15 @@ class FluentClientGenerator( 
            if (smithyRuntimeMode.generateOrchestrator) {

                val orchestratorScope = arrayOf(

                    *preludeScope,

                    "CustomizableOperation" to symbolProvider.moduleForBuilder(operation).toType()

                        .resolve("CustomizableOperation"),

                    "CustomizableOperation" to ClientRustModule.Client.customize.toType()

                        .resolve("orchestrator::CustomizableOperation"),

                    "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig)

                        .resolve("client::orchestrator::HttpResponse"),

                    "Operation" to operationSymbol,

                    "OperationError" to errorType,

                    "OperationOutput" to outputType,

                    "SendResult" to ClientRustModule.Client.customize.toType()

                        .resolve("internal::SendResult"),

                    "SdkError" to RuntimeType.sdkError(runtimeConfig),

                )

                rustTemplate(
@@ -469,8 +471,24 @@ class FluentClientGenerator( 


                    ##[doc(hidden)]

                    // TODO(enableNewSmithyRuntime): Remove `async` once we switch to orchestrator

                    pub async fn customize_orchestrator(self) -> #{CustomizableOperation} {

                        #{CustomizableOperation} { fluent_builder: self, config_override: None, interceptors: vec![] }

                    pub async fn customize_orchestrator(

                        self,

                    ) -> #{CustomizableOperation}<

                        #{OperationOutput},

                        #{OperationError},

                    >

                    {

                        #{CustomizableOperation} {

                            customizable_send: #{Box}::new(move |config_override| {

                                #{Box}::pin(async {

                                    self.config_override(config_override)

                                        .send_orchestrator()

                                        .await

                                })

                            }),

                            config_override: None,

                            interceptors: vec![],

                        }

                    }

                    """,

                    *orchestratorScope,
@@ -493,10 +511,16 @@ class FluentClientGenerator( 
                        /// Consumes this builder, creating a customizable operation that can be modified before being

                        /// sent.

                        // TODO(enableNewSmithyRuntime): Remove `async` and `Result` once we switch to orchestrator

                        pub async fn customize(self) -> #{Result}<

                            #{CustomizableOperation},

                            #{SdkError}<#{OperationError}>

                        > {

                        pub async fn customize(

                            self,

                        ) -> #{Result}<

                            #{CustomizableOperation}<

                                #{OperationOutput},

                                #{OperationError},

                            >,

                            #{SdkError}<#{OperationError}>,

                        >

                        {

                            #{Ok}(self.customize_orchestrator().await)

                        }

                        """,