do_memory_core/retry/
budget.rs

1use std::sync::Arc;
2use std::sync::atomic::{AtomicU32, Ordering};
3use std::time::{Duration, Instant};
4
5use parking_lot::Mutex;
6
7/// Process-wide retry budget with a sliding window.
8///
9/// Tokens refill at the start of each window. Multiple callers share
10/// the same [`Arc<RetryBudget>`] to coordinate retry pressure across
11/// the entire process.
12pub struct RetryBudget {
13    max_tokens: u32,
14    window: Duration,
15    tokens: AtomicU32,
16    window_start: Mutex<Instant>,
17}
18
19impl RetryBudget {
20    /// Create a shared retry budget.
21    ///
22    /// - `max_tokens`: how many retries are allowed per window.
23    /// - `window`: the sliding window duration.
24    pub fn new(max_tokens: u32, window: Duration) -> Arc<Self> {
25        Arc::new(Self {
26            max_tokens,
27            window,
28            tokens: AtomicU32::new(max_tokens),
29            window_start: Mutex::new(Instant::now()),
30        })
31    }
32
33    /// Try to acquire one retry token.
34    ///
35    /// Returns `true` if a token was available (and consumed), `false` if
36    /// the budget is exhausted for the current window.
37    pub fn acquire(&self) -> bool {
38        self.maybe_refill();
39        let prev = self.tokens.fetch_sub(1, Ordering::SeqCst);
40        if prev > 0 {
41            true
42        } else {
43            // We went below zero — put it back (saturating).
44            self.tokens.fetch_add(1, Ordering::SeqCst);
45            false
46        }
47    }
48
49    /// Number of tokens remaining in the current window.
50    pub fn available(&self) -> u32 {
51        self.maybe_refill();
52        self.tokens.load(Ordering::SeqCst)
53    }
54
55    fn maybe_refill(&self) {
56        let mut start = self.window_start.lock();
57        let elapsed = start.elapsed();
58        if elapsed >= self.window {
59            self.tokens.store(self.max_tokens, Ordering::SeqCst);
60            *start = Instant::now();
61        }
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn token_exhaustion() {
71        let budget = RetryBudget::new(2, Duration::from_secs(10));
72        assert!(budget.acquire());
73        assert!(budget.acquire());
74        assert!(!budget.acquire());
75        assert_eq!(budget.available(), 0);
76    }
77
78    #[test]
79    fn window_refill() {
80        let budget = RetryBudget::new(1, Duration::from_millis(50));
81        assert!(budget.acquire());
82        assert!(!budget.acquire());
83
84        std::thread::sleep(Duration::from_millis(60));
85        assert!(budget.acquire());
86        assert!(!budget.acquire());
87    }
88
89    #[test]
90    fn concurrent_acquisition() {
91        use std::sync::Arc;
92        use std::thread;
93
94        let budget = RetryBudget::new(5, Duration::from_secs(10));
95        let budget = Arc::new(budget);
96
97        let handles: Vec<_> = (0..10)
98            .map(|_| {
99                let b = Arc::clone(&budget);
100                thread::spawn(move || b.acquire())
101            })
102            .collect();
103
104        let mut acquired = 0;
105        for h in handles {
106            if h.join().unwrap() {
107                acquired += 1;
108            }
109        }
110        assert_eq!(acquired, 5);
111        assert_eq!(budget.available(), 0);
112    }
113}
do_memory_core/retry/budget.rs

do_memory_core/retry/
budget.rs
