From 0beb8905013248d447b30ebf6cea607d09592c21 Mon Sep 17 00:00:00 2001 From: ysaito1001 Date: Fri, 26 May 2023 11:17:56 -0500 Subject: [PATCH] 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: Yuki Saito --- .../AwsCustomizableOperationDecorator.kt | 47 ++-- .../client/CustomizableOperationDecorator.kt | 3 +- .../client/CustomizableOperationGenerator.kt | 264 ++++++++++-------- .../client/FluentClientGenerator.kt | 46 ++- 4 files changed, 210 insertions(+), 150 deletions(-) diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt index 70cc99e52..884fefc82 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCustomizableOperationDecorator.kt @@ -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, + ) } } } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt index 93a01e41b..8859cc598 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationDecorator.kt @@ -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") } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt index ee4b0f3c7..eab249a9c 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationGenerator.kt @@ -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,120 +157,165 @@ 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( - """ - /// A wrapper type for [`$builderName`]($builderName) that allows for configuring a single - /// operation invocation. - pub struct CustomizableOperation { - pub(crate) fluent_builder: $builderName, - pub(crate) config_override: #{Option}, - pub(crate) interceptors: Vec<#{SharedInterceptor}>, - } + val customizeModule = ClientRustModule.Client.customize - impl CustomizableOperation { - /// 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`, - /// `map_request`, and `mutate_request` (the last two are implemented via interceptors under the hood). - /// The order in which those user-specified operation interceptors are invoked should not be relied upon - /// as it is an implementation detail. - pub fn interceptor(mut self, interceptor: impl #{Interceptor} + #{Send} + #{Sync} + 'static) -> Self { - self.interceptors.push(#{SharedInterceptor}::new(interceptor)); - self - } + crate.withModule(customizeModule) { + renderConvenienceAliases(customizeModule, this) - /// Allows for customizing the operation's request. - pub fn map_request(mut self, f: F) -> Self - where - F: #{Fn}(#{HttpRequest}) -> #{Result}<#{HttpRequest}, E> - + #{Send} - + #{Sync} - + 'static, - E: ::std::error::Error + #{Send} + #{Sync} + 'static, - { - self.interceptors.push( - #{SharedInterceptor}::new( - #{MapRequestInterceptor}::new(f), - ), - ); - self - } + // 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( + """ + /// `CustomizableOperation` allows for configuring a single operation invocation before it is sent. + pub struct CustomizableOperation { + pub(crate) customizable_send: #{Box}>, + pub(crate) config_override: #{Option}, + pub(crate) interceptors: Vec<#{SharedInterceptor}>, + } - /// Convenience for `map_request` where infallible direct mutation of request is acceptable. - pub fn mutate_request(mut self, f: F) -> Self - where - F: #{Fn}(&mut http::Request<#{SdkBody}>) + #{Send} + #{Sync} + 'static, - { - self.interceptors.push( - #{SharedInterceptor}::new( - #{MutateRequestInterceptor}::new(f), - ), - ); - self - } + impl CustomizableOperation { + /// 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`, + /// `map_request`, and `mutate_request` (the last two are implemented via interceptors under the hood). + /// The order in which those user-specified operation interceptors are invoked should not be relied upon + /// as it is an implementation detail. + pub fn interceptor(mut self, interceptor: impl #{Interceptor} + #{Send} + #{Sync} + 'static) -> Self { + self.interceptors.push(#{SharedInterceptor}::new(interceptor)); + self + } - /// Overrides config for a single operation invocation. - /// - /// `config_override` is applied to the operation configuration level. - /// The fields in the builder that are `Some` override those applied to the service - /// configuration level. For instance, - /// - /// Config A overridden by Config B == Config C - /// field_1: None, field_1: Some(v2), field_1: Some(v2), - /// field_2: Some(v1), field_2: Some(v2), field_2: Some(v2), - /// field_3: Some(v1), field_3: None, field_3: Some(v1), - pub fn config_override( - mut self, - config_override: impl #{Into}, - ) -> Self { - self.config_override = Some(config_override.into()); - self - } + /// Allows for customizing the operation's request. + pub fn map_request(mut self, f: F) -> Self + where + F: #{Fn}(#{HttpRequest}) -> #{Result}<#{HttpRequest}, MapE> + + #{Send} + + #{Sync} + + 'static, + MapE: ::std::error::Error + #{Send} + #{Sync} + 'static, + { + self.interceptors.push( + #{SharedInterceptor}::new( + #{MapRequestInterceptor}::new(f), + ), + ); + self + } - /// Sends the request and returns the response. - pub async fn send( - self - ) -> #{Result}< - #{OperationOutput}, - #{SdkError}< - #{OperationError}, - #{HttpResponse} - > - > { - let mut config_override = if let Some(config_override) = self.config_override { - config_override - } else { - crate::config::Builder::new() - }; - - self.interceptors.into_iter().for_each(|interceptor| { - config_override.add_interceptor(interceptor); - }); - - self.fluent_builder - .config_override(config_override) - .send_orchestrator() - .await - } + /// Convenience for `map_request` where infallible direct mutation of request is acceptable. + pub fn mutate_request(mut self, f: F) -> Self + where + F: #{Fn}(&mut http::Request<#{SdkBody}>) + #{Send} + #{Sync} + 'static, + { + self.interceptors.push( + #{SharedInterceptor}::new( + #{MutateRequestInterceptor}::new(f), + ), + ); + self + } - #{additional_methods} + /// Overrides config for a single operation invocation. + /// + /// `config_override` is applied to the operation configuration level. + /// The fields in the builder that are `Some` override those applied to the service + /// configuration level. For instance, + /// + /// Config A overridden by Config B == Config C + /// field_1: None, field_1: Some(v2), field_1: Some(v2), + /// field_2: Some(v1), field_2: Some(v2), field_2: Some(v2), + /// field_3: Some(v1), field_3: None, field_3: Some(v1), + pub fn config_override( + mut self, + config_override: impl #{Into}, + ) -> Self { + self.config_override = Some(config_override.into()); + self + } + + /// Sends the request and returns the response. + pub async fn send( + self, + ) -> #{SendResult} + where + E: std::error::Error + #{Send} + #{Sync} + 'static, + { + let mut config_override = if let Some(config_override) = self.config_override { + config_override + } else { + crate::config::Builder::new() + }; + + self.interceptors.into_iter().for_each(|interceptor| { + config_override.add_interceptor(interceptor); + }); + + (self.customizable_send)(config_override).await + } + + #{additional_methods} + } + """, + *codegenScope, + "additional_methods" to writable { + writeCustomizations( + customizations, + CustomizableOperationSection.CustomizableOperationImpl(true), + ) + }, + ) } - """, - *codegenScope, - "additional_methods" to writable { - writeCustomizations(customizations, CustomizableOperationSection.CustomizableOperationImpl(operation)) - }, - ) + } + } + + private fun renderConvenienceAliases(parentModule: RustModule, writer: RustWriter) { + writer.withInlineModule(RustModule.new("internal", Visibility.PUBCRATE, true, parentModule), null) { + rustTemplate( + """ + pub type BoxFuture = ::std::pin::Pin<#{Box} + #{Send}>>; + + pub type SendResult = #{Result}< + T, + #{SdkError}< + E, + #{HttpResponse}, + >, + >; + + pub trait CustomizableSend: + #{FnOnce}(crate::config::Builder) -> BoxFuture> + {} + + impl CustomizableSend for F + where + F: #{FnOnce}(crate::config::Builder) -> BoxFuture> + {} + """, + *preludeScope, + "HttpResponse" to RuntimeType.smithyRuntimeApi(runtimeConfig) + .resolve("client::orchestrator::HttpResponse"), + "SdkError" to RuntimeType.sdkError(runtimeConfig), + ) + } } } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt index 40e43425c..546197947 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt @@ -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) } """, -- GitLab