From c0b5d7e427ef8e31cd95d053f851599b381200df Mon Sep 17 00:00:00 2001 From: Russell Cohen Date: Tue, 17 Oct 2023 10:40:29 -0400 Subject: [PATCH] Set a default for idempotency_token in the service runtime plugin This commit also refactors idempotency tokens to be centralized in a decorator --- .../client/smithy/RustClientCodegenPlugin.kt | 2 + .../IdempotencyTokenDecorator.kt | 60 +++++++++++++++++++ .../IdempotencyTokenGenerator.kt | 3 +- .../customize/RequiredCustomizations.kt | 2 - .../ConfigOverrideRuntimePluginGenerator.kt | 3 +- .../generators/OperationCustomization.kt | 4 +- .../ServiceRuntimePluginGenerator.kt | 41 ++++++++++++- .../IdempotencyTokenProviderCustomization.kt | 25 -------- .../config/ServiceConfigGenerator.kt | 3 - .../src/client_idempotency_token.rs | 8 ++- .../inlineable/src/idempotency_token.rs | 2 +- 11 files changed, 113 insertions(+), 40 deletions(-) create mode 100644 codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenDecorator.kt diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt index 39bc4633f2c..08814fadac1 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt @@ -12,6 +12,7 @@ import software.amazon.smithy.model.shapes.ServiceShape import software.amazon.smithy.rust.codegen.client.smithy.customizations.ClientCustomizations import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpAuthDecorator import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpConnectorConfigDecorator +import software.amazon.smithy.rust.codegen.client.smithy.customizations.IdempotencyTokenDecorator import software.amazon.smithy.rust.codegen.client.smithy.customizations.NoAuthDecorator import software.amazon.smithy.rust.codegen.client.smithy.customizations.SensitiveOutputDecorator import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator @@ -66,6 +67,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() { HttpAuthDecorator(), HttpConnectorConfigDecorator(), SensitiveOutputDecorator(), + IdempotencyTokenDecorator(), *decorator, ) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenDecorator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenDecorator.kt new file mode 100644 index 00000000000..b33ee8fbbc2 --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenDecorator.kt @@ -0,0 +1,60 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.rust.codegen.client.smithy.customizations + +import software.amazon.smithy.model.shapes.OperationShape +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.ServiceRuntimePluginCustomization +import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection +import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization +import software.amazon.smithy.rust.codegen.client.smithy.generators.config.IdempotencyTokenProviderCustomization +import software.amazon.smithy.rust.codegen.client.smithy.generators.config.needsIdempotencyToken +import software.amazon.smithy.rust.codegen.core.rustlang.rust +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.util.extendIf + +class IdempotencyTokenDecorator : ClientCodegenDecorator { + override val name: String = "IdempotencyToken" + override val order: Byte = 0 + + private fun enabled(ctx: ClientCodegenContext) = ctx.serviceShape.needsIdempotencyToken(ctx.model) + override fun configCustomizations( + codegenContext: ClientCodegenContext, + baseCustomizations: List, + ): List = baseCustomizations.extendIf(enabled(codegenContext)) { + IdempotencyTokenProviderCustomization(codegenContext) + } + + override fun operationCustomizations( + codegenContext: ClientCodegenContext, + operation: OperationShape, + baseCustomizations: List, + ): List { + return baseCustomizations + IdempotencyTokenGenerator(codegenContext, operation) + } + + override fun serviceRuntimePluginCustomizations( + codegenContext: ClientCodegenContext, + baseCustomizations: List, + ): List { + return baseCustomizations.extendIf(enabled(codegenContext)) { + object : ServiceRuntimePluginCustomization() { + override fun section(section: ServiceRuntimePluginSection) = writable { + if (section is ServiceRuntimePluginSection.AdditionalConfig) { + section.putConfigValue(this, defaultTokenProvider((codegenContext.runtimeConfig))) + } + } + } + } + } +} + +private fun defaultTokenProvider(runtimeConfig: RuntimeConfig) = + writable { rust("#T()", RuntimeType.idempotencyToken(runtimeConfig).resolve("default_provider")) } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt index dd1bcfef08f..cfcb11e8193 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt @@ -47,12 +47,13 @@ class IdempotencyTokenGenerator( "/inlineable/src/client_idempotency_token.rs", CargoDependency.smithyRuntimeApi(runtimeConfig), CargoDependency.smithyTypes(runtimeConfig), + InlineDependency.idempotencyToken(runtimeConfig), ).toType().resolve("IdempotencyTokenRuntimePlugin"), ) return when (section) { is OperationSection.AdditionalRuntimePlugins -> writable { - section.addOperationRuntimePlugin(this) { + section.addClientPlugin(this) { if (symbolProvider.toSymbol(idempotencyTokenMember).isOptional()) { // An idempotency token is optional. If the user didn't specify a token // then we'll generate one and set it. diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt index a1f11cf4a26..2123d802868 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt @@ -10,7 +10,6 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule import software.amazon.smithy.rust.codegen.client.smithy.customizations.ConnectionPoisoningRuntimePluginCustomization import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpChecksumRequiredGenerator -import software.amazon.smithy.rust.codegen.client.smithy.customizations.IdempotencyTokenGenerator import software.amazon.smithy.rust.codegen.client.smithy.customizations.InterceptorConfigCustomization import software.amazon.smithy.rust.codegen.client.smithy.customizations.MetadataCustomization import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyConfigCustomization @@ -50,7 +49,6 @@ class RequiredCustomizations : ClientCodegenDecorator { ): List = baseCustomizations + MetadataCustomization(codegenContext, operation) + - IdempotencyTokenGenerator(codegenContext, operation) + HttpChecksumRequiredGenerator(codegenContext, operation) + RetryClassifierOperationCustomization(codegenContext, operation) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGenerator.kt index aa3566edabc..e83ba1e0e87 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGenerator.kt @@ -56,7 +56,8 @@ class ConfigOverrideRuntimePluginGenerator( ) -> Self { let mut layer = config_override.config; let mut components = config_override.runtime_components; - let resolver = #{Resolver}::overrid(initial_config, initial_components, &mut layer, &mut components); + ##[allow(unused_mut)] + let mut resolver = #{Resolver}::overrid(initial_config, initial_components, &mut layer, &mut components); #{config} diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt index db534910ef5..176b85b1c8f 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt @@ -95,8 +95,8 @@ sealed class OperationSection(name: String) : Section(name) { override val customizations: List, val operationShape: OperationShape, ) : OperationSection("AdditionalRuntimePlugins") { - fun addServiceRuntimePlugin(writer: RustWriter, plugin: Writable) { - writer.rustTemplate(".with_service_plugin(#{plugin})", "plugin" to plugin) + fun addClientPlugin(writer: RustWriter, plugin: Writable) { + writer.rustTemplate(".with_client_plugin(#{plugin})", "plugin" to plugin) } fun addOperationRuntimePlugin(writer: RustWriter, plugin: Writable) { diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ServiceRuntimePluginGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ServiceRuntimePluginGenerator.kt index afcd491ae3d..60015d897cc 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ServiceRuntimePluginGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ServiceRuntimePluginGenerator.kt @@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.client.smithy.generators import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter import software.amazon.smithy.rust.codegen.core.rustlang.Writable +import software.amazon.smithy.rust.codegen.core.rustlang.isNotEmpty 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 @@ -16,6 +17,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.pre import software.amazon.smithy.rust.codegen.core.smithy.customize.NamedCustomization import software.amazon.smithy.rust.codegen.core.smithy.customize.Section import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations +import software.amazon.smithy.rust.codegen.core.util.dq sealed class ServiceRuntimePluginSection(name: String) : Section(name) { /** @@ -25,6 +27,16 @@ sealed class ServiceRuntimePluginSection(name: String) : Section(name) { */ class DeclareSingletons : ServiceRuntimePluginSection("DeclareSingletons") + /** + * Hook for adding additional things to config inside service runtime plugins. + */ + data class AdditionalConfig(val newLayerName: String, val serviceConfigName: String) : ServiceRuntimePluginSection("AdditionalConfig") { + /** Adds a value to the config bag */ + fun putConfigValue(writer: RustWriter, value: Writable) { + writer.rust("$newLayerName.store_put(#T);", value) + } + } + data class RegisterRuntimeComponents(val serviceConfigName: String) : ServiceRuntimePluginSection("RegisterRuntimeComponents") { /** Generates the code to register an interceptor */ fun registerInterceptor(writer: RustWriter, interceptor: Writable) { @@ -69,6 +81,7 @@ class ServiceRuntimePluginGenerator( "Layer" to smithyTypes.resolve("config_bag::Layer"), "RuntimeComponentsBuilder" to RuntimeType.runtimeComponentsBuilder(rc), "RuntimePlugin" to RuntimeType.runtimePlugin(rc), + "Order" to runtimeApi.resolve("client::runtime_plugin::Order"), ) } @@ -76,24 +89,33 @@ class ServiceRuntimePluginGenerator( writer: RustWriter, customizations: List, ) { + val additionalConfig = writable { + writeCustomizations(customizations, ServiceRuntimePluginSection.AdditionalConfig("cfg", "_service_config")) + } writer.rustTemplate( """ ##[derive(::std::fmt::Debug)] pub(crate) struct ServiceRuntimePlugin { + config: #{Option}<#{FrozenLayer}>, runtime_components: #{RuntimeComponentsBuilder}, } impl ServiceRuntimePlugin { pub fn new(_service_config: crate::config::Config) -> Self { + let config = { #{config} }; let mut runtime_components = #{RuntimeComponentsBuilder}::new("ServiceRuntimePlugin"); #{runtime_components} - Self { runtime_components } + Self { config, runtime_components } } } impl #{RuntimePlugin} for ServiceRuntimePlugin { fn config(&self) -> #{Option}<#{FrozenLayer}> { - None + self.config.clone() + } + + fn order(&self) -> #{Order} { + #{Order}::Defaults } fn runtime_components(&self, _: &#{RuntimeComponentsBuilder}) -> #{Cow}<'_, #{RuntimeComponentsBuilder}> { @@ -105,6 +127,21 @@ class ServiceRuntimePluginGenerator( #{declare_singletons} """, *codegenScope, + "config" to writable { + if (additionalConfig.isNotEmpty()) { + rustTemplate( + """ + let mut cfg = #{Layer}::new(${codegenContext.serviceShape.id.name.dq()}); + #{additional_config} + #{Some}(cfg.freeze()) + """, + *codegenScope, + "additional_config" to additionalConfig, + ) + } else { + rust("None") + } + }, "runtime_components" to writable { writeCustomizations(customizations, ServiceRuntimePluginSection.RegisterRuntimeComponents("_service_config")) }, diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/IdempotencyTokenProviderCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/IdempotencyTokenProviderCustomization.kt index 0408c72d8e6..bac532cd6ba 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/IdempotencyTokenProviderCustomization.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/IdempotencyTokenProviderCustomization.kt @@ -27,20 +27,6 @@ class IdempotencyTokenProviderCustomization(codegenContext: ClientCodegenContext override fun section(section: ServiceConfig): Writable { return when (section) { - ServiceConfig.ConfigImpl -> writable { - rustTemplate( - """ - /// Returns a copy of the idempotency token provider. - /// If a random token provider was configured, - /// a newly-randomized token provider will be returned. - pub fn idempotency_token_provider(&self) -> #{IdempotencyTokenProvider} { - self.config.load::<#{IdempotencyTokenProvider}>().expect("the idempotency provider should be set").clone() - } - """, - *codegenScope, - ) - } - ServiceConfig.BuilderImpl -> writable { rustTemplate( """ @@ -65,17 +51,6 @@ class IdempotencyTokenProviderCustomization(codegenContext: ClientCodegenContext ) } - ServiceConfig.BuilderBuild -> writable { - rustTemplate( - """ - if !resolver.is_set::<#{IdempotencyTokenProvider}>() { - resolver.config_mut().store_put(#{default_provider}()); - } - """, - *codegenScope, - ) - } - is ServiceConfig.DefaultForTests -> writable { rust("""${section.configBuilderRef}.set_idempotency_token_provider(Some("00000000-0000-4000-8000-000000000000".into()));""") } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGenerator.kt index 07a267786d8..bc3eea97622 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGenerator.kt @@ -257,9 +257,6 @@ class ServiceConfigGenerator( extraCustomizations: List, ): ServiceConfigGenerator { val baseFeatures = mutableListOf() - if (codegenContext.serviceShape.needsIdempotencyToken(codegenContext.model)) { - baseFeatures.add(IdempotencyTokenProviderCustomization(codegenContext)) - } return ServiceConfigGenerator(codegenContext, baseFeatures + extraCustomizations) } } diff --git a/rust-runtime/inlineable/src/client_idempotency_token.rs b/rust-runtime/inlineable/src/client_idempotency_token.rs index 6d857471f75..40f361278d1 100644 --- a/rust-runtime/inlineable/src/client_idempotency_token.rs +++ b/rust-runtime/inlineable/src/client_idempotency_token.rs @@ -3,7 +3,9 @@ * SPDX-License-Identifier: Apache-2.0 */ -use crate::idempotency_token::IdempotencyTokenProvider; +use std::borrow::Cow; +use std::fmt; + use aws_smithy_runtime_api::box_error::BoxError; use aws_smithy_runtime_api::client::interceptors::context::{ BeforeSerializationInterceptorContextMut, Input, @@ -14,8 +16,8 @@ use aws_smithy_runtime_api::client::runtime_components::{ }; use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin; use aws_smithy_types::config_bag::ConfigBag; -use std::borrow::Cow; -use std::fmt; + +use crate::idempotency_token::IdempotencyTokenProvider; #[derive(Debug)] pub(crate) struct IdempotencyTokenRuntimePlugin { diff --git a/rust-runtime/inlineable/src/idempotency_token.rs b/rust-runtime/inlineable/src/idempotency_token.rs index 12921f02af8..69425abc80b 100644 --- a/rust-runtime/inlineable/src/idempotency_token.rs +++ b/rust-runtime/inlineable/src/idempotency_token.rs @@ -31,7 +31,7 @@ pub(crate) fn uuid_v4(input: u128) -> String { out } -/// IdempotencyTokenProvider generates idempotency tokens for idempotency API requests +/// IdempotencyTokenProvider generates idempotency tokens for idempotent API requests /// /// Generally, customers will not need to interact with this at all. A sensible default will be /// provided automatically during config construction. However, if you need deterministic behavior