wickra_exchange_core/
ratelimiter.rs

1//! Weight-based rate limiting.
2//!
3//! Venues meter requests by **weight** per time window, not a flat count: a
4//! `GET /depth?limit=1000` costs far more than a ticker. [`WeightedRateLimiter`]
5//! is a windowed weight budget that also honours an explicit cool-off when the
6//! venue returns `429`/`418` with a `Retry-After`. Like the rest of the crate it
7//! takes `now_ms` as an argument, so it is deterministic under test.
8
9/// The result of asking the limiter for budget.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Acquire {
12    /// The request fits within budget; the weight has been charged.
13    Allowed,
14    /// The request must wait; `retry_after_ms` is the advised delay.
15    Throttled {
16        /// Milliseconds to wait before retrying.
17        retry_after_ms: i64,
18    },
19}
20
21/// A windowed, weight-based rate limiter with venue cool-off support.
22#[derive(Debug)]
23pub struct WeightedRateLimiter {
24    capacity: u32,
25    window_ms: i64,
26    used: u32,
27    window_start_ms: Option<i64>,
28    banned_until_ms: Option<i64>,
29}
30
31impl WeightedRateLimiter {
32    /// A limiter allowing `capacity` units of weight per `window_ms`.
33    #[must_use]
34    pub fn new(capacity: u32, window_ms: i64) -> Self {
35        Self {
36            capacity,
37            window_ms,
38            used: 0,
39            window_start_ms: None,
40            banned_until_ms: None,
41        }
42    }
43
44    /// The weight charged in the current window.
45    #[must_use]
46    pub fn used(&self) -> u32 {
47        self.used
48    }
49
50    /// Record that the venue rate-limited us (`429`/`418`): refuse all requests
51    /// until `now_ms + retry_after_ms`.
52    pub fn note_rate_limited(&mut self, retry_after_ms: i64, now_ms: i64) {
53        self.banned_until_ms = Some(now_ms + retry_after_ms);
54    }
55
56    /// Try to charge `weight` against the budget at `now_ms`.
57    pub fn try_acquire(&mut self, weight: u32, now_ms: i64) -> Acquire {
58        // Honour an explicit venue cool-off first.
59        if let Some(until) = self.banned_until_ms {
60            if now_ms < until {
61                return Acquire::Throttled {
62                    retry_after_ms: until - now_ms,
63                };
64            }
65            self.banned_until_ms = None;
66        }
67
68        // Roll the window if it has elapsed (or start it on first use).
69        let start = match self.window_start_ms {
70            Some(start) if now_ms < start + self.window_ms => start,
71            _ => {
72                self.window_start_ms = Some(now_ms);
73                self.used = 0;
74                now_ms
75            }
76        };
77
78        if self.used + weight > self.capacity {
79            return Acquire::Throttled {
80                retry_after_ms: start + self.window_ms - now_ms,
81            };
82        }
83        self.used += weight;
84        Acquire::Allowed
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn charges_weight_until_the_budget_is_spent() {
94        let mut rl = WeightedRateLimiter::new(100, 60_000);
95        assert_eq!(rl.try_acquire(40, 1_000), Acquire::Allowed);
96        assert_eq!(rl.used(), 40);
97        assert_eq!(rl.try_acquire(40, 1_500), Acquire::Allowed);
98        assert_eq!(rl.used(), 80);
99        // 80 + 30 > 100 -> throttled, advised wait is to the window end.
100        assert_eq!(
101            rl.try_acquire(30, 2_000),
102            Acquire::Throttled {
103                retry_after_ms: 1_000 + 60_000 - 2_000
104            }
105        );
106    }
107
108    #[test]
109    fn window_rolls_and_resets_budget() {
110        let mut rl = WeightedRateLimiter::new(100, 60_000);
111        assert_eq!(rl.try_acquire(100, 0), Acquire::Allowed);
112        // Still inside the window: full.
113        assert!(matches!(
114            rl.try_acquire(1, 30_000),
115            Acquire::Throttled { .. }
116        ));
117        // After the window: budget resets.
118        assert_eq!(rl.try_acquire(50, 60_000), Acquire::Allowed);
119        assert_eq!(rl.used(), 50);
120    }
121
122    #[test]
123    fn venue_cool_off_blocks_until_it_expires() {
124        let mut rl = WeightedRateLimiter::new(100, 60_000);
125        rl.note_rate_limited(5_000, 1_000);
126        assert_eq!(
127            rl.try_acquire(1, 2_000),
128            Acquire::Throttled {
129                retry_after_ms: 4_000
130            }
131        );
132        // After the cool-off, requests flow again.
133        assert_eq!(rl.try_acquire(1, 6_000), Acquire::Allowed);
134    }
135}
wickra_exchange_core/ratelimiter.rs

wickra_exchange_core/
ratelimiter.rs
