Simplify default config and add default async sleep (#3071)

While implementing identity caching, I noticed we don't have a default async sleep impl wired up for generic clients, which was causing caching to panic in many tests since it needs a sleep impl for timeouts. I likely need to figure out what to do other than panic if there's no sleep impl, but that's a problem for a different PR. This PR adds a default sleep impl to generic clients, and also simplifies how default config works. Previously, the generated config `Builder::build` method set all the defaults with a series of "if not set, then set default" statements. In this PR, defaults are registered via default ordered runtime plugins. Additionally, I cleaned up the standard retry strategy: - The `TokenBucketPartition` didn't appear to be used at all, so I deleted it. - `StandardRetryStrategy` was taking retry config at construction, which means it isn't possible to config override the retry config unless the strategy is recreated. It now doesn't take any config at all during construction. - The adaptive retry client rate limiter was created at construction based on retry config at that point in time. This means config overrides wouldn't work with it, so it is also no longer set up at construction time. - Removed some unused runtime plugins. ## Checklist  - [ ] I have updated `CHANGELOG.next.toml` if I made changes to the smithy-rs codegen or runtime crates - [ ] I have updated `CHANGELOG.next.toml` if I made changes to the AWS SDK, generated SDK code, or SDK runtime crates ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._
smithy-lang · Oct 17, 2023 · bcfc211 · bcfc211
1 parent d48acae
commit bcfc211
Show file tree

Hide file tree

Showing 18 changed files with 352 additions and 331 deletions.
diff --git a/aws/rust-runtime/aws-config/src/imds/client.rs b/aws/rust-runtime/aws-config/src/imds/client.rs
@@ -238,13 +238,14 @@ impl ImdsCommonRuntimePlugin {
     fn new(
         config: &ProviderConfig,
         endpoint_resolver: ImdsEndpointResolver,
-        retry_config: &RetryConfig,
+        retry_config: RetryConfig,
         timeout_config: TimeoutConfig,
     ) -> Self {
         let mut layer = Layer::new("ImdsCommonRuntimePlugin");
         layer.store_put(AuthSchemeOptionResolverParams::new(()));
         layer.store_put(EndpointResolverParams::new(()));
         layer.store_put(SensitiveOutput);
+        layer.store_put(retry_config);
         layer.store_put(timeout_config);
         layer.store_put(user_agent());
 
@@ -255,7 +256,7 @@ impl ImdsCommonRuntimePlugin {
                 .with_endpoint_resolver(Some(endpoint_resolver))
                 .with_interceptor(UserAgentInterceptor::new())
                 .with_retry_classifier(SharedRetryClassifier::new(ImdsResponseRetryClassifier))
-                .with_retry_strategy(Some(StandardRetryStrategy::new(retry_config)))
+                .with_retry_strategy(Some(StandardRetryStrategy::new()))
                 .with_time_source(Some(config.time_source()))
                 .with_sleep_impl(config.sleep_impl()),
         }
@@ -423,7 +424,7 @@ impl Builder {
         let common_plugin = SharedRuntimePlugin::new(ImdsCommonRuntimePlugin::new(
             &config,
             endpoint_resolver,
-            &retry_config,
+            retry_config,
             timeout_config,
         ));
         let operation = Operation::builder()

diff --git a/.../amazon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomization.kt b/.../amazon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomization.kt
@@ -7,20 +7,16 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations
 
 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.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.ServiceConfig
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
-import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 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.util.sdkId
 
-class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenContext) : ConfigCustomization() {
+class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : ConfigCustomization() {
     private val runtimeConfig = codegenContext.runtimeConfig
     private val retryConfig = RuntimeType.smithyTypes(runtimeConfig).resolve("retry")
     private val sleepModule = RuntimeType.smithyAsync(runtimeConfig).resolve("rt::sleep")
@@ -44,8 +40,6 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
         "StandardRetryStrategy" to retries.resolve("strategy::StandardRetryStrategy"),
         "SystemTime" to RuntimeType.std.resolve("time::SystemTime"),
         "TimeoutConfig" to timeoutModule.resolve("TimeoutConfig"),
-        "TokenBucket" to retries.resolve("TokenBucket"),
-        "TokenBucketPartition" to retries.resolve("TokenBucketPartition"),
     )
 
     override fun section(section: ServiceConfig) =
@@ -281,57 +275,6 @@ class ResiliencyConfigCustomization(private val codegenContext: ClientCodegenCon
                     )
                 }
 
-                is ServiceConfig.BuilderBuild -> {
-                    rustTemplate(
-                        """
-                        if layer.load::<#{RetryConfig}>().is_none() {
-                            layer.store_put(#{RetryConfig}::disabled());
-                        }
-                        let retry_config = layer.load::<#{RetryConfig}>().expect("set to default above").clone();
-
-                        if layer.load::<#{RetryPartition}>().is_none() {
-                            layer.store_put(#{RetryPartition}::new("${codegenContext.serviceShape.sdkId()}"));
-                        }
-                        let retry_partition = layer.load::<#{RetryPartition}>().expect("set to default above").clone();
-
-                        if retry_config.has_retry() {
-                            #{debug}!("using retry strategy with partition '{}'", retry_partition);
-                        }
-
-                        if retry_config.mode() == #{RetryMode}::Adaptive {
-                            if let #{Some}(time_source) = self.runtime_components.time_source() {
-                                let seconds_since_unix_epoch = time_source
-                                    .now()
-                                    .duration_since(#{SystemTime}::UNIX_EPOCH)
-                                    .expect("the present takes place after the UNIX_EPOCH")
-                                    .as_secs_f64();
-                                let client_rate_limiter_partition = #{ClientRateLimiterPartition}::new(retry_partition.clone());
-                                let client_rate_limiter = CLIENT_RATE_LIMITER.get_or_init(client_rate_limiter_partition, || {
-                                    #{ClientRateLimiter}::new(seconds_since_unix_epoch)
-                                });
-                                layer.store_put(client_rate_limiter);
-                            }
-                        }
-
-                        // The token bucket is used for both standard AND adaptive retries.
-                        let token_bucket_partition = #{TokenBucketPartition}::new(retry_partition);
-                        let token_bucket = TOKEN_BUCKET.get_or_init(token_bucket_partition, #{TokenBucket}::default);
-                        layer.store_put(token_bucket);
-
-                        // TODO(enableNewSmithyRuntimeCleanup): Should not need to provide a default once smithy-rs##2770
-                        //  is resolved
-                        if layer.load::<#{TimeoutConfig}>().is_none() {
-                            layer.store_put(#{TimeoutConfig}::disabled());
-                        }
-
-                        self.runtime_components.set_retry_strategy(#{Some}(
-                            #{SharedRetryStrategy}::new(#{StandardRetryStrategy}::new(&retry_config)))
-                        );
-                        """,
-                        *codegenScope,
-                    )
-                }
-
                 else -> emptySection
             }
         }
@@ -366,32 +309,3 @@ class ResiliencyReExportCustomization(codegenContext: ClientCodegenContext) {
         }
     }
 }
-
-class ResiliencyServiceRuntimePluginCustomization(codegenContext: ClientCodegenContext) : ServiceRuntimePluginCustomization() {
-    private val runtimeConfig = codegenContext.runtimeConfig
-    private val smithyRuntime = RuntimeType.smithyRuntime(runtimeConfig)
-    private val retries = smithyRuntime.resolve("client::retries")
-    private val codegenScope = arrayOf(
-        "TokenBucket" to retries.resolve("TokenBucket"),
-        "TokenBucketPartition" to retries.resolve("TokenBucketPartition"),
-        "ClientRateLimiter" to retries.resolve("ClientRateLimiter"),
-        "ClientRateLimiterPartition" to retries.resolve("ClientRateLimiterPartition"),
-        "StaticPartitionMap" to smithyRuntime.resolve("static_partition_map::StaticPartitionMap"),
-    )
-
-    override fun section(section: ServiceRuntimePluginSection): Writable = writable {
-        when (section) {
-            is ServiceRuntimePluginSection.DeclareSingletons -> {
-                rustTemplate(
-                    """
-                    static TOKEN_BUCKET: #{StaticPartitionMap}<#{TokenBucketPartition}, #{TokenBucket}> = #{StaticPartitionMap}::new();
-                    static CLIENT_RATE_LIMITER: #{StaticPartitionMap}<#{ClientRateLimiterPartition}, #{ClientRateLimiter}> = #{StaticPartitionMap}::new();
-                    """,
-                    *codegenScope,
-                )
-            }
-
-            else -> emptySection
-        }
-    }
-}
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
@@ -14,7 +14,6 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.Intercep
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.MetadataCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyReExportCustomization
-import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyServiceRuntimePluginCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryClassifierConfigCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryClassifierOperationCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.customizations.RetryClassifierServiceRuntimePluginCustomization
@@ -113,7 +112,6 @@ class RequiredCustomizations : ClientCodegenDecorator {
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<ServiceRuntimePluginCustomization>,
     ): List<ServiceRuntimePluginCustomization> = baseCustomizations +
-        ResiliencyServiceRuntimePluginCustomization(codegenContext) +
         ConnectionPoisoningRuntimePluginCustomization(codegenContext) +
         RetryClassifierServiceRuntimePluginCustomization(codegenContext)
 }
diff --git a/...tware/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt b/...tware/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt
@@ -40,7 +40,6 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.stripOuter
 import software.amazon.smithy.rust.codegen.core.rustlang.withBlockTemplate
 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.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
@@ -50,9 +49,11 @@ import software.amazon.smithy.rust.codegen.core.smithy.expectRustMetadata
 import software.amazon.smithy.rust.codegen.core.smithy.generators.getterName
 import software.amazon.smithy.rust.codegen.core.smithy.generators.setterName
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
+import software.amazon.smithy.rust.codegen.core.util.dq
 import software.amazon.smithy.rust.codegen.core.util.inputShape
 import software.amazon.smithy.rust.codegen.core.util.orNull
 import software.amazon.smithy.rust.codegen.core.util.outputShape
+import software.amazon.smithy.rust.codegen.core.util.sdkId
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 
 class FluentClientGenerator(
@@ -161,7 +162,7 @@ class FluentClientGenerator(
                 }
                 """,
                 *clientScope,
-                "base_client_runtime_plugins" to baseClientRuntimePluginsFn(runtimeConfig),
+                "base_client_runtime_plugins" to baseClientRuntimePluginsFn(codegenContext),
             )
         }
 
@@ -446,22 +447,36 @@ class FluentClientGenerator(
     }
 }
 
-private fun baseClientRuntimePluginsFn(runtimeConfig: RuntimeConfig): RuntimeType =
+private fun baseClientRuntimePluginsFn(codegenContext: ClientCodegenContext): RuntimeType = codegenContext.runtimeConfig.let { rc ->
     RuntimeType.forInlineFun("base_client_runtime_plugins", ClientRustModule.config) {
+        val api = RuntimeType.smithyRuntimeApi(rc)
+        val rt = RuntimeType.smithyRuntime(rc)
         rustTemplate(
             """
             pub(crate) fn base_client_runtime_plugins(
                 mut config: crate::Config,
             ) -> #{RuntimePlugins} {
                 let mut configured_plugins = #{Vec}::new();
                 ::std::mem::swap(&mut config.runtime_plugins, &mut configured_plugins);
+
+                let defaults = [
+                    #{default_http_client_plugin}(),
+                    #{default_retry_config_plugin}(${codegenContext.serviceShape.sdkId().dq()}),
+                    #{default_sleep_impl_plugin}(),
+                    #{default_time_source_plugin}(),
+                    #{default_timeout_config_plugin}(),
+                ].into_iter().flatten();
+
                 let mut plugins = #{RuntimePlugins}::new()
-                    .with_client_plugin(#{default_http_client_plugin}())
+                    // defaults
+                    .with_client_plugins(defaults)
+                    // user config
                     .with_client_plugin(
                         #{StaticRuntimePlugin}::new()
                             .with_config(config.config.clone())
                             .with_runtime_components(config.runtime_components.clone())
                     )
+                    // codegen config
                     .with_client_plugin(crate::config::ServiceRuntimePlugin::new(config))
                     .with_client_plugin(#{NoAuthRuntimePlugin}::new());
                 for plugin in configured_plugins {
@@ -471,15 +486,17 @@ private fun baseClientRuntimePluginsFn(runtimeConfig: RuntimeConfig): RuntimeTyp
             }
             """,
             *preludeScope,
-            "RuntimePlugins" to RuntimeType.runtimePlugins(runtimeConfig),
-            "NoAuthRuntimePlugin" to RuntimeType.smithyRuntime(runtimeConfig)
-                .resolve("client::auth::no_auth::NoAuthRuntimePlugin"),
-            "StaticRuntimePlugin" to RuntimeType.smithyRuntimeApi(runtimeConfig)
-                .resolve("client::runtime_plugin::StaticRuntimePlugin"),
-            "default_http_client_plugin" to RuntimeType.smithyRuntime(runtimeConfig)
-                .resolve("client::http::default_http_client_plugin"),
+            "default_http_client_plugin" to rt.resolve("client::defaults::default_http_client_plugin"),
+            "default_retry_config_plugin" to rt.resolve("client::defaults::default_retry_config_plugin"),
+            "default_sleep_impl_plugin" to rt.resolve("client::defaults::default_sleep_impl_plugin"),
+            "default_timeout_config_plugin" to rt.resolve("client::defaults::default_timeout_config_plugin"),
+            "default_time_source_plugin" to rt.resolve("client::defaults::default_time_source_plugin"),
+            "NoAuthRuntimePlugin" to rt.resolve("client::auth::no_auth::NoAuthRuntimePlugin"),
+            "RuntimePlugins" to RuntimeType.runtimePlugins(rc),
+            "StaticRuntimePlugin" to api.resolve("client::runtime_plugin::StaticRuntimePlugin"),
         )
     }
+}
 
 /**
  * For a given `operation` shape, return a list of strings where each string describes the name and input type of one of

diff --git a/...zon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomizationTest.kt b/...zon/smithy/rust/codegen/client/smithy/customizations/ResiliencyConfigCustomizationTest.kt
@@ -45,7 +45,7 @@ internal class ResiliencyConfigCustomizationTest {
         project.withModule(ClientRustModule.config) {
             ServiceRuntimePluginGenerator(codegenContext).render(
                 this,
-                listOf(ResiliencyServiceRuntimePluginCustomization(codegenContext)),
+                emptyList(),
             )
         }
         ResiliencyReExportCustomization(codegenContext).extras(project)

diff --git a/...e/amazon/smithy/rust/codegen/client/smithy/customizations/SensitiveOutputDecoratorTest.kt b/...e/amazon/smithy/rust/codegen/client/smithy/customizations/SensitiveOutputDecoratorTest.kt
@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations
 import org.junit.jupiter.api.Test
 import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
+import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
@@ -16,6 +17,8 @@ import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
 
 class SensitiveOutputDecoratorTest {
     private fun codegenScope(runtimeConfig: RuntimeConfig): Array<Pair<String, Any>> = arrayOf(
+        "capture_test_logs" to CargoDependency.smithyRuntimeTestUtil(runtimeConfig).toType()
+            .resolve("test_util::capture_test_logs::capture_test_logs"),
         "capture_request" to RuntimeType.captureRequest(runtimeConfig),
         "SdkBody" to RuntimeType.sdkBody(runtimeConfig),
     )
@@ -48,10 +51,10 @@ class SensitiveOutputDecoratorTest {
             rustCrate.integrationTest("redacting_sensitive_response_body") {
                 val moduleName = codegenContext.moduleUseName()
                 Attribute.TokioTest.render(this)
-                Attribute.TracedTest.render(this)
                 rustTemplate(
                     """
                     async fn redacting_sensitive_response_body() {
+                        let (_logs, logs_rx) = #{capture_test_logs}();
                         let (http_client, _r) = #{capture_request}(Some(
                             http::Response::builder()
                                 .status(200)
@@ -69,7 +72,8 @@ class SensitiveOutputDecoratorTest {
                             .await
                             .expect("success");
 
-                        assert!(logs_contain("** REDACTED **"));
+                        let log_contents = logs_rx.contents();
+                        assert!(log_contents.contains("** REDACTED **"));
                     }
                     """,
                     *codegenScope(codegenContext.runtimeConfig),

diff --git a/.../smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGeneratorTest.kt b/.../smithy/rust/codegen/client/smithy/generators/ConfigOverrideRuntimePluginGeneratorTest.kt
@@ -142,7 +142,7 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
     }
 
     @Test
-    fun `operation overrides retry strategy`() {
+    fun `operation overrides retry config`() {
         clientIntegrationTest(model) { clientCodegenContext, rustCrate ->
             val runtimeConfig = clientCodegenContext.runtimeConfig
             val codegenScope = arrayOf(
@@ -164,11 +164,13 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
                     .resolve("client::retries::RetryClassifiers"),
                 "RuntimeComponentsBuilder" to RuntimeType.runtimeComponentsBuilder(runtimeConfig),
                 "RuntimePlugin" to RuntimeType.runtimePlugin(runtimeConfig),
+                "StandardRetryStrategy" to RuntimeType.smithyRuntime(runtimeConfig)
+                    .resolve("client::retries::strategy::StandardRetryStrategy"),
                 "ShouldAttempt" to RuntimeType.smithyRuntimeApi(runtimeConfig)
                     .resolve("client::retries::ShouldAttempt"),
             )
             rustCrate.testModule {
-                unitTest("test_operation_overrides_retry_strategy") {
+                unitTest("test_operation_overrides_retry_config") {
                     rustTemplate(
                         """
                         use #{RuntimePlugin};
@@ -193,6 +195,8 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
 
                         // Emulate the merging of runtime components from runtime plugins that the orchestrator does
                         let runtime_components = #{RuntimeComponentsBuilder}::for_tests()
+                            // emulate the default retry config plugin by setting a retry strategy
+                            .with_retry_strategy(#{Some}(#{StandardRetryStrategy}::new()))
                             .merge_from(&client_config.runtime_components)
                             .merge_from(&retry_classifiers_component)
                             .build()
@@ -219,6 +223,8 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
 
                         // Emulate the merging of runtime components from runtime plugins that the orchestrator does
                         let runtime_components = #{RuntimeComponentsBuilder}::for_tests()
+                            // emulate the default retry config plugin by setting a retry strategy
+                            .with_retry_strategy(#{Some}(#{StandardRetryStrategy}::new()))
                             .merge_from(&client_config.runtime_components)
                             .merge_from(&retry_classifiers_component)
                             .merge_from(&config_override.runtime_components)