execution_policy/retry/
budget.rs

1//! Shared retry budget: a lock-free token bucket that bounds retries to a
2//! fraction of total traffic, protecting dependencies from retry storms.
3//!
4//! Each top-level call deposits tokens; each retry withdraws one. When the
5//! bucket is empty, retries are denied even if attempts remain.
6
7use std::sync::Arc;
8use std::sync::atomic::{AtomicI64, Ordering};
9
10const SCALE: i64 = 1000;
11
12#[derive(Debug)]
13struct Inner {
14    tokens: AtomicI64,
15    max_tokens: i64,
16    deposit: i64,
17    retry_cost: i64,
18}
19
20/// A shared retry budget. Clone it into multiple policies to share one budget.
21#[derive(Debug, Clone)]
22pub struct RetryBudget(Arc<Inner>);
23
24impl RetryBudget {
25    /// Allow retries up to `ratio` of total calls (e.g. `0.2` = 20%), with a
26    /// burst of up to `burst` retries when the bucket is full.
27    pub fn new(ratio: f64, burst: u32) -> Self {
28        let ratio = ratio.clamp(0.0, 1.0);
29        let deposit = (ratio * SCALE as f64) as i64;
30        let max_tokens = (burst.max(1) as i64) * SCALE;
31        Self(Arc::new(Inner {
32            tokens: AtomicI64::new(max_tokens),
33            max_tokens,
34            deposit,
35            retry_cost: SCALE,
36        }))
37    }
38
39    /// Sensible default: 20% retry ratio, burst of 10.
40    pub fn standard() -> Self {
41        Self::new(0.2, 10)
42    }
43
44    /// Record a top-level call (deposits tokens, capped at the burst max).
45    pub(crate) fn deposit(&self) {
46        let inner = &self.0;
47        let mut cur = inner.tokens.load(Ordering::Relaxed);
48        loop {
49            let next = (cur + inner.deposit).min(inner.max_tokens);
50            match inner
51                .tokens
52                .compare_exchange_weak(cur, next, Ordering::AcqRel, Ordering::Relaxed)
53            {
54                Ok(_) => return,
55                Err(observed) => cur = observed,
56            }
57        }
58    }
59
60    /// Try to spend one retry token. Returns `false` if the budget is exhausted.
61    pub(crate) fn try_withdraw(&self) -> bool {
62        let inner = &self.0;
63        let mut cur = inner.tokens.load(Ordering::Relaxed);
64        loop {
65            if cur < inner.retry_cost {
66                return false;
67            }
68            match inner.tokens.compare_exchange_weak(
69                cur,
70                cur - inner.retry_cost,
71                Ordering::AcqRel,
72                Ordering::Relaxed,
73            ) {
74                Ok(_) => return true,
75                Err(observed) => cur = observed,
76            }
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn fresh_budget_allows_burst_then_denies() {
87        let b = RetryBudget::new(0.0, 3); // no replenishment, burst 3
88        assert!(b.try_withdraw());
89        assert!(b.try_withdraw());
90        assert!(b.try_withdraw());
91        assert!(!b.try_withdraw(), "burst exhausted with no deposits");
92    }
93
94    #[test]
95    fn deposits_replenish_at_ratio() {
96        let b = RetryBudget::new(0.5, 1); // 50% ratio, small burst
97        // Drain the burst.
98        assert!(b.try_withdraw());
99        assert!(!b.try_withdraw());
100        // Two calls deposit 0.5 + 0.5 = 1.0 token → one more retry allowed.
101        b.deposit();
102        b.deposit();
103        assert!(b.try_withdraw());
104        assert!(!b.try_withdraw());
105    }
106
107    #[test]
108    fn shared_clones_share_one_bucket() {
109        let a = RetryBudget::new(0.0, 1);
110        let b = a.clone();
111        assert!(a.try_withdraw());
112        assert!(!b.try_withdraw(), "clone sees the same drained bucket");
113    }
114
115    #[test]
116    fn deposit_caps_at_max() {
117        let b = RetryBudget::new(1.0, 2); // deposit 1 token/call, max 2
118        for _ in 0..100 {
119            b.deposit();
120        }
121        assert!(b.try_withdraw());
122        assert!(b.try_withdraw());
123        assert!(
124            !b.try_withdraw(),
125            "capped at burst max regardless of deposits"
126        );
127    }
128}
execution_policy/retry/budget.rs

execution_policy/retry/
budget.rs
