typesafe_jev/
limiter.rs

1use std::fmt;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::{Condvar, Mutex, MutexGuard, PoisonError};
4use std::time::{Duration, Instant};
5
6/// Uniform in [0, 1). Only used for retry and pause jitter, so a tiny xorshift is plenty.
7pub(crate) fn jitter() -> f64 {
8    static STATE: AtomicU64 = AtomicU64::new(0);
9    let mut x = STATE.load(Ordering::Relaxed);
10    if x == 0 {
11        let nanos = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(1, |d| d.as_nanos() as u64);
12        x = nanos ^ (u64::from(std::process::id()) << 32) | 1;
13    }
14    x ^= x << 13;
15    x ^= x >> 7;
16    x ^= x << 17;
17    STATE.store(x, Ordering::Relaxed);
18    (x >> 11) as f64 / (1u64 << 53) as f64
19}
20
21struct Gate {
22    limit: f64,
23    in_flight: usize,
24    paused_until: Option<Instant>,
25}
26
27/// Shared concurrency gate with additive-increase, multiplicative-decrease backoff.
28///
29/// TypeSafe's rate limit is a sustained token budget with no quota headers, and several
30/// processes may share one key. So on a throttled reply every thread pauses briefly and the
31/// allowed concurrency halves; each success grows it back by roughly one slot per round.
32///
33/// A [`Client`](crate::Client) owns one and wraps every attempt in [`acquire`](Self::acquire)
34/// and [`release`](Self::release). The type is public so that the gate can be inspected, and so
35/// that other transports can share the policy.
36pub struct AdaptiveLimiter {
37    max: f64,
38    pause: Duration,
39    gate: Mutex<Gate>,
40    cv: Condvar,
41}
42
43impl AdaptiveLimiter {
44    /// `max_concurrency` is the ceiling the limit grows back to (at least one). `pause` is how
45    /// long every thread waits after a throttled reply, scaled by a random factor in [1, 2).
46    pub fn new(max_concurrency: usize, pause: Duration) -> Self {
47        let max = max_concurrency.max(1) as f64;
48        AdaptiveLimiter { max, pause, gate: Mutex::new(Gate { limit: max, in_flight: 0, paused_until: None }), cv: Condvar::new() }
49    }
50
51    /// The gate is three plain numbers that every critical section leaves consistent, so a lock
52    /// poisoned by a panicking thread is still good to use.
53    fn gate(&self) -> MutexGuard<'_, Gate> {
54        self.gate.lock().unwrap_or_else(PoisonError::into_inner)
55    }
56
57    /// The current concurrency limit. Starts at the ceiling, halves on throttling, grows back on
58    /// success; never below one. Fractional, because it grows by a fraction of a slot per reply;
59    /// the whole part is the number of attempts allowed at once.
60    pub fn limit(&self) -> f64 {
61        self.gate().limit
62    }
63
64    /// Attempts between `acquire` and `release` right now.
65    pub fn in_flight(&self) -> usize {
66        self.gate().in_flight
67    }
68
69    /// Blocks until there is a free slot under the current limit and no pause is in effect.
70    /// Every `acquire` must be matched by one [`release`](Self::release).
71    pub fn acquire(&self) {
72        let mut gate = self.gate();
73        loop {
74            let wait = gate.paused_until.and_then(|t| t.checked_duration_since(Instant::now())).filter(|d| !d.is_zero());
75            match wait {
76                None if gate.in_flight < (gate.limit as usize).max(1) => {
77                    gate.in_flight += 1;
78                    return;
79                }
80                None => gate = self.cv.wait(gate).unwrap_or_else(PoisonError::into_inner),
81                Some(d) => {
82                    gate = self.cv.wait_timeout(gate, d.max(Duration::from_millis(50))).unwrap_or_else(PoisonError::into_inner).0;
83                }
84            }
85        }
86    }
87
88    /// Frees the slot. `throttled` reports a rate-limited reply (HTTP 429 or 529): the limit
89    /// halves and every thread pauses, once per throttling episode rather than once per reply.
90    /// Any other outcome grows the limit by about one slot per round of requests. A `release`
91    /// without an `acquire` frees nothing.
92    pub fn release(&self, throttled: bool) {
93        let mut gate = self.gate();
94        gate.in_flight = gate.in_flight.saturating_sub(1);
95        let now = Instant::now();
96        if throttled {
97            if gate.paused_until.is_none_or(|t| now >= t) {
98                gate.limit = (gate.limit / 2.0).max(1.0);
99                // A pause too long to represent is no pause, not a panic with the lock held.
100                let pause = Duration::try_from_secs_f64(self.pause.as_secs_f64() * (1.0 + jitter())).ok();
101                gate.paused_until = pause.and_then(|pause| now.checked_add(pause));
102            }
103        } else {
104            gate.limit = self.max.min(gate.limit + 1.0 / gate.limit.max(1.0));
105        }
106        self.cv.notify_all();
107    }
108}
109
110impl fmt::Debug for AdaptiveLimiter {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        let gate = self.gate();
113        f.debug_struct("AdaptiveLimiter")
114            .field("max", &self.max)
115            .field("pause", &self.pause)
116            .field("limit", &gate.limit)
117            .field("in_flight", &gate.in_flight)
118            .finish_non_exhaustive()
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use std::sync::atomic::AtomicUsize;
126    use std::sync::Arc;
127
128    #[test]
129    fn jitter_is_in_the_unit_interval_and_varies() {
130        let draws: Vec<f64> = (0..64).map(|_| jitter()).collect();
131        assert!(draws.iter().all(|&x| (0.0..1.0).contains(&x)));
132        assert!(draws.windows(2).any(|w| w[0] != w[1]));
133    }
134
135    #[test]
136    fn throttling_halves_the_limit_once_per_episode_and_success_grows_it_back() {
137        let limiter = AdaptiveLimiter::new(8, Duration::ZERO);
138        assert_eq!(limiter.limit(), 8.0);
139        for _ in 0..3 {
140            limiter.acquire();
141        }
142        assert_eq!(limiter.in_flight(), 3);
143        limiter.release(true);
144        limiter.release(true); // same episode: no second cut while the pause is not over
145        assert!(limiter.limit() <= 4.0 && limiter.limit() >= 2.0, "{}", limiter.limit());
146        limiter.release(false);
147        assert_eq!(limiter.in_flight(), 0);
148        for _ in 0..200 {
149            limiter.acquire();
150            limiter.release(false);
151        }
152        assert_eq!(limiter.limit(), 8.0);
153    }
154
155    #[test]
156    fn the_limit_never_drops_below_one() {
157        let limiter = AdaptiveLimiter::new(0, Duration::ZERO);
158        assert_eq!(limiter.limit(), 1.0);
159        for _ in 0..5 {
160            limiter.acquire();
161            std::thread::sleep(Duration::from_millis(1));
162            limiter.release(true);
163        }
164        assert_eq!(limiter.limit(), 1.0);
165    }
166
167    #[test]
168    fn an_unmatched_release_and_an_unrepresentable_pause_do_not_panic() {
169        let limiter = AdaptiveLimiter::new(2, Duration::MAX);
170        limiter.release(false);
171        assert_eq!(limiter.in_flight(), 0);
172        limiter.acquire();
173        limiter.release(true);
174        assert_eq!((limiter.limit(), limiter.in_flight()), (1.0, 0));
175        limiter.acquire(); // no pause took effect, so this returns
176        limiter.release(false);
177        let shown = format!("{limiter:?}");
178        assert!(shown.starts_with("AdaptiveLimiter {") && shown.contains("in_flight: 0"), "{shown}");
179    }
180
181    #[test]
182    fn acquire_blocks_while_the_gate_is_full() {
183        let limiter = Arc::new(AdaptiveLimiter::new(2, Duration::ZERO));
184        let peak = Arc::new(AtomicUsize::new(0));
185        let running = Arc::new(AtomicUsize::new(0));
186        let workers: Vec<_> = (0..8)
187            .map(|_| {
188                let (limiter, peak, running) = (Arc::clone(&limiter), Arc::clone(&peak), Arc::clone(&running));
189                std::thread::spawn(move || {
190                    limiter.acquire();
191                    let now = running.fetch_add(1, Ordering::SeqCst) + 1;
192                    peak.fetch_max(now, Ordering::SeqCst);
193                    std::thread::sleep(Duration::from_millis(5));
194                    running.fetch_sub(1, Ordering::SeqCst);
195                    limiter.release(false);
196                })
197            })
198            .collect();
199        for worker in workers {
200            worker.join().unwrap();
201        }
202        assert!(peak.load(Ordering::SeqCst) <= 2, "{}", peak.load(Ordering::SeqCst));
203        assert_eq!(limiter.in_flight(), 0);
204    }
205}
typesafe_jev/limiter.rs

typesafe_jev/
limiter.rs
