aws_smithy_runtime/client/
retries.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6/// Smithy retry classifiers.
7pub mod classifiers;
8
9/// Smithy retry strategies.
10pub mod strategy;
11
12mod client_rate_limiter;
13pub(crate) mod token_bucket;
14
15use aws_smithy_types::config_bag::{Storable, StoreReplace};
16use std::fmt;
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19
20// Shared slot for the retry strategy to communicate a backoff delay to the
21// orchestrator when returning `ShouldAttempt::No`.
22//
23// Used for long-polling operations: when the token bucket is empty, the
24// strategy writes the backoff here and returns `No`. The orchestrator
25// sleeps for the delay before returning to the caller, preventing request
26// amplification in polling loops.
27//
28// Uses `Arc<Mutex>` interior mutability so the strategy can write through
29// `&ConfigBag` (immutable). The orchestrator seeds this before the retry
30// loop and reads it after the loop exits on `No`.
31#[derive(Clone, Debug, Default)]
32pub(crate) struct LongPollingBackoff(Arc<Mutex<Option<Duration>>>);
33
34impl LongPollingBackoff {
35    pub(crate) fn set(&self, delay: Duration) {
36        *self.0.lock().expect("lock is acquired") = Some(delay);
37    }
38    pub(crate) fn take(&self) -> Option<Duration> {
39        self.0.lock().expect("lock is acquired").take()
40    }
41}
42
43impl Storable for LongPollingBackoff {
44    type Storer = StoreReplace<Self>;
45}
46
47pub use client_rate_limiter::{
48    ClientRateLimiter, ClientRateLimiterBuilder, ClientRateLimiterPartition,
49};
50pub use token_bucket::{TokenBucket, TokenBucketBuilder, MAXIMUM_CAPACITY};
51
52use std::borrow::Cow;
53
54/// Represents the retry partition, e.g. an endpoint, a region
55///
56/// A retry partition created with [`RetryPartition::new`] uses built-in
57/// token bucket and rate limiter settings, with no option for customization.
58/// Default partitions with the same name share the same token bucket
59/// and client rate limiter.
60///
61/// To customize these components, use a custom retry partition via [`RetryPartition::custom`].
62/// A custom partition owns its token bucket and rate limiter, which:
63/// - Are independent from those in any default partition.
64/// - Are not shared with other custom partitions, even if they have the same name.
65///
66/// To share a token bucket and rate limiter among custom partitions,
67/// either clone the custom partition itself or clone these components
68/// beforehand and pass them to each custom partition.
69#[non_exhaustive]
70#[derive(Clone, Debug)]
71pub struct RetryPartition {
72    pub(crate) inner: RetryPartitionInner,
73}
74
75#[derive(Clone, Debug)]
76pub(crate) enum RetryPartitionInner {
77    Default(Cow<'static, str>),
78    Custom {
79        name: Cow<'static, str>,
80        token_bucket: TokenBucket,
81        client_rate_limiter: ClientRateLimiter,
82    },
83}
84
85impl RetryPartition {
86    /// Creates a new `RetryPartition` from the given `name`.
87    pub fn new(name: impl Into<Cow<'static, str>>) -> Self {
88        Self {
89            inner: RetryPartitionInner::Default(name.into()),
90        }
91    }
92
93    /// Creates a builder for a custom `RetryPartition`.
94    pub fn custom(name: impl Into<Cow<'static, str>>) -> RetryPartitionBuilder {
95        RetryPartitionBuilder {
96            name: name.into(),
97            token_bucket: None,
98            client_rate_limiter: None,
99        }
100    }
101
102    fn name(&self) -> &str {
103        match &self.inner {
104            RetryPartitionInner::Default(name) => name,
105            RetryPartitionInner::Custom { name, .. } => name,
106        }
107    }
108}
109
110impl PartialEq for RetryPartition {
111    fn eq(&self, other: &Self) -> bool {
112        match (&self.inner, &other.inner) {
113            (RetryPartitionInner::Default(name1), RetryPartitionInner::Default(name2)) => {
114                name1 == name2
115            }
116            (
117                RetryPartitionInner::Custom { name: name1, .. },
118                RetryPartitionInner::Custom { name: name2, .. },
119            ) => name1 == name2,
120            // Different variants: not equal
121            _ => false,
122        }
123    }
124}
125
126impl Eq for RetryPartition {}
127
128impl std::hash::Hash for RetryPartition {
129    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
130        match &self.inner {
131            RetryPartitionInner::Default(name) => {
132                // Hash discriminant for Default variant
133                0u8.hash(state);
134                name.hash(state);
135            }
136            RetryPartitionInner::Custom { name, .. } => {
137                // Hash discriminant for Configured variant
138                1u8.hash(state);
139                name.hash(state);
140            }
141        }
142    }
143}
144
145impl fmt::Display for RetryPartition {
146    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
147        f.write_str(self.name())
148    }
149}
150
151impl Storable for RetryPartition {
152    type Storer = StoreReplace<RetryPartition>;
153}
154
155/// Builder for creating custom retry partitions.
156pub struct RetryPartitionBuilder {
157    name: Cow<'static, str>,
158    token_bucket: Option<TokenBucket>,
159    client_rate_limiter: Option<ClientRateLimiter>,
160}
161
162impl RetryPartitionBuilder {
163    /// Sets the token bucket.
164    pub fn token_bucket(mut self, token_bucket: TokenBucket) -> Self {
165        self.token_bucket = Some(token_bucket);
166        self
167    }
168
169    /// Sets the client rate limiter.
170    pub fn client_rate_limiter(mut self, client_rate_limiter: ClientRateLimiter) -> Self {
171        self.client_rate_limiter = Some(client_rate_limiter);
172        self
173    }
174
175    /// Builds the custom retry partition.
176    pub fn build(self) -> RetryPartition {
177        RetryPartition {
178            inner: RetryPartitionInner::Custom {
179                name: self.name,
180                token_bucket: self.token_bucket.unwrap_or_default(),
181                client_rate_limiter: self.client_rate_limiter.unwrap_or_default(),
182            },
183        }
184    }
185}
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use std::collections::hash_map::DefaultHasher;
190    use std::hash::{Hash, Hasher};
191
192    fn hash_value<T: Hash>(t: &T) -> u64 {
193        let mut hasher = DefaultHasher::new();
194        t.hash(&mut hasher);
195        hasher.finish()
196    }
197
198    #[test]
199    fn test_retry_partition_equality() {
200        let default1 = RetryPartition::new("test");
201        let default2 = RetryPartition::new("test");
202        let default3 = RetryPartition::new("other");
203
204        let configured1 = RetryPartition::custom("test").build();
205        let configured2 = RetryPartition::custom("test").build();
206        let configured3 = RetryPartition::custom("other").build();
207
208        // Same variant, same name
209        assert_eq!(default1, default2);
210        assert_eq!(configured1, configured2);
211
212        // Same variant, different name
213        assert_ne!(default1, default3);
214        assert_ne!(configured1, configured3);
215
216        // Different variant, same name
217        assert_ne!(default1, configured1);
218    }
219
220    #[test]
221    fn test_retry_partition_hash() {
222        let default = RetryPartition::new("test");
223        let configured = RetryPartition::custom("test").build();
224
225        // Different variants with same name should have different hashes
226        assert_ne!(hash_value(&default), hash_value(&configured));
227
228        // Same variants with same name should have same hashes
229        let default2 = RetryPartition::new("test");
230        assert_eq!(hash_value(&default), hash_value(&default2));
231    }
232}
aws_smithy_runtime/client/retries.rs

aws_smithy_runtime/client/
retries.rs
