aws_smithy_runtime/client/retries/
token_bucket.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use aws_smithy_async::time::TimeSource;
7use aws_smithy_types::config_bag::{Storable, StoreReplace};
8use aws_smithy_types::retry::ErrorKind;
9use std::fmt;
10use std::sync::atomic::AtomicU32;
11use std::sync::atomic::Ordering;
12use std::sync::Arc;
13use std::time::{Duration, SystemTime};
14use tokio::sync::{OwnedSemaphorePermit, Semaphore};
15
16pub(crate) const DEFAULT_CAPACITY: usize = 500;
17// On a 32 bit architecture, the value of Semaphore::MAX_PERMITS is 536,870,911.
18// Therefore, we will enforce a value lower than that to ensure behavior is
19// identical across platforms.
20// This also allows room for slight bucket overfill in the case where a bucket
21// is at maximum capacity and another thread drops a permit it was holding.
22/// The maximum number of permits a token bucket can have.
23pub const MAXIMUM_CAPACITY: usize = 500_000_000;
24#[allow(dead_code)]
25pub(crate) const DEFAULT_RETRY_COST: u32 = 14;
26#[allow(dead_code)]
27pub(crate) const DEFAULT_RETRY_TIMEOUT_COST: u32 = 14;
28#[allow(dead_code)]
29pub(crate) const THROTTLING_RETRY_COST: u32 = 5;
30
31// Legacy (Retry 2.0) costs
32const LEGACY_RETRY_COST: u32 = 5;
33const LEGACY_RETRY_TIMEOUT_COST: u32 = LEGACY_RETRY_COST * 2;
34const PERMIT_REGENERATION_AMOUNT: usize = 1;
35const DEFAULT_SUCCESS_REWARD: f32 = 0.0;
36
37/// Token bucket that bounds retries — a client-side retry quota.
38///
39/// Each retry attempt acquires tokens from the bucket; successful requests return or refill
40/// tokens. When the bucket is empty, the retry strategy stops retrying and returns the error to
41/// the caller, even if `max_attempts` has not been reached.
42///
43/// A single token bucket is shared by every operation in a [`RetryPartition`]. Use
44/// [`TokenBucket::builder`] to configure its [`capacity`](TokenBucketBuilder::capacity) and
45/// [`refill_rate`](TokenBucketBuilder::refill_rate), and attach it to a custom [`RetryPartition`]
46/// to give a workload its own bucket instead of sharing the default one.
47///
48/// ```
49/// use aws_smithy_runtime::client::retries::TokenBucket;
50///
51/// let token_bucket = TokenBucket::builder()
52///     .capacity(5000)
53///     .refill_rate(100.0) // tokens regenerated per second
54///     .build();
55/// ```
56///
57/// A custom [`RetryPartition`] carrying this bucket can then be set on a generated client's config
58/// builder via its `retry_partition` method.
59///
60/// [`RetryPartition`]: crate::client::retries::RetryPartition
61#[derive(Clone, Debug)]
62pub struct TokenBucket {
63    semaphore: Arc<Semaphore>,
64    max_permits: usize,
65    timeout_retry_cost: u32,
66    retry_cost: u32,
67    throttling_retry_cost: u32,
68    success_reward: f32,
69    fractional_tokens: Arc<AtomicF32>,
70    refill_rate: f32,
71    // Note this value is only an AtomicU32 so it works on 32bit powerpc architectures.
72    // If we ever remove the need for that compatibility it should become an AtomicU64
73    last_refill_time_secs: Arc<AtomicU32>,
74}
75
76impl std::panic::UnwindSafe for AtomicF32 {}
77impl std::panic::RefUnwindSafe for AtomicF32 {}
78struct AtomicF32 {
79    storage: AtomicU32,
80}
81impl AtomicF32 {
82    fn new(value: f32) -> Self {
83        let as_u32 = value.to_bits();
84        Self {
85            storage: AtomicU32::new(as_u32),
86        }
87    }
88    fn store(&self, value: f32) {
89        let as_u32 = value.to_bits();
90        self.storage.store(as_u32, Ordering::Relaxed)
91    }
92    fn load(&self) -> f32 {
93        let as_u32 = self.storage.load(Ordering::Relaxed);
94        f32::from_bits(as_u32)
95    }
96}
97
98impl fmt::Debug for AtomicF32 {
99    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100        // Use debug_struct, debug_tuple, or write! for formatting
101        f.debug_struct("AtomicF32")
102            .field("value", &self.load())
103            .finish()
104    }
105}
106
107impl Clone for AtomicF32 {
108    fn clone(&self) -> Self {
109        // Manually clone each field
110        AtomicF32 {
111            storage: AtomicU32::new(self.storage.load(Ordering::Relaxed)),
112        }
113    }
114}
115
116impl Storable for TokenBucket {
117    type Storer = StoreReplace<Self>;
118}
119
120impl Default for TokenBucket {
121    fn default() -> Self {
122        Self {
123            semaphore: Arc::new(Semaphore::new(DEFAULT_CAPACITY)),
124            max_permits: DEFAULT_CAPACITY,
125            timeout_retry_cost: LEGACY_RETRY_TIMEOUT_COST,
126            retry_cost: LEGACY_RETRY_COST,
127            throttling_retry_cost: LEGACY_RETRY_COST,
128            success_reward: DEFAULT_SUCCESS_REWARD,
129            fractional_tokens: Arc::new(AtomicF32::new(0.0)),
130            refill_rate: 0.0,
131            last_refill_time_secs: Arc::new(AtomicU32::new(0)),
132        }
133    }
134}
135
136impl TokenBucket {
137    /// Creates a new `TokenBucket` with the given initial quota.
138    pub fn new(initial_quota: usize) -> Self {
139        Self {
140            semaphore: Arc::new(Semaphore::new(initial_quota)),
141            max_permits: initial_quota,
142            ..Default::default()
143        }
144    }
145
146    /// A token bucket with unlimited capacity that allows retries at no cost.
147    pub fn unlimited() -> Self {
148        Self {
149            semaphore: Arc::new(Semaphore::new(MAXIMUM_CAPACITY)),
150            max_permits: MAXIMUM_CAPACITY,
151            timeout_retry_cost: 0,
152            retry_cost: 0,
153            throttling_retry_cost: 0,
154            success_reward: 0.0,
155            fractional_tokens: Arc::new(AtomicF32::new(0.0)),
156            refill_rate: 0.0,
157            last_refill_time_secs: Arc::new(AtomicU32::new(0)),
158        }
159    }
160
161    /// Creates a builder for constructing a `TokenBucket`.
162    pub fn builder() -> TokenBucketBuilder {
163        TokenBucketBuilder::default()
164    }
165
166    pub(crate) fn acquire(
167        &self,
168        err: &ErrorKind,
169        time_source: &impl TimeSource,
170    ) -> Option<OwnedSemaphorePermit> {
171        // Add time-based tokens to fractional accumulator
172        self.refill_tokens_based_on_time(time_source);
173        // Convert accumulated fractional tokens to whole tokens
174        self.convert_fractional_tokens();
175
176        let retry_cost = match err {
177            ErrorKind::TransientError => self.timeout_retry_cost,
178            ErrorKind::ThrottlingError => self.throttling_retry_cost,
179            _ => self.retry_cost,
180        };
181
182        self.semaphore
183            .clone()
184            .try_acquire_many_owned(retry_cost)
185            .ok()
186    }
187
188    pub(crate) fn success_reward(&self) -> f32 {
189        self.success_reward
190    }
191
192    pub(crate) fn regenerate_a_token(&self) {
193        self.add_permits(PERMIT_REGENERATION_AMOUNT);
194    }
195
196    /// Converts accumulated fractional tokens to whole tokens and adds them as permits.
197    /// Stores the remaining fractional amount back.
198    /// This is shared by both time-based refill and success rewards.
199    #[inline]
200    fn convert_fractional_tokens(&self) {
201        let mut calc_fractional_tokens = self.fractional_tokens.load();
202        // Verify that fractional tokens have not become corrupted - if they have, reset to zero
203        if !calc_fractional_tokens.is_finite() {
204            tracing::error!(
205                "Fractional tokens corrupted to: {}, resetting to 0.0",
206                calc_fractional_tokens
207            );
208            self.fractional_tokens.store(0.0);
209            return;
210        }
211
212        let full_tokens_accumulated = calc_fractional_tokens.floor();
213        if full_tokens_accumulated >= 1.0 {
214            self.add_permits(full_tokens_accumulated as usize);
215            calc_fractional_tokens -= full_tokens_accumulated;
216        }
217        // Always store the updated fractional tokens back, even if no conversion happened
218        self.fractional_tokens.store(calc_fractional_tokens);
219    }
220
221    /// Refills tokens based on elapsed time since last refill.
222    /// This method implements lazy evaluation - tokens are only calculated when accessed.
223    /// Uses a single compare-and-swap to ensure only one thread processes each time window.
224    #[inline]
225    fn refill_tokens_based_on_time(&self, time_source: &impl TimeSource) {
226        if self.refill_rate > 0.0 {
227            // The cast to u32 here is safe until 2106, and I will be long dead then so ¯\_(ツ)_/¯
228            let current_time_secs = time_source
229                .now()
230                .duration_since(SystemTime::UNIX_EPOCH)
231                .unwrap_or(Duration::ZERO)
232                .as_secs() as u32;
233
234            let last_refill_secs = self.last_refill_time_secs.load(Ordering::Relaxed);
235
236            // Early exit if no time elapsed - most threads take this path
237            if current_time_secs == last_refill_secs {
238                return;
239            }
240
241            // Try to atomically claim this time window with a single CAS
242            // If we lose, another thread is handling the refill, so we can exit
243            if self
244                .last_refill_time_secs
245                .compare_exchange(
246                    last_refill_secs,
247                    current_time_secs,
248                    Ordering::Relaxed,
249                    Ordering::Relaxed,
250                )
251                .is_err()
252            {
253                // Another thread claimed this time window, we're done
254                return;
255            }
256
257            // We won the CAS - we're responsible for adding tokens for this time window
258            let current_fractional = self.fractional_tokens.load();
259            let max_fractional = self.max_permits as f32;
260
261            // Skip token addition if already at cap
262            if current_fractional >= max_fractional {
263                return;
264            }
265
266            let elapsed_secs = current_time_secs.saturating_sub(last_refill_secs);
267            let tokens_to_add = elapsed_secs as f32 * self.refill_rate;
268
269            // Add tokens to fractional accumulator, capping at max_permits to prevent unbounded growth
270            let new_fractional = (current_fractional + tokens_to_add).min(max_fractional);
271            self.fractional_tokens.store(new_fractional);
272        }
273    }
274
275    #[inline]
276    pub(crate) fn reward_success(&self) {
277        if self.success_reward > 0.0 {
278            let current = self.fractional_tokens.load();
279            let max_fractional = self.max_permits as f32;
280            // Early exit if already at cap - no point calculating
281            if current >= max_fractional {
282                return;
283            }
284            // Cap fractional tokens at max_permits to prevent unbounded growth
285            let new_fractional = (current + self.success_reward).min(max_fractional);
286            self.fractional_tokens.store(new_fractional);
287        }
288    }
289
290    pub(crate) fn add_permits(&self, amount: usize) {
291        let available = self.semaphore.available_permits();
292        if available >= self.max_permits {
293            return;
294        }
295        self.semaphore
296            .add_permits(amount.min(self.max_permits - available));
297    }
298
299    /// Returns true if the token bucket is full, false otherwise
300    pub fn is_full(&self) -> bool {
301        self.convert_fractional_tokens();
302        self.semaphore.available_permits() >= self.max_permits
303    }
304
305    /// Returns true if the token bucket is empty, false otherwise
306    pub fn is_empty(&self) -> bool {
307        self.convert_fractional_tokens();
308        self.semaphore.available_permits() == 0
309    }
310
311    #[allow(dead_code)] // only used in tests
312    #[cfg(any(test, feature = "test-util", feature = "legacy-test-util"))]
313    pub(crate) fn available_permits(&self) -> usize {
314        self.semaphore.available_permits()
315    }
316
317    /// Only used in tests
318    #[allow(dead_code)]
319    #[doc(hidden)]
320    #[cfg(any(test, feature = "test-util", feature = "legacy-test-util"))]
321    pub fn last_refill_time_secs(&self) -> Arc<AtomicU32> {
322        self.last_refill_time_secs.clone()
323    }
324}
325
326/// Builder for constructing a `TokenBucket`.
327#[derive(Clone, Debug, Default)]
328pub struct TokenBucketBuilder {
329    capacity: Option<usize>,
330    retry_cost: Option<u32>,
331    throttling_retry_cost: Option<u32>,
332    timeout_retry_cost: Option<u32>,
333    success_reward: Option<f32>,
334    refill_rate: Option<f32>,
335}
336
337impl TokenBucketBuilder {
338    /// Creates a new `TokenBucketBuilder` with default values.
339    pub fn new() -> Self {
340        Self::default()
341    }
342
343    /// Sets the maximum bucket capacity for the builder.
344    pub fn capacity(mut self, mut capacity: usize) -> Self {
345        if capacity > MAXIMUM_CAPACITY {
346            capacity = MAXIMUM_CAPACITY;
347        }
348        self.capacity = Some(capacity);
349        self
350    }
351
352    /// Sets the specified retry cost for the builder.
353    pub fn retry_cost(mut self, retry_cost: u32) -> Self {
354        self.retry_cost = Some(retry_cost);
355        self
356    }
357
358    /// Sets the throttling retry cost for the builder.
359    pub fn throttling_retry_cost(mut self, throttling_retry_cost: u32) -> Self {
360        self.throttling_retry_cost = Some(throttling_retry_cost);
361        self
362    }
363
364    /// Sets the specified timeout retry cost for the builder.
365    pub fn timeout_retry_cost(mut self, timeout_retry_cost: u32) -> Self {
366        self.timeout_retry_cost = Some(timeout_retry_cost);
367        self
368    }
369
370    /// Sets the reward for any successful request for the builder.
371    pub fn success_reward(mut self, reward: f32) -> Self {
372        self.success_reward = Some(reward);
373        self
374    }
375
376    /// Sets the refill rate (tokens per second) for time-based token regeneration.
377    ///
378    /// Negative values are clamped to 0.0. A refill rate of 0.0 disables time-based regeneration.
379    /// Non-finite values (NaN, infinity) are treated as 0.0.
380    pub fn refill_rate(mut self, rate: f32) -> Self {
381        let validated_rate = if rate.is_finite() { rate.max(0.0) } else { 0.0 };
382        self.refill_rate = Some(validated_rate);
383        self
384    }
385
386    /// Builds a `TokenBucket`.
387    pub fn build(self) -> TokenBucket {
388        TokenBucket {
389            semaphore: Arc::new(Semaphore::new(self.capacity.unwrap_or(DEFAULT_CAPACITY))),
390            max_permits: self.capacity.unwrap_or(DEFAULT_CAPACITY),
391            retry_cost: self.retry_cost.unwrap_or(LEGACY_RETRY_COST),
392            throttling_retry_cost: self.throttling_retry_cost.unwrap_or(LEGACY_RETRY_COST),
393            timeout_retry_cost: self.timeout_retry_cost.unwrap_or(LEGACY_RETRY_TIMEOUT_COST),
394            success_reward: self.success_reward.unwrap_or(DEFAULT_SUCCESS_REWARD),
395            fractional_tokens: Arc::new(AtomicF32::new(0.0)),
396            refill_rate: self.refill_rate.unwrap_or(0.0),
397            last_refill_time_secs: Arc::new(AtomicU32::new(0)),
398        }
399    }
400}
401
402#[cfg(test)]
403mod tests {
404
405    use super::*;
406    use aws_smithy_async::test_util::ManualTimeSource;
407    use std::{sync::LazyLock, time::UNIX_EPOCH};
408
409    static TIME_SOURCE: LazyLock<ManualTimeSource> =
410        LazyLock::new(|| ManualTimeSource::new(UNIX_EPOCH + Duration::from_secs(12344321)));
411
412    #[test]
413    fn test_unlimited_token_bucket() {
414        let bucket = TokenBucket::unlimited();
415
416        // Should always acquire permits regardless of error type
417        assert!(bucket
418            .acquire(&ErrorKind::ThrottlingError, &*TIME_SOURCE)
419            .is_some());
420        assert!(bucket
421            .acquire(&ErrorKind::TransientError, &*TIME_SOURCE)
422            .is_some());
423
424        // Should have maximum capacity
425        assert_eq!(bucket.max_permits, MAXIMUM_CAPACITY);
426
427        // Should have zero retry costs
428        assert_eq!(bucket.retry_cost, 0);
429        assert_eq!(bucket.timeout_retry_cost, 0);
430
431        // The loop count is arbitrary; should obtain permits without limit
432        let mut permits = Vec::new();
433        for _ in 0..100 {
434            let permit = bucket.acquire(&ErrorKind::ThrottlingError, &*TIME_SOURCE);
435            assert!(permit.is_some());
436            permits.push(permit);
437            // Available permits should stay constant
438            assert_eq!(MAXIMUM_CAPACITY, bucket.semaphore.available_permits());
439        }
440    }
441
442    #[test]
443    fn test_bounded_permits_exhaustion() {
444        let bucket = TokenBucket::new(10);
445        let mut permits = Vec::new();
446
447        for _ in 0..100 {
448            let permit = bucket.acquire(&ErrorKind::ThrottlingError, &*TIME_SOURCE);
449            if let Some(p) = permit {
450                permits.push(p);
451            } else {
452                break;
453            }
454        }
455
456        assert_eq!(permits.len(), 2); // 10 capacity / 5 retry cost = 2 permits
457
458        // Verify next acquisition fails
459        assert!(bucket
460            .acquire(&ErrorKind::ThrottlingError, &*TIME_SOURCE)
461            .is_none());
462    }
463
464    #[test]
465    fn test_fractional_tokens_accumulate_and_convert() {
466        let bucket = TokenBucket::builder()
467            .capacity(10)
468            .success_reward(0.4)
469            .build();
470
471        // acquire 10 tokens to bring capacity below max so we can test accumulation
472        let _hold_permit = bucket.acquire(&ErrorKind::TransientError, &*TIME_SOURCE);
473        assert_eq!(bucket.semaphore.available_permits(), 0);
474
475        // First success: 0.4 fractional tokens
476        bucket.reward_success();
477        bucket.convert_fractional_tokens();
478        assert_eq!(bucket.semaphore.available_permits(), 0);
479
480        // Second success: 0.8 fractional tokens
481        bucket.reward_success();
482        bucket.convert_fractional_tokens();
483        assert_eq!(bucket.semaphore.available_permits(), 0);
484
485        // Third success: 1.2 fractional tokens -> 1 full token added
486        bucket.reward_success();
487        bucket.convert_fractional_tokens();
488        assert_eq!(bucket.semaphore.available_permits(), 1);
489    }
490
491    #[test]
492    fn test_fractional_tokens_respect_max_capacity() {
493        let bucket = TokenBucket::builder()
494            .capacity(10)
495            .success_reward(2.0)
496            .build();
497
498        for _ in 0..20 {
499            bucket.reward_success();
500        }
501
502        assert!(bucket.semaphore.available_permits() == 10);
503    }
504
505    #[test]
506    fn test_convert_fractional_tokens() {
507        // (input, expected_permits_added, expected_remaining)
508        let test_cases = [
509            (0.7, 0, 0.7),
510            (1.0, 1, 0.0),
511            (2.3, 2, 0.3),
512            (5.8, 5, 0.8),
513            (10.0, 10, 0.0),
514            // verify that if fractional permits are corrupted, we reset to 0 gracefully
515            (f32::NAN, 0, 0.0),
516            (f32::INFINITY, 0, 0.0),
517        ];
518
519        for (input, expected_permits, expected_remaining) in test_cases {
520            let bucket = TokenBucket::builder().capacity(10).build();
521            let _hold_permit = bucket.acquire(&ErrorKind::TransientError, &*TIME_SOURCE);
522            let initial = bucket.semaphore.available_permits();
523
524            bucket.fractional_tokens.store(input);
525            bucket.convert_fractional_tokens();
526
527            assert_eq!(
528                bucket.semaphore.available_permits() - initial,
529                expected_permits
530            );
531            assert!((bucket.fractional_tokens.load() - expected_remaining).abs() < 0.0001);
532        }
533    }
534
535    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
536    #[test]
537    fn test_builder_with_custom_values() {
538        let bucket = TokenBucket::builder()
539            .capacity(100)
540            .retry_cost(10)
541            .timeout_retry_cost(20)
542            .success_reward(0.5)
543            .refill_rate(2.5)
544            .build();
545
546        assert_eq!(bucket.max_permits, 100);
547        assert_eq!(bucket.retry_cost, 10);
548        assert_eq!(bucket.timeout_retry_cost, 20);
549        assert_eq!(bucket.success_reward, 0.5);
550        assert_eq!(bucket.refill_rate, 2.5);
551    }
552
553    #[test]
554    fn test_builder_refill_rate_validation() {
555        // Test negative values are clamped to 0.0
556        let bucket = TokenBucket::builder().refill_rate(-5.0).build();
557        assert_eq!(bucket.refill_rate, 0.0);
558
559        // Test valid positive value
560        let bucket = TokenBucket::builder().refill_rate(1.5).build();
561        assert_eq!(bucket.refill_rate, 1.5);
562
563        // Test zero is valid
564        let bucket = TokenBucket::builder().refill_rate(0.0).build();
565        assert_eq!(bucket.refill_rate, 0.0);
566    }
567
568    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
569    #[test]
570    fn test_builder_custom_time_source() {
571        use aws_smithy_async::test_util::ManualTimeSource;
572        use std::time::UNIX_EPOCH;
573
574        // Test that TokenBucket uses provided TimeSource when specified via builder
575        let manual_time = ManualTimeSource::new(UNIX_EPOCH);
576        let bucket = TokenBucket::builder()
577            .capacity(100)
578            .refill_rate(1.0)
579            .build();
580
581        // Consume all tokens to test refill from empty state
582        let _permits = bucket.semaphore.try_acquire_many(100).unwrap();
583        assert_eq!(bucket.available_permits(), 0);
584
585        // Advance time and verify tokens are added based on manual time
586        manual_time.advance(Duration::from_secs(5));
587
588        bucket.refill_tokens_based_on_time(&manual_time);
589        bucket.convert_fractional_tokens();
590
591        // Should have 5 tokens (5 seconds * 1 token/sec)
592        assert_eq!(bucket.available_permits(), 5);
593    }
594
595    #[test]
596    fn test_atomicf32_f32_to_bits_conversion_correctness() {
597        // This is the core functionality
598        let test_values = vec![
599            0.0,
600            -0.0,
601            1.0,
602            -1.0,
603            f32::INFINITY,
604            f32::NEG_INFINITY,
605            f32::NAN,
606            f32::MIN,
607            f32::MAX,
608            f32::MIN_POSITIVE,
609            f32::EPSILON,
610            std::f32::consts::PI,
611            std::f32::consts::E,
612            // Test values that could expose bit manipulation bugs
613            1.23456789e-38, // Very small normal number
614            1.23456789e38,  // Very large number (within f32 range)
615            1.1754944e-38,  // Near MIN_POSITIVE for f32
616        ];
617
618        for &expected in &test_values {
619            let atomic = AtomicF32::new(expected);
620            let actual = atomic.load();
621
622            // For NaN, we can't use == but must check bit patterns
623            if expected.is_nan() {
624                assert!(actual.is_nan(), "Expected NaN, got {}", actual);
625                // Different NaN bit patterns should be preserved exactly
626                assert_eq!(expected.to_bits(), actual.to_bits());
627            } else {
628                assert_eq!(expected.to_bits(), actual.to_bits());
629            }
630        }
631    }
632
633    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
634    #[test]
635    fn test_atomicf32_store_load_preserves_exact_bits() {
636        let atomic = AtomicF32::new(0.0);
637
638        // Test that store/load cycle preserves EXACT bit patterns
639        // This would catch bugs in the to_bits/from_bits conversion
640        let critical_bit_patterns = vec![
641            0x00000000u32, // +0.0
642            0x80000000u32, // -0.0
643            0x7F800000u32, // +infinity
644            0xFF800000u32, // -infinity
645            0x7FC00000u32, // Quiet NaN
646            0x7FA00000u32, // Signaling NaN
647            0x00000001u32, // Smallest positive subnormal
648            0x007FFFFFu32, // Largest subnormal
649            0x00800000u32, // Smallest positive normal (MIN_POSITIVE)
650        ];
651
652        for &expected_bits in &critical_bit_patterns {
653            let expected_f32 = f32::from_bits(expected_bits);
654            atomic.store(expected_f32);
655            let loaded_f32 = atomic.load();
656            let actual_bits = loaded_f32.to_bits();
657
658            assert_eq!(expected_bits, actual_bits);
659        }
660    }
661
662    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
663    #[test]
664    fn test_atomicf32_concurrent_store_load_safety() {
665        use std::sync::Arc;
666        use std::thread;
667
668        let atomic = Arc::new(AtomicF32::new(0.0));
669        let test_values = vec![1.0, 2.0, 3.0, 4.0, 5.0];
670        let mut handles = Vec::new();
671
672        // Start multiple threads that continuously write different values
673        for &value in &test_values {
674            let atomic_clone = Arc::clone(&atomic);
675            let handle = thread::spawn(move || {
676                for _ in 0..1000 {
677                    atomic_clone.store(value);
678                }
679            });
680            handles.push(handle);
681        }
682
683        // Start a reader thread that continuously reads
684        let atomic_reader = Arc::clone(&atomic);
685        let reader_handle = thread::spawn(move || {
686            let mut readings = Vec::new();
687            for _ in 0..5000 {
688                let value = atomic_reader.load();
689                readings.push(value);
690            }
691            readings
692        });
693
694        // Wait for all writers to complete
695        for handle in handles {
696            handle.join().expect("Writer thread panicked");
697        }
698
699        let readings = reader_handle.join().expect("Reader thread panicked");
700
701        // Verify that all read values are valid (one of the written values)
702        // This tests that there's no data corruption from concurrent access
703        for &reading in &readings {
704            assert!(test_values.contains(&reading) || reading == 0.0);
705
706            // More importantly, verify the reading is a valid f32
707            // (not corrupted bits that happen to parse as valid)
708            assert!(
709                reading.is_finite() || reading == 0.0,
710                "Corrupted reading detected"
711            );
712        }
713    }
714
715    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
716    #[test]
717    fn test_atomicf32_stress_concurrent_access() {
718        use std::sync::{Arc, Barrier};
719        use std::thread;
720
721        let expected_values = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
722        let atomic = Arc::new(AtomicF32::new(0.0));
723        let barrier = Arc::new(Barrier::new(10)); // Synchronize all threads
724        let mut handles = Vec::new();
725
726        // Launch threads that all start simultaneously
727        for i in 0..10 {
728            let atomic_clone = Arc::clone(&atomic);
729            let barrier_clone = Arc::clone(&barrier);
730            let handle = thread::spawn(move || {
731                barrier_clone.wait(); // All threads start at same time
732
733                // Tight loop increases chance of race conditions
734                for _ in 0..10000 {
735                    let value = i as f32;
736                    atomic_clone.store(value);
737                    let loaded = atomic_clone.load();
738                    // Verify no corruption occurred
739                    assert!(loaded >= 0.0 && loaded <= 9.0);
740                    assert!(
741                        expected_values.contains(&loaded),
742                        "Got unexpected value: {}, expected one of {:?}",
743                        loaded,
744                        expected_values
745                    );
746                }
747            });
748            handles.push(handle);
749        }
750
751        for handle in handles {
752            handle.join().unwrap();
753        }
754    }
755
756    #[test]
757    fn test_atomicf32_integration_with_token_bucket_usage() {
758        let atomic = AtomicF32::new(0.0);
759        let success_reward = 0.3;
760        let iterations = 5;
761
762        // Accumulate fractional tokens
763        for _ in 1..=iterations {
764            let current = atomic.load();
765            atomic.store(current + success_reward);
766        }
767
768        let accumulated = atomic.load();
769        let expected_total = iterations as f32 * success_reward; // 1.5
770
771        // Test the floor() operation pattern
772        let full_tokens = accumulated.floor();
773        atomic.store(accumulated - full_tokens);
774        let remaining = atomic.load();
775
776        // These assertions should be general:
777        assert_eq!(full_tokens, expected_total.floor()); // Could be 1.0, 2.0, 3.0, etc.
778        assert!(remaining >= 0.0 && remaining < 1.0);
779        assert_eq!(remaining, expected_total - expected_total.floor());
780    }
781
782    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
783    #[test]
784    fn test_atomicf32_clone_creates_independent_copy() {
785        let original = AtomicF32::new(123.456);
786        let cloned = original.clone();
787
788        // Verify they start with the same value
789        assert_eq!(original.load(), cloned.load());
790
791        // Verify they're independent - modifying one doesn't affect the other
792        original.store(999.0);
793        assert_eq!(
794            cloned.load(),
795            123.456,
796            "Clone should be unaffected by original changes"
797        );
798        assert_eq!(original.load(), 999.0, "Original should have new value");
799    }
800
801    #[test]
802    fn test_combined_time_and_success_rewards() {
803        use aws_smithy_async::test_util::ManualTimeSource;
804        use std::time::UNIX_EPOCH;
805
806        let time_source = ManualTimeSource::new(UNIX_EPOCH);
807        let current_time_secs = UNIX_EPOCH
808            .duration_since(SystemTime::UNIX_EPOCH)
809            .unwrap()
810            .as_secs() as u32;
811
812        let bucket = TokenBucket {
813            refill_rate: 1.0,
814            success_reward: 0.5,
815            last_refill_time_secs: Arc::new(AtomicU32::new(current_time_secs)),
816            semaphore: Arc::new(Semaphore::new(0)),
817            max_permits: 100,
818            ..Default::default()
819        };
820
821        // Add success rewards: 2 * 0.5 = 1.0 token
822        bucket.reward_success();
823        bucket.reward_success();
824
825        // Advance time by 2 seconds
826        time_source.advance(Duration::from_secs(2));
827
828        // Trigger time-based refill: 2 sec * 1.0 = 2.0 tokens
829        // Total: 1.0 + 2.0 = 3.0 tokens
830        bucket.refill_tokens_based_on_time(&time_source);
831        bucket.convert_fractional_tokens();
832
833        assert_eq!(bucket.available_permits(), 3);
834        assert!(bucket.fractional_tokens.load().abs() < 0.0001);
835    }
836
837    #[test]
838    fn test_refill_rates() {
839        use aws_smithy_async::test_util::ManualTimeSource;
840        use std::time::UNIX_EPOCH;
841        // (refill_rate, elapsed_secs, expected_permits, expected_fractional)
842        let test_cases = [
843            (10.0, 2, 20, 0.0),      // Basic: 2 sec * 10 tokens/sec = 20 tokens
844            (0.001, 1100, 1, 0.1),   // Small: 1100 * 0.001 = 1.1 tokens
845            (0.0001, 11000, 1, 0.1), // Tiny: 11000 * 0.0001 = 1.1 tokens
846            (0.001, 1200, 1, 0.2),   // 1200 * 0.001 = 1.2 tokens
847            (0.0001, 10000, 1, 0.0), // 10000 * 0.0001 = 1.0 tokens
848            (0.001, 500, 0, 0.5),    // Fractional only: 500 * 0.001 = 0.5 tokens
849        ];
850
851        for (refill_rate, elapsed_secs, expected_permits, expected_fractional) in test_cases {
852            let time_source = ManualTimeSource::new(UNIX_EPOCH);
853            let current_time_secs = UNIX_EPOCH
854                .duration_since(SystemTime::UNIX_EPOCH)
855                .unwrap()
856                .as_secs() as u32;
857
858            let bucket = TokenBucket {
859                refill_rate,
860                last_refill_time_secs: Arc::new(AtomicU32::new(current_time_secs)),
861                semaphore: Arc::new(Semaphore::new(0)),
862                max_permits: 100,
863                ..Default::default()
864            };
865
866            // Advance time by the specified duration
867            time_source.advance(Duration::from_secs(elapsed_secs));
868
869            bucket.refill_tokens_based_on_time(&time_source);
870            bucket.convert_fractional_tokens();
871
872            assert_eq!(
873                bucket.available_permits(),
874                expected_permits,
875                "Rate {}: After {}s expected {} permits",
876                refill_rate,
877                elapsed_secs,
878                expected_permits
879            );
880            assert!(
881                (bucket.fractional_tokens.load() - expected_fractional).abs() < 0.0001,
882                "Rate {}: After {}s expected {} fractional, got {}",
883                refill_rate,
884                elapsed_secs,
885                expected_fractional,
886                bucket.fractional_tokens.load()
887            );
888        }
889    }
890
891    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
892    #[test]
893    fn test_rewards_capped_at_max_capacity() {
894        use aws_smithy_async::test_util::ManualTimeSource;
895        use std::time::UNIX_EPOCH;
896
897        let time_source = ManualTimeSource::new(UNIX_EPOCH);
898        let current_time_secs = UNIX_EPOCH
899            .duration_since(SystemTime::UNIX_EPOCH)
900            .unwrap()
901            .as_secs() as u32;
902
903        let bucket = TokenBucket {
904            refill_rate: 50.0,
905            success_reward: 2.0,
906            last_refill_time_secs: Arc::new(AtomicU32::new(current_time_secs)),
907            semaphore: Arc::new(Semaphore::new(5)),
908            max_permits: 10,
909            ..Default::default()
910        };
911
912        // Add success rewards: 50 * 2.0 = 100 tokens (without cap)
913        for _ in 0..50 {
914            bucket.reward_success();
915        }
916
917        // Fractional tokens capped at 10 from success rewards
918        assert_eq!(bucket.fractional_tokens.load(), 10.0);
919
920        // Advance time by 100 seconds
921        time_source.advance(Duration::from_secs(100));
922
923        // Time-based refill: 100 * 50 = 5000 tokens (without cap)
924        // But fractional is already at 10, so it stays at 10
925        bucket.refill_tokens_based_on_time(&time_source);
926
927        // Fractional tokens should be capped at max_permits (10)
928        assert_eq!(
929            bucket.fractional_tokens.load(),
930            10.0,
931            "Fractional tokens should be capped at max_permits"
932        );
933        // Convert should add 5 tokens (bucket at 5, can add 5 more to reach max 10)
934        bucket.convert_fractional_tokens();
935        assert_eq!(bucket.available_permits(), 10);
936    }
937
938    #[cfg(any(feature = "test-util", feature = "legacy-test-util"))]
939    #[test]
940    fn test_concurrent_time_based_refill_no_over_generation() {
941        use aws_smithy_async::test_util::ManualTimeSource;
942        use std::sync::{Arc, Barrier};
943        use std::thread;
944        use std::time::UNIX_EPOCH;
945
946        let time_source = ManualTimeSource::new(UNIX_EPOCH);
947        let current_time_secs = UNIX_EPOCH
948            .duration_since(SystemTime::UNIX_EPOCH)
949            .unwrap()
950            .as_secs() as u32;
951
952        // Create bucket with 1 token/sec refill
953        let bucket = Arc::new(TokenBucket {
954            refill_rate: 1.0,
955            last_refill_time_secs: Arc::new(AtomicU32::new(current_time_secs)),
956            semaphore: Arc::new(Semaphore::new(0)),
957            max_permits: 100,
958            ..Default::default()
959        });
960
961        // Advance time by 10 seconds
962        time_source.advance(Duration::from_secs(10));
963        let shared_time_source = aws_smithy_async::time::SharedTimeSource::new(time_source);
964
965        // Launch 100 threads that all try to refill simultaneously
966        let barrier = Arc::new(Barrier::new(100));
967        let mut handles = Vec::new();
968
969        for _ in 0..100 {
970            let bucket_clone1 = Arc::clone(&bucket);
971            let barrier_clone1 = Arc::clone(&barrier);
972            let time_source_clone1 = shared_time_source.clone();
973            let bucket_clone2 = Arc::clone(&bucket);
974            let barrier_clone2 = Arc::clone(&barrier);
975            let time_source_clone2 = shared_time_source.clone();
976
977            let handle1 = thread::spawn(move || {
978                // Wait for all threads to be ready
979                barrier_clone1.wait();
980
981                // All threads call refill at the same time
982                bucket_clone1.refill_tokens_based_on_time(&time_source_clone1);
983            });
984
985            let handle2 = thread::spawn(move || {
986                // Wait for all threads to be ready
987                barrier_clone2.wait();
988
989                // All threads call refill at the same time
990                bucket_clone2.refill_tokens_based_on_time(&time_source_clone2);
991            });
992            handles.push(handle1);
993            handles.push(handle2);
994        }
995
996        // Wait for all threads to complete
997        for handle in handles {
998            handle.join().unwrap();
999        }
1000
1001        // Convert fractional tokens to whole tokens
1002        bucket.convert_fractional_tokens();
1003
1004        // Should have exactly 10 tokens (10 seconds * 1 token/sec)
1005        // Not 1000 tokens (100 threads * 10 tokens each)
1006        assert_eq!(
1007            bucket.available_permits(),
1008            10,
1009            "Only one thread should have added tokens, not all 100"
1010        );
1011
1012        // Fractional should be 0 after conversion
1013        assert!(bucket.fractional_tokens.load().abs() < 0.0001);
1014    }
1015
1016    /// Regression test for https://github.com/awslabs/aws-sdk-rust/issues/1423
1017    #[test]
1018    fn test_is_full_accounts_for_fractional_tokens() {
1019        let bucket = TokenBucket::builder()
1020            .capacity(2)
1021            .retry_cost(1)
1022            .success_reward(0.9)
1023            .build();
1024
1025        assert!(bucket.is_full());
1026
1027        let _p1 = bucket
1028            .acquire(&ErrorKind::ServerError, &*TIME_SOURCE)
1029            .unwrap();
1030        let _p2 = bucket
1031            .acquire(&ErrorKind::ServerError, &*TIME_SOURCE)
1032            .unwrap();
1033
1034        assert!(bucket.is_empty());
1035
1036        // 3 rewards of 0.9 = 2.7 fractional tokens, which converts to 2 whole
1037        // permits — enough to fill the bucket (capacity 2).
1038        bucket.reward_success();
1039        bucket.reward_success();
1040        bucket.reward_success();
1041
1042        // Before the fix, is_full() returned false here because fractional
1043        // tokens hadn't been converted to real permits.
1044        assert!(bucket.is_full());
1045        assert!(!bucket.is_empty());
1046    }
1047
1048    #[test]
1049    fn test_is_empty_accounts_for_fractional_tokens() {
1050        let bucket = TokenBucket::builder()
1051            .capacity(10)
1052            .retry_cost(10)
1053            .success_reward(0.5)
1054            .build();
1055
1056        let _p = bucket
1057            .acquire(&ErrorKind::ServerError, &*TIME_SOURCE)
1058            .unwrap();
1059        assert_eq!(bucket.semaphore.available_permits(), 0);
1060
1061        // 0.5 fractional tokens can't convert to a whole permit
1062        bucket.reward_success();
1063        assert!(bucket.is_empty());
1064
1065        // 1.0 fractional tokens converts to a permit
1066        bucket.reward_success();
1067        assert!(!bucket.is_empty());
1068    }
1069}
aws_smithy_runtime/client/retries/token_bucket.rs

aws_smithy_runtime/client/retries/
token_bucket.rs
