update generic clients to support user-configurable runtime plugins (#2864)

_This PR also updates `pre-commit ktlint` runner. Now it won't spit out
a bazillion debug logs when run_

----

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

.pre-commit-config.yaml

@@ -20,7 +20,7 @@ repos:
     files: ^.*$
     pass_filenames: false
 - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-  rev: v2.6.0
+  rev: v2.10.0
   hooks:
   - id: pretty-format-kotlin
     args: [--autofix, --ktlint-version, 0.48.2]

aws/sdk/build.gradle.kts

@@ -105,7 +105,8 @@ fun generateSmithyBuild(services: AwsServices): String {
                             "renameErrors": false,
                             "debugMode": $debugMode,
                             "eventStreamAllowList": [$eventStreamAllowListMembers],
-                            "enableNewSmithyRuntime": "${getSmithyRuntimeMode()}"
+                            "enableNewSmithyRuntime": "${getSmithyRuntimeMode()}",
+                            "enableUserConfigurableRuntimePlugins": false
                         },
                         "service": "${service.service}",
                         "module": "$moduleName",

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

@@ -36,4 +36,5 @@ data class ClientCodegenContext(
     model, symbolProvider, moduleDocProvider, serviceShape, protocol, settings, CodegenTarget.CLIENT,
 ) {
     val smithyRuntimeMode: SmithyRuntimeMode get() = settings.codegenConfig.enableNewSmithyRuntime
+    val enableUserConfigurableRuntimePlugins: Boolean get() = settings.codegenConfig.enableUserConfigurableRuntimePlugins
 }

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

@@ -74,25 +74,25 @@ data class ClientRustSettings(
 
 // TODO(enableNewSmithyRuntimeCleanup): Remove this mode after switching to the orchestrator
 enum class SmithyRuntimeMode {
-    Middleware,
-    BothDefaultMiddleware,
-    BothDefaultOrchestrator,
-    Orchestrator,
+    Middleware, BothDefaultMiddleware, BothDefaultOrchestrator, Orchestrator,
     ;
 
     val exclusivelyGenerateMiddleware: Boolean get() = generateMiddleware && !generateOrchestrator
 
-    val generateMiddleware: Boolean get() = when (this) {
+    val generateMiddleware: Boolean
+        get() = when (this) {
             Middleware, BothDefaultMiddleware, BothDefaultOrchestrator -> true
             else -> false
         }
 
-    val generateOrchestrator: Boolean get() = when (this) {
+    val generateOrchestrator: Boolean
+        get() = when (this) {
             Orchestrator, BothDefaultMiddleware, BothDefaultOrchestrator -> true
             else -> false
         }
 
-    val defaultToMiddleware: Boolean get() = when (this) {
+    val defaultToMiddleware: Boolean
+        get() = when (this) {
             Middleware, BothDefaultMiddleware -> true
             else -> false
         }
@@ -127,6 +127,7 @@ data class ClientCodegenConfig(
     val enableNewSmithyRuntime: SmithyRuntimeMode = defaultEnableNewSmithyRuntime,
     /** If true, adds `endpoint_url`/`set_endpoint_url` methods to the service config */
     val includeEndpointUrlConfig: Boolean = defaultIncludeEndpointUrlConfig,
+    val enableUserConfigurableRuntimePlugins: Boolean = defaultEnableUserConfigurableRuntimePlugins,
 ) : CoreCodegenConfig(
     formatTimeoutSeconds, debugMode,
 ) {
@@ -137,25 +138,24 @@ data class ClientCodegenConfig(
         private val defaultEventStreamAllowList: Set<String> = emptySet()
         private val defaultEnableNewSmithyRuntime = SmithyRuntimeMode.Orchestrator
         private const val defaultIncludeEndpointUrlConfig = true
+        private const val defaultEnableUserConfigurableRuntimePlugins = true
 
         fun fromCodegenConfigAndNode(coreCodegenConfig: CoreCodegenConfig, node: Optional<ObjectNode>) =
             if (node.isPresent) {
                 ClientCodegenConfig(
                     formatTimeoutSeconds = coreCodegenConfig.formatTimeoutSeconds,
                     debugMode = coreCodegenConfig.debugMode,
-                    eventStreamAllowList = node.get().getArrayMember("eventStreamAllowList")
-                        .map { array -> array.toList().mapNotNull { node -> node.asStringNode().orNull()?.value } }
-                        .orNull()?.toSet() ?: defaultEventStreamAllowList,
+                    eventStreamAllowList = node.get().getArrayMember("eventStreamAllowList").map { array ->
+                        array.toList().mapNotNull { node ->
+                            node.asStringNode().orNull()?.value
+                        }
+                    }.orNull()?.toSet() ?: defaultEventStreamAllowList,
                     renameExceptions = node.get().getBooleanMemberOrDefault("renameErrors", defaultRenameExceptions),
-                    includeFluentClient = node.get()
-                        .getBooleanMemberOrDefault("includeFluentClient", defaultIncludeFluentClient),
-                    addMessageToErrors = node.get()
-                        .getBooleanMemberOrDefault("addMessageToErrors", defaultAddMessageToErrors),
-                    enableNewSmithyRuntime = SmithyRuntimeMode.fromString(
-                        node.get().getStringMemberOrDefault("enableNewSmithyRuntime", "middleware"),
-                    ),
-                    includeEndpointUrlConfig = node.get()
-                        .getBooleanMemberOrDefault("includeEndpointUrlConfig", defaultIncludeEndpointUrlConfig),
+                    includeFluentClient = node.get().getBooleanMemberOrDefault("includeFluentClient", defaultIncludeFluentClient),
+                    addMessageToErrors = node.get().getBooleanMemberOrDefault("addMessageToErrors", defaultAddMessageToErrors),
+                    enableNewSmithyRuntime = SmithyRuntimeMode.fromString(node.get().getStringMemberOrDefault("enableNewSmithyRuntime", "middleware")),
+                    includeEndpointUrlConfig = node.get().getBooleanMemberOrDefault("includeEndpointUrlConfig", defaultIncludeEndpointUrlConfig),
+                    enableUserConfigurableRuntimePlugins = node.get().getBooleanMemberOrDefault("userConfigurableRuntimePlugins", defaultEnableUserConfigurableRuntimePlugins),
                 )
             } else {
                 ClientCodegenConfig(

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

@@ -34,6 +34,17 @@ class ClientRuntimeTypesReExportGenerator(
                 "Interceptor" to RuntimeType.interceptor(rc),
                 "SharedInterceptor" to RuntimeType.sharedInterceptor(rc),
             )
+
+            if (codegenContext.enableUserConfigurableRuntimePlugins) {
+                rustTemplate(
+                    """
+                    pub use #{runtime_plugin}::{RuntimePlugin, SharedRuntimePlugin};
+                    pub use #{config_bag}::FrozenLayer;
+                    """,
+                    "runtime_plugin" to RuntimeType.smithyRuntimeApi(rc).resolve("client::runtime_plugin"),
+                    "config_bag" to RuntimeType.smithyTypes(rc).resolve("config_bag"),
+                )
+            }
         }
         rustCrate.withModule(ClientRustModule.endpoint(codegenContext)) {
             rustTemplate(

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

@@ -19,7 +19,6 @@ 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.docs
 import software.amazon.smithy.rust.codegen.core.rustlang.docsOrFallback
-import software.amazon.smithy.rust.codegen.core.rustlang.raw
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate
@@ -305,26 +304,29 @@ class ServiceConfigGenerator(
         }
     }
 
-    private val smithyTypes = RuntimeType.smithyTypes(codegenContext.runtimeConfig)
+    private val moduleUseName = codegenContext.moduleUseName()
+    private val runtimeMode = codegenContext.smithyRuntimeMode
+    private val runtimeConfig = codegenContext.runtimeConfig
+    private val enableUserConfigurableRuntimePlugins = codegenContext.enableUserConfigurableRuntimePlugins
+    private val smithyTypes = RuntimeType.smithyTypes(runtimeConfig)
     val codegenScope = arrayOf(
         *preludeScope,
-        "BoxError" to RuntimeType.boxError(codegenContext.runtimeConfig),
+        "BoxError" to RuntimeType.boxError(runtimeConfig),
         "CloneableLayer" to smithyTypes.resolve("config_bag::CloneableLayer"),
-        "ConfigBag" to RuntimeType.configBag(codegenContext.runtimeConfig),
-        "ConfigBagAccessors" to RuntimeType.configBagAccessors(codegenContext.runtimeConfig),
+        "ConfigBag" to RuntimeType.configBag(runtimeConfig),
+        "ConfigBagAccessors" to RuntimeType.configBagAccessors(runtimeConfig),
         "Cow" to RuntimeType.Cow,
         "FrozenLayer" to smithyTypes.resolve("config_bag::FrozenLayer"),
         "Layer" to smithyTypes.resolve("config_bag::Layer"),
-        "Resolver" to RuntimeType.smithyRuntime(codegenContext.runtimeConfig).resolve("client::config_override::Resolver"),
-        "RuntimeComponentsBuilder" to RuntimeType.runtimeComponentsBuilder(codegenContext.runtimeConfig),
-        "RuntimePlugin" to RuntimeType.runtimePlugin(codegenContext.runtimeConfig),
-        "SharedRuntimePlugin" to RuntimeType.sharedRuntimePlugin(codegenContext.runtimeConfig),
+        "Resolver" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::config_override::Resolver"),
+        "RuntimeComponentsBuilder" to RuntimeType.runtimeComponentsBuilder(runtimeConfig),
+        "RuntimePlugin" to RuntimeType.runtimePlugin(runtimeConfig),
+        "SharedRuntimePlugin" to RuntimeType.sharedRuntimePlugin(runtimeConfig),
+        "runtime_plugin" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::runtime_plugin"),
     )
-    private val moduleUseName = codegenContext.moduleUseName()
-    private val runtimeMode = codegenContext.smithyRuntimeMode
 
     fun render(writer: RustWriter) {
-        writer.docs("Service config.\n")
+        writer.docs("Configuration for a $moduleUseName service client.\n")
         customizations.forEach {
             it.section(ServiceConfig.ConfigStructAdditionalDocs)(writer)
         }
@@ -394,9 +396,9 @@ class ServiceConfigGenerator(
 
         writer.docs("Builder for creating a `Config`.")
         if (runtimeMode.defaultToMiddleware) {
-            writer.raw("#[derive(Clone, Default)]")
+            Attribute(Attribute.derive(RuntimeType.Clone, RuntimeType.Default)).render(writer)
         } else {
-            writer.raw("#[derive(Clone, Debug)]")
+            Attribute(Attribute.derive(RuntimeType.Clone, RuntimeType.Debug)).render(writer)
         }
         writer.rustBlock("pub struct Builder") {
             if (runtimeMode.defaultToOrchestrator) {
@@ -451,19 +453,25 @@ class ServiceConfigGenerator(
                 it.section(ServiceConfig.BuilderImpl)(this)
             }
 
-            if (runtimeMode.defaultToOrchestrator) {
+            if (runtimeMode.generateOrchestrator) {
+                val visibility = if (enableUserConfigurableRuntimePlugins) { "pub" } else { "pub(crate)" }
+
+                docs("Adds a runtime plugin to the config.")
+                if (!enableUserConfigurableRuntimePlugins) { Attribute.AllowUnused.render(this) }
                 rustTemplate(
                     """
-                    /// Adds a runtime plugin to the config.
-                    ##[allow(unused)]
-                    pub(crate) fn runtime_plugin(mut self, plugin: impl #{RuntimePlugin} + 'static) -> Self {
+                    $visibility fn runtime_plugin(mut self, plugin: impl #{RuntimePlugin} + 'static) -> Self {
                         self.push_runtime_plugin(#{SharedRuntimePlugin}::new(plugin));
                         self
                     }
-
-                    /// Adds a runtime plugin to the config.
-                    ##[allow(unused)]
-                    pub(crate) fn push_runtime_plugin(&mut self, plugin: #{SharedRuntimePlugin}) -> &mut Self {
+                    """,
+                    *codegenScope,
+                )
+                docs("Adds a runtime plugin to the config.")
+                if (!enableUserConfigurableRuntimePlugins) { Attribute.AllowUnused.render(this) }
+                rustTemplate(
+                    """
+                    $visibility fn push_runtime_plugin(&mut self, plugin: #{SharedRuntimePlugin}) -> &mut Self {
                         self.runtime_plugins.push(plugin);
                         self
                     }
@@ -528,6 +536,7 @@ class ServiceConfigGenerator(
                     }
                 }
             }
+
             customizations.forEach {
                 it.section(ServiceConfig.Extras)(writer)
             }

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/testutil/TestHelpers.kt

@@ -94,6 +94,9 @@ fun testClientCodegenContext(
 fun ClientCodegenContext.withSmithyRuntimeMode(smithyRuntimeMode: SmithyRuntimeMode): ClientCodegenContext =
     copy(settings = settings.copy(codegenConfig = settings.codegenConfig.copy(enableNewSmithyRuntime = smithyRuntimeMode)))
 
+fun ClientCodegenContext.withEnableUserConfigurableRuntimePlugins(enableUserConfigurableRuntimePlugins: Boolean): ClientCodegenContext =
+    copy(settings = settings.copy(codegenConfig = settings.codegenConfig.copy(enableUserConfigurableRuntimePlugins = enableUserConfigurableRuntimePlugins)))
+
 fun TestWriterDelegator.clientRustSettings() =
     testClientRustSettings(
         service = ShapeId.from("fake#Fake"),

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/config/ServiceConfigGeneratorTest.kt

@@ -14,6 +14,7 @@ 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.SmithyRuntimeMode
 import software.amazon.smithy.rust.codegen.client.testutil.testClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.testutil.withEnableUserConfigurableRuntimePlugins
 import software.amazon.smithy.rust.codegen.client.testutil.withSmithyRuntimeMode
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
@@ -161,7 +162,9 @@ internal class ServiceConfigGeneratorTest {
 
         val model = "namespace empty".asSmithyModel()
         val smithyRuntimeMode = SmithyRuntimeMode.fromString(smithyRuntimeModeStr)
-        val codegenContext = testClientCodegenContext(model).withSmithyRuntimeMode(smithyRuntimeMode)
+        val codegenContext = testClientCodegenContext(model)
+            .withSmithyRuntimeMode(smithyRuntimeMode)
+            .withEnableUserConfigurableRuntimePlugins(true)
         val sut = ServiceConfigGenerator(codegenContext, listOf(ServiceCustomizer(codegenContext)))
         val symbolProvider = codegenContext.symbolProvider
         val project = TestWorkspace.testProject(symbolProvider)
@@ -176,6 +179,28 @@ internal class ServiceConfigGeneratorTest {
                     assert_eq!(config.config_field(), 99);
                     """,
                 )
+
+                unitTest(
+                    "set_runtime_plugin",
+                    """
+                    use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
+                    use aws_smithy_types::config_bag::FrozenLayer;
+
+                    #[derive(Debug)]
+                    struct TestRuntimePlugin;
+
+                    impl RuntimePlugin for TestRuntimePlugin {
+                        fn config(&self) -> Option<FrozenLayer> {
+                            todo!("ExampleRuntimePlugin.config")
+                        }
+                    }
+
+                    let config = Config::builder()
+                        .runtime_plugin(TestRuntimePlugin)
+                        .build();
+                    assert_eq!(config.runtime_plugins.len(), 1);
+                    """,
+                )
             } else {
                 unitTest(
                     "set_config_fields",

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustType.kt

@@ -509,6 +509,7 @@ class Attribute(val inner: Writable, val isDeriveHelper: Boolean = false) {
         val AllowNonSnakeCase = Attribute(allow("non_snake_case"))
         val AllowUnreachableCode = Attribute(allow("unreachable_code"))
         val AllowUnreachablePatterns = Attribute(allow("unreachable_patterns"))
+        val AllowUnused = Attribute(allow("unused"))
         val AllowUnusedImports = Attribute(allow("unused_imports"))
         val AllowUnusedMut = Attribute(allow("unused_mut"))
         val AllowUnusedVariables = Attribute(allow("unused_variables"))