Make TokenBucket and ClientRateLimiter configurable

[ysaito1001]

Motivation and Context

The SDK provides basic retry configuration through the RetryConfig struct, allowing users to set parameters like retry mode, maximum attempts, and backoff durations. However, there has been a need for more advanced configuration options, particularly for token buckets and adaptive rate limiting.

Description

This PR makes TokenBucket and ClientRateLimiter configurable through the RetryPartition struct.

In most cases, a "default" RetryPartition is used where TokenBucket and ClientRateLimiter are constructed behind the scenes and registered in static maps. If a customer wishes to configure a TokenBucket and/or ClientRateLimiter, it can be done by a custom RetryPartition:

Config::builder().retry_partition(
    RetryPartition::custom("test")
        .token_bucket(...)
        .build()
)

For more details, refer to the API docs on the setter and the RetryPartition struct itself.

End-to-end customer experience

Customers can use a custom partition with an SDK like so

    let shared_config = aws_config::defaults(BehaviorVersion::latest()).load().await;
    let unlimited_token_bucket = TokenBucket::unlimited();
    let custom_retry_partition = RetryPartition::custom("my-partition")
        .token_bucket(unlimited_token_bucket)
        .build();

    let conf1 = aws_sdk_s3::config::Builder::from(&shared_config)
        .retry_partition(custom_retry_partition.clone())
        .build();
    let client1 = aws_sdk_s3::Client::from_conf(conf1);
    // call operation on `client1`, which uses `unlimited_token_bucket` in the custom retry partition


    let conf2 = aws_sdk_s3::config::Builder::from(&shared_config)
        // use the same `custom_retry_partition` as `conf1`
        .retry_partition(custom_retry_partition)
        // additional configs different from `conf1` go here...
        .build();
    let client2 = aws_sdk_s3::Client::from_conf(conf2);
    // call operation on `client2`, which uses `unlimited_token_bucket` being shared as `client1`


    let conf3 = aws_sdk_s3::config::Builder::from(&shared_config)
        // custom partition also named `my-partition` but is a different instance than `custom_retry_partition` above
        .retry_partition(
            RetryPartition::custom("my-partition")
                .token_bucket(TokenBucket::default())
                .build(),
        )
        .build();
    let client3 = aws_sdk_s3::Client::from_conf(conf3);
    // call operation on `client3`, which uses the default `TokenBucket` in the custom partition.


    let client4 = aws_sdk_s3::Client::new(&shared_config);
    // call operation on `client4`, which uses the default `TokenBucket` associated with a default retry partition named "s3-<region>"

Dismissed alternatives

We considered alternative approaches for making these configurable but dismissed them for the following reasons:

Alt1: Make them configurable through RetryConfig

RetryConfig::standard().token_bucket(...).client_rate_limiter(...)

RetryConfig is defined in aws-smithy-types whereas TokenBucket and ClientRateLimiter in aws-smithy-runtime. Since aws-smithy-types is higher in the dependency hierarchy, it cannot depend on aws-smithy-runtime to refer to TokenBucket and ClientRateLimiter. Moving TokenBucket and ClientRateLimiter to aws-smithy-types was probably not ideal either, as they are more closely tied to the smithy runtime.

Alt2: Make them configurable through individual setters on config builder

Config::builder().retry_config(...).retry_partition(...).token_bucket(...).client_rate_limiter(...)

This significantly increases the cognitive load for customers, as they see multiple configuration options at the same level.

Testing

• Added unit tests for TokenBucket in rust-runtime/aws-smithy-runtime/src/client/retries/token_bucket.rs
• Added codegen tests for configurable TokenBucket and ClientRateLimiter in aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/RetryPartitionTest.kt.

Checklist

• 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.
• 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.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.

[aajtodd]

This still leaves the possibility that a user creating two different clients that wants to assign them different token buckets or rate limiters results in "first access wins" right?

e.g.

// configure client 1 with custom token bucket
let client1 = Client::builder().retry_partition(RetryPartition::custom("part1").token_bucket(...).build());

// configure second client with different token bucket but same partition name
let client2 = Client::builder().retry_partition(RetryPartition::custom("part1").token_bucket(...).build());

It may not be obvious that client2 will never use the given token bucket.

---

I'm not necessarily saying we shouldn't do this I'm just trying to understand the sharp edges of this approach (all of the approaches discussed have some sort of sharp edge). Also wondering how to best mitigate this if possible.

[ysaito1001]

two different clients that wants to assign them different token buckets or rate limiters results in "first access wins" right?

Yeah, this is just like any other configurations. Instead of sharing a config value, if two clients are provided a separate config value, whichever makes to the config bag wins.

[ysaito1001]

8118cab clarifies the rustdoc per this discussion.

[ysaito1001]

Alternatively, we can consider changing RetryPartitionInner to be

pub(crate) enum RetryPartitionInner {
    Named(Cow<'static, str>),
    Anonymous {
        token_bucket: Option<TokenBucket>,
        client_rate_limiter: Option<ClientRateLimiter>,
    },
}

Anonymous RetryPartitions aren't equal to each other (just like anonymous functions are all different). This way, customers cannot give a name and it may help avoid confusion that they are different even though they have the same custom partition names.

EDIT:
Since RetryPartition implements the Hash trait, the Anonymous variant needs some sort of unique ID, something like

pub(crate) enum RetryPartitionInner {
    Named(Cow<'static, str>),
    Anonymous {
        _id: Arc<()>, // cloned `RetryPartition::Anonymous`s are equal
        token_bucket: Option<TokenBucket>,
        client_rate_limiter: Option<ClientRateLimiter>,
    },
}

[aajtodd]

Hmm need to think on this, that may be worth considering. Let's discuss offline.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.

[github-actions]

A new generated diff is ready to view.

• AWS SDK (ignoring whitespace)
• Client Test (ignoring whitespace)
• No codegen difference in the Server Test
• No codegen difference in the Server Test Python
• No codegen difference in the Server Test Typescript

A new doc preview is ready to view.