Skip to content

Commit fe0a3a3

Browse files
ysaito1001ysaito1001
authored
Make TokenBucket and ClientRateLimiter configurable (#4263)
2 parents d9c5aeb + 1e9af90 commit fe0a3a3

File tree

10 files changed

+573
-97
lines changed

10 files changed

+573
-97
lines changed

‎.changelog/1755104778.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
---
2+
applies_to:
3+
- client
4+
- aws-sdk-rust
5+
authors:
6+
- ysaito1001
7+
references:
8+
- smithy-rs#4263
9+
breaking: false
10+
new_feature: false
11+
bug_fix: false
12+
---
13+
Make [`TokenBucket`](https://docs.rs/aws-smithy-runtime/latest/aws_smithy_runtime/client/retries/struct.TokenBucket.html) and [`ClientRateLimiter`](https://docs.rs/aws-smithy-runtime/latest/aws_smithy_runtime/client/retries/struct.ClientRateLimiter.html) configurable through [`RetryPartition`](https://docs.rs/aws-smithy-runtime/latest/aws_smithy_runtime/client/retries/struct.RetryPartition.html).

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

Lines changed: 85 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,8 @@ class RetryPartitionTest {
5151
.expect("success");
5252
5353
let log_contents = logs_rx.contents();
54-
assert!(log_contents.contains("token bucket for RetryPartition { name: \"dontcare-us-west-2\" } added to config bag"));
54+
let expected = r##"token bucket for RetryPartition { inner: Default("dontcare-us-west-2") } added to config bag"##;
55+
assert!(log_contents.contains(expected));
5556
5657
""",
5758
*codegenScope,
@@ -80,7 +81,8 @@ class RetryPartitionTest {
8081
.expect("success");
8182
8283
let log_contents = logs_rx.contents();
83-
assert!(log_contents.contains("token bucket for RetryPartition { name: \"user-partition\" } added to config bag"));
84+
let expected = r##"token bucket for RetryPartition { inner: Default("user-partition") } added to config bag"##;
85+
assert!(log_contents.contains(expected));
8486
8587
""",
8688
*codegenScope,
@@ -90,4 +92,85 @@ class RetryPartitionTest {
9092
}
9193
}
9294
}
95+
96+
// This test doesn't need to be in "sdk-codegen" but since "default retry partition" test was initially here,
97+
// it is added to this file for consistency.
98+
@Test
99+
fun `custom retry partition`() {
100+
awsSdkIntegrationTest(
101+
SdkCodegenIntegrationTest.model,
102+
) { ctx, crate ->
103+
val codegenScope =
104+
arrayOf(
105+
"BeforeTransmitInterceptorContextRef" to RuntimeType.beforeTransmitInterceptorContextRef(ctx.runtimeConfig),
106+
"BoxError" to RuntimeType.boxError(ctx.runtimeConfig),
107+
"capture_test_logs" to
108+
CargoDependency.smithyRuntimeTestUtil(ctx.runtimeConfig).toType()
109+
.resolve("test_util::capture_test_logs::capture_test_logs"),
110+
"capture_request" to RuntimeType.captureRequest(ctx.runtimeConfig),
111+
"ConfigBag" to RuntimeType.configBag(ctx.runtimeConfig),
112+
"Intercept" to RuntimeType.intercept(ctx.runtimeConfig),
113+
"RetryConfig" to RuntimeType.smithyTypes(ctx.runtimeConfig).resolve("retry::RetryConfig"),
114+
"RetryPartition" to RuntimeType.smithyRuntime(ctx.runtimeConfig).resolve("client::retries::RetryPartition"),
115+
"RuntimeComponents" to RuntimeType.runtimeComponents(ctx.runtimeConfig),
116+
"TokenBucket" to RuntimeType.smithyRuntime(ctx.runtimeConfig).resolve("client::retries::TokenBucket"),
117+
)
118+
crate.integrationTest("custom_retry_partition") {
119+
tokioTest("test_custom_token_bucket") {
120+
val moduleName = ctx.moduleUseName()
121+
rustTemplate(
122+
"""
123+
use std::sync::{Arc, atomic::{AtomicU32, Ordering}};
124+
use $moduleName::{Client, Config};
125+
126+
##[derive(Clone, Debug, Default)]
127+
struct TestInterceptor {
128+
called: Arc<AtomicU32>,
129+
}
130+
impl #{Intercept} for TestInterceptor {
131+
fn name(&self) -> &'static str {
132+
"TestInterceptor"
133+
}
134+
fn read_before_attempt(
135+
&self,
136+
_context: &#{BeforeTransmitInterceptorContextRef}<'_>,
137+
_runtime_components: &#{RuntimeComponents},
138+
cfg: &mut #{ConfigBag},
139+
) -> Result<(), #{BoxError}> {
140+
self.called.fetch_add(1, Ordering::Relaxed);
141+
let token_bucket = cfg.load::<#{TokenBucket}>().unwrap();
142+
let expected = format!("permits: {}", tokio::sync::Semaphore::MAX_PERMITS);
143+
assert!(
144+
format!("{token_bucket:?}").contains(&expected),
145+
"Expected debug output to contain `{expected}`, but got: {token_bucket:?}"
146+
);
147+
Ok(())
148+
}
149+
}
150+
151+
let (http_client, _) = #{capture_request}(None);
152+
let test_interceptor = TestInterceptor::default();
153+
let client_config = Config::builder()
154+
.interceptor(test_interceptor.clone())
155+
.retry_partition(#{RetryPartition}::custom("test")
156+
.token_bucket(#{TokenBucket}::unlimited())
157+
.build()
158+
)
159+
.http_client(http_client)
160+
.build();
161+
162+
let client = Client::from_conf(client_config);
163+
let _ = client.some_operation().send().await;
164+
165+
assert!(
166+
test_interceptor.called.load(Ordering::Relaxed) == 1,
167+
"the interceptor should have been called"
168+
);
169+
""",
170+
*codegenScope,
171+
)
172+
}
173+
}
174+
}
175+
}
93176
}

‎aws/sdk/sdk-external-types.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ allowed_external_types = [
1515

1616
"aws_smithy_runtime::client::identity::cache::IdentityCache",
1717
"aws_smithy_runtime::client::retries::RetryPartition",
18+
"aws_smithy_runtime::client::retries::client_rate_limiter::ClientRateLimiter",
19+
"aws_smithy_runtime::client::retries::token_bucket::TokenBucket",
1820

1921
"aws_runtime::invocation_id::SharedInvocationIdGenerator",
2022
"aws_runtime::invocation_id::InvocationIdGenerator",

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

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.writable
1515
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
1616
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
1717
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
18+
import software.amazon.smithy.rust.codegen.core.util.sdkId
1819

1920
class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : ConfigCustomization() {
2021
private val runtimeConfig = codegenContext.runtimeConfig
@@ -23,6 +24,8 @@ class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : Conf
2324
private val timeoutModule = RuntimeType.smithyTypes(runtimeConfig).resolve("timeout")
2425
private val retries = RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries")
2526
private val moduleUseName = codegenContext.moduleUseName()
27+
private val sdkId = codegenContext.serviceShape.sdkId()
28+
private val defaultRetryPartition = sdkId.lowercase().replace(" ", "")
2629
private val codegenScope =
2730
arrayOf(
2831
*preludeScope,
@@ -266,8 +269,50 @@ class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : Conf
266269
rustTemplate(
267270
"""
268271
/// Set the partition for retry-related state. When clients share a retry partition, they will
269-
/// also share things like token buckets and client rate limiters. By default, all clients
270-
/// for the same service will share a partition.
272+
/// also share components such as token buckets and client rate limiters.
273+
/// See the [`RetryPartition`](#{RetryPartition}) documentation for more details.
274+
///
275+
/// ## Default Behavior
276+
///
277+
/// When no retry partition is explicitly set, the SDK automatically creates a default retry partition named `$defaultRetryPartition`
278+
/// (or `$defaultRetryPartition-<region>` if a region is configured).
279+
/// All $sdkId clients without an explicit retry partition will share this default partition.
280+
///
281+
/// ## Notes
282+
///
283+
/// - This is an advanced setting — most users won't need to modify it.
284+
/// - A configured client rate limiter has no effect unless [`RetryConfig::adaptive`](#{RetryConfig}::adaptive) is used.
285+
///
286+
/// ## Examples
287+
///
288+
/// Creating a custom retry partition with a token bucket:
289+
/// ```no_run
290+
/// use $moduleUseName::config::Config;
291+
/// use $moduleUseName::config::retry::{RetryPartition, TokenBucket};
292+
///
293+
/// let token_bucket = TokenBucket::new(10);
294+
/// let config = Config::builder()
295+
/// .retry_partition(RetryPartition::custom("custom")
296+
/// .token_bucket(token_bucket)
297+
/// .build()
298+
/// )
299+
/// .build();
300+
/// ```
301+
///
302+
/// Configuring a client rate limiter with adaptive retry mode:
303+
/// ```no_run
304+
/// use $moduleUseName::config::Config;
305+
/// use $moduleUseName::config::retry::{ClientRateLimiter, RetryConfig, RetryPartition};
306+
///
307+
/// let client_rate_limiter = ClientRateLimiter::new(10.0);
308+
/// let config = Config::builder()
309+
/// .retry_partition(RetryPartition::custom("custom")
310+
/// .client_rate_limiter(client_rate_limiter)
311+
/// .build()
312+
/// )
313+
/// .retry_config(RetryConfig::adaptive())
314+
/// .build();
315+
/// ```
271316
pub fn retry_partition(mut self, retry_partition: #{RetryPartition}) -> Self {
272317
self.set_retry_partition(Some(retry_partition));
273318
self
@@ -278,9 +323,7 @@ class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : Conf
278323

279324
rustTemplate(
280325
"""
281-
/// Set the partition for retry-related state. When clients share a retry partition, they will
282-
/// also share things like token buckets and client rate limiters. By default, all clients
283-
/// for the same service will share a partition.
326+
/// Like [`Self::retry_partition`], but takes a mutable reference to the builder and an optional `RetryPartition`
284327
pub fn set_retry_partition(&mut self, retry_partition: #{Option}<#{RetryPartition}>) -> &mut Self {
285328
retry_partition.map(|r| self.config.store_put(r));
286329
self
@@ -327,7 +370,7 @@ class ResiliencyReExportCustomization(codegenContext: ClientCodegenContext) {
327370
)
328371

329372
rustTemplate(
330-
"pub use #{types_retry}::RetryPartition;",
373+
"pub use #{types_retry}::{ClientRateLimiter, RetryPartition, TokenBucket};",
331374
"types_retry" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries"),
332375
)
333376
}

‎rust-runtime/Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎rust-runtime/aws-smithy-runtime/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "aws-smithy-runtime"
3-
version = "1.9.0"
3+
version = "1.9.1"
44
authors = ["AWS Rust SDK Team <[email protected]>", "Zelda Hessler <[email protected]>"]
55
description = "The new smithy runtime crate"
66
edition = "2021"

0 commit comments

Comments
 (0)