Set a default for idempotency_token in the service runtime plugin

This commit also refactors idempotency tokens to be centralized in a decorator
smithy-lang · Oct 17, 2023 · c0b5d7e · c0b5d7e
1 parent 149fc7a
commit c0b5d7e
Show file tree

Hide file tree

Showing 11 changed files with 113 additions and 40 deletions.
diff --git a/.../main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt b/.../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/...ware/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenDecorator.kt b/...ware/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<ConfigCustomization>,
+    ): List<ConfigCustomization> = baseCustomizations.extendIf(enabled(codegenContext)) {
+        IdempotencyTokenProviderCustomization(codegenContext)
+    }
+
+    override fun operationCustomizations(
+        codegenContext: ClientCodegenContext,
+        operation: OperationShape,
+        baseCustomizations: List<OperationCustomization>,
+    ): List<OperationCustomization> {
+        return baseCustomizations + IdempotencyTokenGenerator(codegenContext, operation)
+    }
+
+    override fun serviceRuntimePluginCustomizations(
+        codegenContext: ClientCodegenContext,
+        baseCustomizations: List<ServiceRuntimePluginCustomization>,
+    ): List<ServiceRuntimePluginCustomization> {
+        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/...ware/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt b/...ware/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/...lin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt b/...lin/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<OperationCustomization> =
         baseCustomizations +
             MetadataCustomization(codegenContext, operation) +
-            IdempotencyTokenGenerator(codegenContext, operation) +
             HttpChecksumRequiredGenerator(codegenContext, operation) +
             RetryClassifierOperationCustomization(codegenContext, operation)
 

diff --git a/...azon/smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGenerator.kt b/...azon/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/...in/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt b/...in/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<OperationCustomization>,
         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/...ware/amazon/smithy/rust/codegen/client/smithy/generators/ServiceRuntimePluginGenerator.kt b/...ware/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,31 +81,41 @@ 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"),
         )
     }
 
     fun render(
         writer: RustWriter,
         customizations: List<ServiceRuntimePluginCustomization>,
     ) {
+        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/...thy/rust/codegen/client/smithy/generators/config/IdempotencyTokenProviderCustomization.kt b/...thy/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/...ware/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGenerator.kt b/...ware/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGenerator.kt
@@ -257,9 +257,6 @@ class ServiceConfigGenerator(
             extraCustomizations: List<ConfigCustomization>,
         ): ServiceConfigGenerator {
             val baseFeatures = mutableListOf<ConfigCustomization>()
-            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
@@ -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
@@ -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