Move default endpoint resolver into a runtime plugin

CHANGELOG.next.toml

@@ -408,3 +408,9 @@ message = "`RuntimeComponents` have been added as an argument to the `IdentityRe
 references = ["smithy-rs#2917"]
 meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "client"}
 author = "jdisanti"
+
+[[smithy-rs]]
+message = "The `idempotency_provider` field has been removed from config as a public field. If you need access to this field, it is still available from the context of an interceptor."
+references = ["smithy-rs#3072"]
+meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "client" }
+author = "rcoh"

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

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<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")) }

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/IdempotencyTokenGenerator.kt

@@ -20,12 +20,13 @@ import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.smithy.isOptional
+import software.amazon.smithy.rust.codegen.core.util.UNREACHABLE
 import software.amazon.smithy.rust.codegen.core.util.findMemberWithTrait
 import software.amazon.smithy.rust.codegen.core.util.inputShape
 
 class IdempotencyTokenGenerator(
     codegenContext: CodegenContext,
-    operationShape: OperationShape,
+    private val operationShape: OperationShape,
 ) : OperationCustomization() {
     private val model = codegenContext.model
     private val runtimeConfig = codegenContext.runtimeConfig
@@ -47,44 +48,32 @@ 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) {
-                    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.
-                        rustTemplate(
-                            """
-                            #{IdempotencyTokenRuntimePlugin}::new(|token_provider, input| {
-                                let input: &mut #{Input} = input.downcast_mut().expect("correct type");
-                                if input.$memberName.is_none() {
-                                    input.$memberName = #{Some}(token_provider.make_idempotency_token());
-                                }
-                            })
-                            """,
-                            *codegenScope,
-                        )
-                    } else {
-                        // An idempotency token is required, but it'll be set to an empty string if
-                        // the user didn't specify one. If that's the case, then we'll generate one
-                        // and set it.
-                        rustTemplate(
-                            """
-                            #{IdempotencyTokenRuntimePlugin}::new(|token_provider, input| {
-                                let input: &mut #{Input} = input.downcast_mut().expect("correct type");
-                                if input.$memberName.is_empty() {
-                                    input.$memberName = token_provider.make_idempotency_token();
-                                }
-                            })
-                            """,
-                            *codegenScope,
-                        )
+                    if (!symbolProvider.toSymbol(idempotencyTokenMember).isOptional()) {
+                        UNREACHABLE("top level input members are always optional. $operationShape")
                     }
+                    // An idempotency token is optional. If the user didn't specify a token
+                    // then we'll generate one and set it.
+                    rustTemplate(
+                        """
+                        #{IdempotencyTokenRuntimePlugin}::new(|token_provider, input| {
+                            let input: &mut #{Input} = input.downcast_mut().expect("correct type");
+                            if input.$memberName.is_none() {
+                                input.$memberName = #{Some}(token_provider.make_idempotency_token());
+                            }
+                        })
+                        """,
+                        *codegenScope,
+                    )
                 }
             }
+
             else -> emptySection
         }
     }

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<OperationCustomization> =
         baseCustomizations +
             MetadataCustomization(codegenContext, operation) +
-            IdempotencyTokenGenerator(codegenContext, operation) +
             HttpChecksumRequiredGenerator(codegenContext, operation) +
             RetryClassifierOperationCustomization(codegenContext, operation)
 

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointConfigCustomization.kt

@@ -10,6 +10,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientRustModule
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+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
@@ -153,84 +154,17 @@ internal class EndpointConfigCustomization(
                     rustTemplate(
                         """
                         pub fn set_endpoint_resolver(&mut self, endpoint_resolver: #{Option}<$sharedEndpointResolver>) -> &mut Self {
-                            self.config.store_or_unset(endpoint_resolver);
+                            self.runtime_components.set_endpoint_resolver(endpoint_resolver.map(|r|#{wrap_resolver}));
                             self
                         }
                         """,
                         *codegenScope,
-                    )
-                }
-
-                ServiceConfig.BuilderBuild -> {
-                    rustTemplate(
-                        "#{set_endpoint_resolver}(&mut resolver);",
-                        "set_endpoint_resolver" to setEndpointResolverFn(),
-                    )
-                }
-
-                is ServiceConfig.OperationConfigOverride -> {
-                    rustTemplate(
-                        "#{set_endpoint_resolver}(&mut resolver);",
-                        "set_endpoint_resolver" to setEndpointResolverFn(),
+                        "wrap_resolver" to codegenContext.wrapResolver { rust("r") },
                     )
                 }
 
                 else -> emptySection
             }
         }
     }
-
-    private fun defaultResolver(): RuntimeType {
-        // For now, fallback to a default endpoint resolver that always fails. In the future,
-        // the endpoint resolver will be required (so that it can be unwrapped).
-        return typesGenerator.defaultResolver() ?: RuntimeType.forInlineFun(
-            "MissingResolver",
-            ClientRustModule.Config.endpoint,
-        ) {
-            rustTemplate(
-                """
-                ##[derive(Debug)]
-                pub(crate) struct MissingResolver;
-                impl MissingResolver {
-                    pub(crate) fn new() -> Self { Self }
-                }
-                impl<T> #{ResolveEndpoint}<T> for MissingResolver {
-                    fn resolve_endpoint(&self, _params: &T) -> #{Result} {
-                        Err(#{ResolveEndpointError}::message("an endpoint resolver must be provided."))
-                    }
-                }
-                """,
-                "ResolveEndpoint" to types.resolveEndpoint,
-                "ResolveEndpointError" to types.resolveEndpointError,
-                "Result" to types.smithyHttpEndpointModule.resolve("Result"),
-            )
-        }
-    }
-
-    private fun setEndpointResolverFn(): RuntimeType = RuntimeType.forInlineFun("set_endpoint_resolver", ClientRustModule.config) {
-        // TODO(enableNewSmithyRuntimeCleanup): Simplify the endpoint resolvers
-        rustTemplate(
-            """
-            fn set_endpoint_resolver(resolver: &mut #{Resolver}<'_>) {
-                let endpoint_resolver = if resolver.is_initial() {
-                    Some(resolver.resolve_config::<#{OldSharedEndpointResolver}<#{Params}>>().cloned().unwrap_or_else(||
-                        #{OldSharedEndpointResolver}::new(#{DefaultResolver}::new())
-                    ))
-                } else if resolver.is_latest_set::<#{OldSharedEndpointResolver}<#{Params}>>() {
-                    resolver.resolve_config::<#{OldSharedEndpointResolver}<#{Params}>>().cloned()
-                } else {
-                    None
-                };
-                if let Some(endpoint_resolver) = endpoint_resolver {
-                    let shared = #{SharedEndpointResolver}::new(
-                        #{DefaultEndpointResolver}::<#{Params}>::new(endpoint_resolver)
-                    );
-                    resolver.runtime_components_mut().set_endpoint_resolver(#{Some}(shared));
-                }
-            }
-            """,
-            *codegenScope,
-            "DefaultResolver" to defaultResolver(),
-        )
-    }
 }

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt

@@ -13,8 +13,14 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegen
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.CustomRuntimeFunction
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.generators.endpointTestsModule
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.rulesgen.SmithyEndpointsStdLib
+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.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.map
+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
 
 /**
@@ -104,6 +110,24 @@ class EndpointsDecorator : ClientCodegenDecorator {
             EndpointConfigCustomization(codegenContext, EndpointTypesGenerator.fromContext(codegenContext))
     }
 
+    override fun serviceRuntimePluginCustomizations(
+        codegenContext: ClientCodegenContext,
+        baseCustomizations: List<ServiceRuntimePluginCustomization>,
+    ): List<ServiceRuntimePluginCustomization> {
+        return baseCustomizations + object : ServiceRuntimePluginCustomization() {
+            override fun section(section: ServiceRuntimePluginSection): Writable {
+                return when (section) {
+                    is ServiceRuntimePluginSection.RegisterRuntimeComponents -> writable {
+                        codegenContext.defaultEndpointResolver()
+                            ?.let { resolver -> section.registerEndpointResolver(this, resolver) }
+                    }
+
+                    else -> emptySection
+                }
+            }
+        }
+    }
+
     override fun extras(codegenContext: ClientCodegenContext, rustCrate: RustCrate) {
         val generator = EndpointTypesGenerator.fromContext(codegenContext)
         rustCrate.withModule(ClientRustModule.Config.endpoint) {
@@ -113,3 +137,29 @@ class EndpointsDecorator : ClientCodegenDecorator {
         }
     }
 }
+
+private fun ClientCodegenContext.defaultEndpointResolver(): Writable? {
+    val generator = EndpointTypesGenerator.fromContext(this)
+    val defaultResolver = generator.defaultResolver() ?: return null
+    val ctx = arrayOf("DefaultResolver" to defaultResolver)
+    return wrapResolver { rustTemplate("#{DefaultResolver}::new()", *ctx) }
+}
+
+fun ClientCodegenContext.wrapResolver(resolver: Writable): Writable {
+    val generator = EndpointTypesGenerator.fromContext(this)
+    return resolver.map { base ->
+        val types = Types(runtimeConfig)
+        val ctx = arrayOf(
+            "DefaultEndpointResolver" to RuntimeType.smithyRuntime(runtimeConfig)
+                .resolve("client::orchestrator::endpoints::DefaultEndpointResolver"),
+            "Params" to generator.paramsStruct(),
+            "OldSharedEndpointResolver" to types.sharedEndpointResolver,
+        )
+
+        rustTemplate(
+            "#{DefaultEndpointResolver}::<#{Params}>::new(#{OldSharedEndpointResolver}::new(#{base}))",
+            *ctx,
+            "base" to base,
+        )
+    }
+}

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

@@ -56,6 +56,7 @@ class ConfigOverrideRuntimePluginGenerator(
                 ) -> Self {
                     let mut layer = config_override.config;
                     let mut components = config_override.runtime_components;
+                    ##[allow(unused_mut)]
                     let mut resolver = #{Resolver}::overrid(initial_config, initial_components, &mut layer, &mut components);
 
                     #{config}

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