forge_gateway/
rotation.rs

1//! Rotation state: which accounts are exhausted, when their cooldowns
2//! end, and how many 429s are on the streak.
3//!
4//! Four triggers mark an account exhausted (spec, Rotation): a
5//! `rate_limit_event` whose status is not `allowed`, a
6//! `anthropic-ratelimit-unified-status` / `-overage-status` of
7//! `rejected` on the failing response, any window at 100% utilization
8//! with a reset ahead, and five consecutive 429s from one account
9//! inside 60 seconds. The first three are proven exhaustion; the
10//! streak is the heuristic, which is why its numbers are config.
11
12use std::collections::HashMap;
13use std::time::{Duration, SystemTime};
14
15use parking_lot::Mutex;
16
17use crate::account::AccountKey;
18
19/// The default streak: this many consecutive 429s from one account
20/// inside [`STREAK_WINDOW`] fire the rotation.
21pub const STREAK_COUNT: u32 = 5;
22
23/// The default window the streak is counted over.
24pub const STREAK_WINDOW: Duration = Duration::from_secs(60);
25
26/// The default cooldown applied when neither the failing response nor
27/// the account's own usage probe reports a reset time.
28pub const NO_RESET_COOLDOWN: Duration = Duration::from_secs(60);
29
30/// The rotation numbers from `[gateway]`: how many 429s fire the
31/// streak, the window they are counted over, and the cooldown when no
32/// reset time is known.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct RotationNumbers {
35    pub streak_count: u32,
36    pub streak_window: Duration,
37    pub no_reset_cooldown: Duration,
38}
39
40impl Default for RotationNumbers {
41    fn default() -> Self {
42        Self {
43            streak_count: STREAK_COUNT,
44            streak_window: STREAK_WINDOW,
45            no_reset_cooldown: NO_RESET_COOLDOWN,
46        }
47    }
48}
49
50/// A unix-seconds reset time, as the wire frames and headers carry.
51pub fn reset_instant(unix_seconds: u64) -> SystemTime {
52    SystemTime::UNIX_EPOCH + Duration::from_secs(unix_seconds)
53}
54
55/// Per-account rotation state: cooldown end, and the streak.
56pub struct RotationState {
57    /// Streak count, window, and no-reset cooldown, from `[gateway]`.
58    streak_count: u32,
59    streak_window: Duration,
60    no_reset_cooldown: Duration,
61    /// Cooldown end per account; an account inside its cooldown is
62    /// skipped by selection and cannot be re-bound.
63    cooldowns: Mutex<HashMap<AccountKey, SystemTime>>,
64    /// Consecutive 429s per account, with the time of the last one.
65    streaks: Mutex<HashMap<AccountKey, (u32, SystemTime)>>,
66}
67
68impl Default for RotationState {
69    fn default() -> Self {
70        Self::with_numbers(RotationNumbers::default())
71    }
72}
73
74impl RotationState {
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    /// The numbers from `[gateway]`: how many 429s fire the streak,
80    /// the window they are counted over, and the cooldown when no
81    /// reset time is known.
82    pub fn with_numbers(numbers: RotationNumbers) -> Self {
83        Self {
84            streak_count: numbers.streak_count,
85            streak_window: numbers.streak_window,
86            no_reset_cooldown: numbers.no_reset_cooldown,
87            cooldowns: Mutex::new(HashMap::new()),
88            streaks: Mutex::new(HashMap::new()),
89        }
90    }
91
92    /// The cooldown applied when no reset time is known.
93    pub fn no_reset_cooldown(&self) -> Duration {
94        self.no_reset_cooldown
95    }
96
97    /// Mark an account exhausted until `until`.
98    pub fn cool_down(&self, key: &AccountKey, until: SystemTime) {
99        self.cooldowns.lock().insert(key.clone(), until);
100    }
101
102    /// `true` while the account is inside a cooldown.
103    pub fn is_cooling_down(&self, key: &AccountKey, now: SystemTime) -> bool {
104        self.cooldowns.lock().get(key).is_some_and(|until| *until > now)
105    }
106
107    /// When `key`'s cooldown ends, `None` when it never cooled.
108    pub fn cooldown_end(&self, key: &AccountKey) -> Option<SystemTime> {
109        self.cooldowns.lock().get(key).copied()
110    }
111
112    /// The soonest cooldown still live at `now`, across every account. A
113    /// reader sleeps until it to announce the thaw the clock alone decides.
114    pub fn soonest_cooldown_end(&self, now: SystemTime) -> Option<SystemTime> {
115        self.cooldowns.lock().values().copied().filter(|until| *until > now).min()
116    }
117
118    /// The soonest a cooling account among `keys` recovers, for the
119    /// failure message. `None` when none of them is cooling.
120    pub fn soonest_reset_among(&self, keys: &[AccountKey], now: SystemTime) -> Option<SystemTime> {
121        let cooldowns = self.cooldowns.lock();
122        keys.iter().filter_map(|key| cooldowns.get(key).copied()).filter(|until| *until > now).min()
123    }
124
125    /// Record a 429 from `account` at `now`; rotates when the streak
126    /// reaches the configured count inside the configured window.
127    /// `true` when the streak fired.
128    ///
129    /// The window is a max-gap since the last 429, not a sliding
130    /// count: a gap beyond the window restarts at 1, so a slow drip
131    /// still accumulates. That over-rotates a drip marginally rather
132    /// than under-rotating a burst - the conservative direction.
133    pub fn record_429(&self, key: &AccountKey, now: SystemTime) -> bool {
134        let mut streaks = self.streaks.lock();
135        let entry = streaks.entry(key.clone()).or_insert((0, now));
136        if now.duration_since(entry.1).unwrap_or(Duration::ZERO) > self.streak_window {
137            *entry = (1, now);
138        } else {
139            entry.0 += 1;
140            entry.1 = now;
141        }
142        entry.0 >= self.streak_count
143    }
144
145    /// Clear the streak: any non-429 success resets it.
146    pub fn reset_streak(&self, key: &AccountKey) {
147        self.streaks.lock().remove(key);
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn the_streak_fires_on_the_fifth_429_inside_the_window() {
157        let state = RotationState::new();
158        let key = AccountKey("A".to_owned());
159        let mut now = SystemTime::now();
160        for _ in 0..4 {
161            assert!(!state.record_429(&key, now), "the streak must not fire early");
162            now += Duration::from_secs(10);
163        }
164        assert!(state.record_429(&key, now), "the fifth 429 inside the window fires");
165    }
166
167    #[test]
168    fn a_gap_outside_the_window_resets_the_streak() {
169        let state = RotationState::new();
170        let key = AccountKey("A".to_owned());
171        let mut now = SystemTime::now();
172        for _ in 0..4 {
173            state.record_429(&key, now);
174            now += Duration::from_secs(10);
175        }
176        now += Duration::from_secs(61);
177        assert!(!state.record_429(&key, now), "a 429 outside the window restarts the count");
178    }
179
180    #[test]
181    fn cooldowns_expire() {
182        let state = RotationState::new();
183        let key = AccountKey("A".to_owned());
184        let now = SystemTime::now();
185        state.cool_down(&key, now + Duration::from_secs(60));
186        assert!(state.is_cooling_down(&key, now));
187        assert!(!state.is_cooling_down(&key, now + Duration::from_secs(61)));
188    }
189
190    #[test]
191    fn the_soonest_reset_among_is_the_earliest_live_cooldown_asked_about() {
192        let state = RotationState::new();
193        let now = SystemTime::now();
194        state.cool_down(&AccountKey("A".to_owned()), now + Duration::from_secs(90));
195        state.cool_down(&AccountKey("B".to_owned()), now + Duration::from_secs(30));
196        state.cool_down(&AccountKey("C".to_owned()), now - Duration::from_secs(1));
197        let asked = [AccountKey("A".to_owned()), AccountKey("C".to_owned())];
198        assert_eq!(
199            state.soonest_reset_among(&asked, now),
200            Some(now + Duration::from_secs(90)),
201            "an account outside the ask and an expired cooldown are both ignored"
202        );
203        assert_eq!(state.soonest_reset_among(&[], now), None);
204        assert_eq!(
205            state
206                .soonest_reset_among(&[AccountKey("B".to_owned())], now + Duration::from_secs(120)),
207            None
208        );
209    }
210
211    /// **The end a reader sleeps to is a live one.** An expired cooldown is
212    /// not a wake, so the filter drops it; the soonest of the rest is the
213    /// answer.
214    #[test]
215    fn the_soonest_cooldown_end_ignores_expired_cooldowns() {
216        let state = RotationState::new();
217        let now = SystemTime::now();
218        state.cool_down(&AccountKey("Expired".to_owned()), now - Duration::from_secs(1));
219        assert_eq!(state.soonest_cooldown_end(now), None, "an expired cooldown is no wake");
220
221        state.cool_down(&AccountKey("Later".to_owned()), now + Duration::from_secs(90));
222        state.cool_down(&AccountKey("Sooner".to_owned()), now + Duration::from_secs(30));
223        assert_eq!(
224            state.soonest_cooldown_end(now),
225            Some(now + Duration::from_secs(30)),
226            "the soonest live cooldown is what a reader sleeps to",
227        );
228    }
229
230    #[test]
231    fn a_unix_reset_time_maps_from_the_epoch() {
232        let reset = reset_instant(1_000_000);
233        assert_eq!(
234            reset.duration_since(SystemTime::UNIX_EPOCH).unwrap(),
235            Duration::from_secs(1_000_000)
236        );
237    }
238
239    #[test]
240    fn configured_numbers_replace_the_defaults() {
241        let state = RotationState::with_numbers(RotationNumbers {
242            streak_count: 2,
243            streak_window: Duration::from_secs(30),
244            no_reset_cooldown: Duration::from_secs(120),
245        });
246        let key = AccountKey("A".to_owned());
247        let mut now = SystemTime::now();
248        assert!(!state.record_429(&key, now), "the configured count must not fire early");
249        now += Duration::from_secs(10);
250        assert!(
251            state.record_429(&key, now),
252            "the second 429 inside the configured window fires at count 2"
253        );
254        assert_eq!(
255            state.no_reset_cooldown(),
256            Duration::from_secs(120),
257            "the configured no-reset cooldown is what the failure paths apply"
258        );
259        // The configured window is the one the max-gap reads against.
260        let key = AccountKey("B".to_owned());
261        let mut now = SystemTime::now();
262        state.record_429(&key, now);
263        now += Duration::from_secs(31);
264        assert!(
265            !state.record_429(&key, now),
266            "a 429 beyond the configured 30s window restarts the count"
267        );
268    }
269}
forge_gateway/rotation.rs

forge_gateway/
rotation.rs
