Make `TokenBucket` and `ClientRateLimiter` configurable (436fd290) · Commits · Public Repositories / Smithy Rust

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

+49 −6

Original line number	Diff line number	Diff line
		@@ -15,6 +15,7 @@ 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(codegenContext: ClientCodegenContext) : ConfigCustomization() {
		private val runtimeConfig = codegenContext.runtimeConfig
		@@ -23,6 +24,8 @@ class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : Conf
		private val timeoutModule = RuntimeType.smithyTypes(runtimeConfig).resolve("timeout")
		private val retries = RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries")
		private val moduleUseName = codegenContext.moduleUseName()
		private val sdkId = codegenContext.serviceShape.sdkId()
		private val defaultRetryPartition = sdkId.lowercase().replace(" ", "")
		private val codegenScope =
		arrayOf(
		*preludeScope,
		@@ -266,8 +269,50 @@ class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : Conf
		rustTemplate(
		"""
		/// Set the partition for retry-related state. When clients share a retry partition, they will
		/// also share things like token buckets and client rate limiters. By default, all clients
		/// for the same service will share a partition.
		/// also share components such as token buckets and client rate limiters.
		/// See the [`RetryPartition`](#{RetryPartition}) documentation for more details.
		///
		/// ## Default Behavior
		///
		/// When no retry partition is explicitly set, the SDK automatically creates a default retry partition named `$defaultRetryPartition`
		/// (or `$defaultRetryPartition-<region>` if a region is configured).
		/// All $sdkId clients without an explicit retry partition will share this default partition.
		///
		/// ## Notes
		///
		/// - This is an advanced setting — most users won't need to modify it.
		/// - A configured client rate limiter has no effect unless [`RetryConfig::adaptive`](#{RetryConfig}::adaptive) is used.
		///
		/// ## Examples
		///
		/// Creating a custom retry partition with a token bucket:
		/// ```no_run
		/// use $moduleUseName::config::Config;
		/// use $moduleUseName::config::retry::{RetryPartition, TokenBucket};
		///
		/// let token_bucket = TokenBucket::new(10);
		/// let config = Config::builder()
		/// .retry_partition(RetryPartition::custom("custom")
		/// .token_bucket(token_bucket)
		/// .build()
		/// )
		/// .build();
		/// ```
		///
		/// Configuring a client rate limiter with adaptive retry mode:
		/// ```no_run
		/// use $moduleUseName::config::Config;
		/// use $moduleUseName::config::retry::{ClientRateLimiter, RetryConfig, RetryPartition};
		///
		/// let client_rate_limiter = ClientRateLimiter::new(10.0);
		/// let config = Config::builder()
		/// .retry_partition(RetryPartition::custom("custom")
		/// .client_rate_limiter(client_rate_limiter)
		/// .build()
		/// )
		/// .retry_config(RetryConfig::adaptive())
		/// .build();
		/// ```
		pub fn retry_partition(mut self, retry_partition: #{RetryPartition}) -> Self {
		self.set_retry_partition(Some(retry_partition));
		self
		@@ -278,9 +323,7 @@ class ResiliencyConfigCustomization(codegenContext: ClientCodegenContext) : Conf

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

		rustTemplate(
		"pub use #{types_retry}::RetryPartition;",
		"pub use #{types_retry}::{ClientRateLimiter, RetryPartition, TokenBucket};",
		"types_retry" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::retries"),
		)
		}

rust-runtime/Cargo.lock

+1 −1

Original line number	Diff line number	Diff line
		@@ -614,7 +614,7 @@ dependencies = [

		[[package]]
		name = "aws-smithy-runtime"
		version = "1.8.6"
		version = "1.8.7"
		dependencies = [
		"approx",
		"aws-smithy-async",

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

+1 −1

Original line number	Diff line number	Diff line
		[package]
		name = "aws-smithy-runtime"
		version = "1.8.6"
		version = "1.8.7"
		authors = ["AWS Rust SDK Team <aws-sdk-rust@amazon.com>", "Zelda Hessler <zhessler@amazon.com>"]
		description = "The new smithy runtime crate"
		edition = "2021"

rust-runtime/aws-smithy-runtime/src/client/retries.rs

+156 −4

Original line number	Diff line number	Diff line
		@@ -22,25 +22,177 @@ pub use client_rate_limiter::ClientRateLimiterPartition;
		use std::borrow::Cow;

		/// Represents the retry partition, e.g. an endpoint, a region
		///
		/// By default, a retry partition created via [`RetryPartition::new`] uses built-in token bucket and rate limiter settings,
		/// with no option for customization.
		///
		/// To configure these components, use a custom retry partition created via [`RetryPartition::custom`].
		/// Custom partitions allow full control over token bucket and rate limiter.
		///
		/// Two `RetryPartition`s that are equal, as defined by the `Eq` trait, share the same token bucket and client rate limiter.
		/// This means that the token bucket and rate limiter in a custom retry partition are independent
		/// of those in a default retry partition, even if the custom partition has the same name as the default partition.
		#[non_exhaustive]
		#[derive(Clone, Debug, Hash, PartialEq, Eq)]
		#[derive(Clone, Debug)]
		pub struct RetryPartition {
		pub(crate) inner: RetryPartitionInner,
		}

		#[derive(Clone, Debug)]
		pub(crate) enum RetryPartitionInner {
		Default(Cow<'static, str>),
		Configured {
		name: Cow<'static, str>,
		token_bucket: Option<TokenBucket>,
		client_rate_limiter: Option<ClientRateLimiter>,
		},
		}

		impl RetryPartition {
		/// Creates a new `RetryPartition` from the given `name`.
		pub fn new(name: impl Into<Cow<'static, str>>) -> Self {
		Self { name: name.into() }
		Self {
		inner: RetryPartitionInner::Default(name.into()),
		}
		}

		/// Creates a builder for a custom `RetryPartition`.
		pub fn custom(name: impl Into<Cow<'static, str>>) -> RetryPartitionBuilder {
		RetryPartitionBuilder {
		name: name.into(),
		token_bucket: None,
		client_rate_limiter: None,
		}
		}

		fn name(&self) -> &str {
		match &self.inner {
		RetryPartitionInner::Default(name) => name,
		RetryPartitionInner::Configured { name, .. } => name,
		}
		}
		}

		impl PartialEq for RetryPartition {
		fn eq(&self, other: &Self) -> bool {
		match (&self.inner, &other.inner) {
		(RetryPartitionInner::Default(name1), RetryPartitionInner::Default(name2)) => {
		name1 == name2
		}
		(
		RetryPartitionInner::Configured { name: name1, .. },
		RetryPartitionInner::Configured { name: name2, .. },
		) => name1 == name2,
		// Different variants: not equal
		_ => false,
		}
		}
		}

		impl Eq for RetryPartition {}

		impl std::hash::Hash for RetryPartition {
		fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
		match &self.inner {
		RetryPartitionInner::Default(name) => {
		// Hash discriminant for Default variant
		0u8.hash(state);
		name.hash(state);
		}
		RetryPartitionInner::Configured { name, .. } => {
		// Hash discriminant for Configured variant
		1u8.hash(state);
		name.hash(state);
		}
		}
		}
		}

		impl fmt::Display for RetryPartition {
		fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.write_str(&self.name)
		f.write_str(self.name())
		}
		}

		impl Storable for RetryPartition {
		type Storer = StoreReplace<RetryPartition>;
		}

		/// Builder for creating custom retry partitions.
		pub struct RetryPartitionBuilder {
		name: Cow<'static, str>,
		token_bucket: Option<TokenBucket>,
		client_rate_limiter: Option<ClientRateLimiter>,
		}

		impl RetryPartitionBuilder {
		/// Sets the token bucket.
		pub fn token_bucket(mut self, token_bucket: TokenBucket) -> Self {
		self.token_bucket = Some(token_bucket);
		self
		}

		/// Sets the client rate limiter.
		pub fn client_rate_limiter(mut self, client_rate_limiter: ClientRateLimiter) -> Self {
		self.client_rate_limiter = Some(client_rate_limiter);
		self
		}

		/// Builds the custom retry partition.
		pub fn build(self) -> RetryPartition {
		RetryPartition {
		inner: RetryPartitionInner::Configured {
		name: self.name,
		token_bucket: self.token_bucket,
		client_rate_limiter: self.client_rate_limiter,
		},
		}
		}
		}
		#[cfg(test)]
		mod tests {
		use super::*;
		use std::collections::hash_map::DefaultHasher;
		use std::hash::{Hash, Hasher};

		fn hash_value<T: Hash>(t: &T) -> u64 {
		let mut hasher = DefaultHasher::new();
		t.hash(&mut hasher);
		hasher.finish()
		}

		#[test]
		fn test_retry_partition_equality() {
		let default1 = RetryPartition::new("test");
		let default2 = RetryPartition::new("test");
		let default3 = RetryPartition::new("other");

		let configured1 = RetryPartition::custom("test").build();
		let configured2 = RetryPartition::custom("test").build();
		let configured3 = RetryPartition::custom("other").build();

		// Same variant, same name
		assert_eq!(default1, default2);
		assert_eq!(configured1, configured2);

		// Same variant, different name
		assert_ne!(default1, default3);
		assert_ne!(configured1, configured3);

		// Different variant, same name
		assert_ne!(default1, configured1);
		}

		#[test]
		fn test_retry_partition_hash() {
		let default = RetryPartition::new("test");
		let configured = RetryPartition::custom("test").build();

		// Different variants with same name should have different hashes
		assert_ne!(hash_value(&default), hash_value(&configured));

		// Same variants with same name should have same hashes
		let default2 = RetryPartition::new("test");
		assert_eq!(hash_value(&default), hash_value(&default2));
		}
		}

rust-runtime/aws-smithy-runtime/src/client/retries/client_rate_limiter.rs

+6 −0

Original line number	Diff line number	Diff line
		@@ -77,6 +77,12 @@ pub(crate) enum RequestReason {
		InitialRequest,
		}

		impl Default for ClientRateLimiter {
		fn default() -> Self {
		Self::builder().build()
		}
		}

		impl ClientRateLimiter {
		/// Creates a new `ClientRateLimiter`
		pub fn new(seconds_since_unix_epoch: f64) -> Self {