fix token bucket not being set for both standard and adaptive retry modes (#3964)

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here -->
awslabs/aws-sdk-rust#1234

## Description
<!--- Describe your changes in detail -->

PR adds a new interceptor registered as part of the default retry plugin
components that ensures a token bucket is _always_ present and available
to the retry strategy. The buckets are partitioned off the retry
partition (which defaults to the service name and is already set by the
default plugin). We use a `static` variable in the runtime for this
which means that token buckets can and will apply to every single client
that uses the same retry partition. The implementation tries to avoid
contention on this new global lock by only consulting it if the retry
partition is overridden after client creation.

For AWS SDK clients I've updated the default retry partition clients are
created with to include the region when set. Now the default partition
for a client will be `{service}-{region}` (e.g. `sts-us-west-2`) rather
than just the service name (e.g. `sts`). This partitioning is a little
more granular and closer to what we want/expect as failures in one
region should not cause throttling to another (and vice versa for
success in one should not increase available quota in another).

I also updated the implementation to follow the SEP a little more
literally/closely as far as structure which fixes some subtle bugs.
State is updated in one place and we ensure that the token bucket is
always consulted (before the token bucket could be skipped in the case
of adaptive retries returning a delay and the adaptive rate limit was
updated in multiple branches).


## Testing
<!--- Please describe in detail how you tested your changes -->
<!--- Include details of your testing environment, and the tests you ran
to -->
<!--- see how your change affects other areas of the code, etc. -->

## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [x ] For changes to the smithy-rs codegen or runtime crates, I have
created a changelog entry Markdown file in the `.changelog` directory,
specifying "client," "server," or both in the `applies_to` key.
- [ x] For changes to the AWS SDK, generated SDK code, or SDK runtime
crates, I have created a changelog entry Markdown file in the
`.changelog` directory, specifying "aws-sdk-rust" in the `applies_to`
key.

----

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

Author: aajtodd

.changelog/1736370747.md

@@ -0,0 +1,13 @@
+---
+applies_to:
+- aws-sdk-rust
+- client
+authors:
+- aajtodd
+references:
+- aws-sdk-rust#1234
+breaking: false
+new_feature: false
+bug_fix: true
+---
+Fix token bucket not being set for standard and adaptive retry modes

aws/rust-runtime/Cargo.lock

@@ -24,6 +24,7 @@ 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.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
 import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsSection
@@ -58,6 +59,7 @@ class AwsFluentClientDecorator : ClientCodegenDecorator {
                 listOf(
                     AwsPresignedFluentBuilderMethod(codegenContext),
                     AwsFluentClientDocs(codegenContext),
+                    AwsFluentClientRetryPartition(codegenContext),
                 ).letIf(codegenContext.serviceShape.id == ShapeId.from("com.amazonaws.s3#AmazonS3")) {
                     it + S3ExpressFluentClientCustomization(codegenContext)
                 },
@@ -166,3 +168,28 @@ private class AwsFluentClientDocs(private val codegenContext: ClientCodegenConte
         }
     }
 }
+
+/**
+ * Replaces the default retry partition for all operations to include the AWS region if set
+ */
+private class AwsFluentClientRetryPartition(private val codegenContext: ClientCodegenContext) : FluentClientCustomization() {
+    override fun section(section: FluentClientSection): Writable {
+        return when {
+            section is FluentClientSection.BeforeBaseClientPluginSetup && usesRegion(codegenContext) -> {
+                writable {
+                    rustTemplate(
+                        """
+                        let default_retry_partition = match config.region() {
+                            Some(region) => #{Cow}::from(format!("{default_retry_partition}-{}", region)),
+                            None => #{Cow}::from(default_retry_partition),
+                        };
+                        """,
+                        *preludeScope,
+                        "Cow" to RuntimeType.Cow,
+                    )
+                }
+            }
+            else -> emptySection
+        }
+    }
+}

aws/rust-runtime/aws-config/Cargo.lock

@@ -83,13 +83,6 @@ class RegionDecorator : ClientCodegenDecorator {
     private val envKey = "AWS_REGION".dq()
     private val profileKey = "region".dq()
 
-    // Services that have an endpoint ruleset that references the SDK::Region built in, or
-    // that use SigV4, both need a configurable region.
-    private fun usesRegion(codegenContext: ClientCodegenContext) =
-        codegenContext.getBuiltIn(AwsBuiltIns.REGION) != null ||
-            ServiceIndex.of(codegenContext.model)
-                .getEffectiveAuthSchemes(codegenContext.serviceShape).containsKey(SigV4Trait.ID)
-
     override fun configCustomizations(
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<ConfigCustomization>,
@@ -223,3 +216,14 @@ class RegionProviderConfig(codegenContext: ClientCodegenContext) : ConfigCustomi
 }
 
 fun region(runtimeConfig: RuntimeConfig) = AwsRuntimeType.awsTypes(runtimeConfig).resolve("region")
+
+/**
+ * Test if region is used and configured for a model (and available on a service client).
+ *
+ * Services that have an endpoint ruleset that references the SDK::Region built in, or
+ * that use SigV4, both need a configurable region.
+ */
+fun usesRegion(codegenContext: ClientCodegenContext) =
+    codegenContext.getBuiltIn(AwsBuiltIns.REGION) != null ||
+        ServiceIndex.of(codegenContext.model)
+            .getEffectiveAuthSchemes(codegenContext.serviceShape).containsKey(SigV4Trait.ID)

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsFluentClientDecorator.kt

@@ -0,0 +1,93 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.rustsdk
+
+import org.junit.jupiter.api.Test
+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.RuntimeType
+import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
+import software.amazon.smithy.rust.codegen.core.testutil.tokioTest
+
+class RetryPartitionTest {
+    @Test
+    fun `default retry partition`() {
+        awsSdkIntegrationTest(SdkCodegenIntegrationTest.model) { ctx, rustCrate ->
+            val rc = ctx.runtimeConfig
+            val codegenScope =
+                arrayOf(
+                    *RuntimeType.preludeScope,
+                    "capture_request" to RuntimeType.captureRequest(rc),
+                    "capture_test_logs" to
+                        CargoDependency.smithyRuntimeTestUtil(rc).toType()
+                            .resolve("test_util::capture_test_logs::capture_test_logs"),
+                    "Credentials" to
+                        AwsRuntimeType.awsCredentialTypesTestUtil(rc)
+                            .resolve("Credentials"),
+                    "Region" to AwsRuntimeType.awsTypes(rc).resolve("region::Region"),
+                )
+
+            rustCrate.integrationTest("default_retry_partition") {
+                tokioTest("default_retry_partition_includes_region") {
+                    val moduleName = ctx.moduleUseName()
+                    rustTemplate(
+                        """
+                        let (_logs, logs_rx) = #{capture_test_logs}();
+                        let (http_client, _rx) = #{capture_request}(#{None});
+                        let client_config = $moduleName::Config::builder()
+                            .http_client(http_client)
+                            .region(#{Region}::new("us-west-2"))
+                            .credentials_provider(#{Credentials}::for_tests())
+                            .build();
+
+                        let client = $moduleName::Client::from_conf(client_config);
+
+                        let _ = client
+                            .some_operation()
+                            .send()
+                            .await
+                            .expect("success");
+
+                        let log_contents = logs_rx.contents();
+                        assert!(log_contents.contains("token bucket for RetryPartition { name: \"dontcare-us-west-2\" } added to config bag"));
+
+                        """,
+                        *codegenScope,
+                    )
+                }
+
+                tokioTest("user_config_retry_partition") {
+                    val moduleName = ctx.moduleUseName()
+                    rustTemplate(
+                        """
+                        let (_logs, logs_rx) = #{capture_test_logs}();
+                        let (http_client, _rx) = #{capture_request}(#{None});
+                        let client_config = $moduleName::Config::builder()
+                            .http_client(http_client)
+                            .region(#{Region}::new("us-west-2"))
+                            .credentials_provider(#{Credentials}::for_tests())
+                            .retry_partition(#{RetryPartition}::new("user-partition"))
+                            .build();
+
+                        let client = $moduleName::Client::from_conf(client_config);
+
+                        let _ = client
+                            .some_operation()
+                            .send()
+                            .await
+                            .expect("success");
+
+                        let log_contents = logs_rx.contents();
+                        assert!(log_contents.contains("token bucket for RetryPartition { name: \"user-partition\" } added to config bag"));
+
+                        """,
+                        *codegenScope,
+                        "RetryPartition" to RuntimeType.smithyRuntime(ctx.runtimeConfig).resolve("client::retries::RetryPartition"),
+                    )
+                }
+            }
+        }
+    }
+}

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/RegionDecorator.kt

@@ -89,6 +89,10 @@ sealed class FluentClientSection(name: String) : Section(name) {
     /** Write custom code for adding additional client plugins to base_client_runtime_plugins */
     data class AdditionalBaseClientPlugins(val plugins: String, val config: String) :
         FluentClientSection("AdditionalBaseClientPlugins")
+
+    /** Write additional code before plugins are configured */
+    data class BeforeBaseClientPluginSetup(val config: String) :
+        FluentClientSection("BeforeBaseClientPluginSetup")
 }
 
 abstract class FluentClientCustomization : NamedCustomization<FluentClientSection>()

aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/RetryPartitionTest.kt

@@ -266,11 +266,14 @@ private fun baseClientRuntimePluginsFn(
                     ::std::mem::swap(&mut config.runtime_plugins, &mut configured_plugins);
                     #{update_bmv}
 
+                    let default_retry_partition = ${codegenContext.serviceShape.sdkId().dq()};
+                    #{before_plugin_setup}
+
                     let mut plugins = #{RuntimePlugins}::new()
                         // defaults
                         .with_client_plugins(#{default_plugins}(
                             #{DefaultPluginParams}::new()
-                                .with_retry_partition_name(${codegenContext.serviceShape.sdkId().dq()})
+                                .with_retry_partition_name(default_retry_partition)
                                 .with_behavior_version(config.behavior_version.expect(${behaviorVersionError.dq()}))
                         ))
                         // user config
@@ -299,6 +302,13 @@ private fun baseClientRuntimePluginsFn(
                             FluentClientSection.AdditionalBaseClientPlugins("plugins", "config"),
                         )
                     },
+                "before_plugin_setup" to
+                    writable {
+                        writeCustomizations(
+                            customizations,
+                            FluentClientSection.BeforeBaseClientPluginSetup("config"),
+                        )
+                    },
                 "DefaultPluginParams" to rt.resolve("client::defaults::DefaultPluginParams"),
                 "default_plugins" to rt.resolve("client::defaults::default_plugins"),
                 "NoAuthRuntimePlugin" to rt.resolve("client::auth::no_auth::NoAuthRuntimePlugin"),

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

@@ -182,6 +182,7 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
                     "ShouldAttempt" to
                         RuntimeType.smithyRuntimeApi(runtimeConfig)
                             .resolve("client::retries::ShouldAttempt"),
+                    "TokenBucket" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries::TokenBucket"),
                 )
             rustCrate.testModule {
                 unitTest("test_operation_overrides_retry_config") {
@@ -199,6 +200,7 @@ internal class ConfigOverrideRuntimePluginGeneratorTest {
 
                         let mut layer = #{Layer}::new("test");
                         layer.store_put(#{RequestAttempts}::new(1));
+                        layer.store_put(#{TokenBucket}::default());
 
                         let mut cfg = #{ConfigBag}::of_layers(vec![layer]);
                         let client_config_layer = client_config.config;