Skip to main content

projectx_client/
rate_limit.rs

1// SPDX-FileCopyrightText: 2026 Kevin Monaghan
2// SPDX-License-Identifier: MIT-0
3
4//! Provider REST rate-limit configuration and shared admission control.
5
6use std::{collections::VecDeque, fmt, time::Duration};
7
8use parking_lot::Mutex;
9use tokio::time::{Instant, sleep};
10
11use crate::Error;
12
13const DEFAULT_HISTORY_REQUESTS: usize = 50;
14const DEFAULT_HISTORY_WINDOW: Duration = Duration::from_secs(30);
15const DEFAULT_GENERAL_REQUESTS: usize = 200;
16const DEFAULT_GENERAL_WINDOW: Duration = Duration::from_mins(1);
17const MAX_INITIAL_CAPACITY: usize = 1_024;
18
19/// Identifies one provider REST rate-limit budget.
20#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
21#[non_exhaustive]
22pub enum RateLimitKind {
23    /// The dedicated `/api/History/retrieveBars` budget.
24    History,
25    /// The budget shared by all other authenticated REST endpoints.
26    General,
27}
28
29impl fmt::Display for RateLimitKind {
30    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
31        match self {
32            Self::History => formatter.write_str("history"),
33            Self::General => formatter.write_str("general"),
34        }
35    }
36}
37
38/// One rolling-window request limit.
39#[derive(Clone, Copy, Debug, Eq, PartialEq)]
40#[non_exhaustive]
41pub struct RateLimit {
42    max_requests: usize,
43    window: Duration,
44}
45
46impl RateLimit {
47    /// Creates a non-zero rolling-window request limit.
48    ///
49    /// # Errors
50    ///
51    /// Returns [`Error::Configuration`] when `max_requests` or `window` is
52    /// zero, or when the duration cannot be represented by a monotonic clock.
53    pub fn new(max_requests: usize, window: Duration) -> Result<Self, Error> {
54        if max_requests == 0
55            || window.is_zero()
56            || std::time::Instant::now().checked_add(window).is_none()
57        {
58            return Err(Error::Configuration(
59                "rate limits require a positive request count and representable window".to_owned(),
60            ));
61        }
62        Ok(Self {
63            max_requests,
64            window,
65        })
66    }
67
68    const fn trusted(max_requests: usize, window: Duration) -> Self {
69        Self {
70            max_requests,
71            window,
72        }
73    }
74
75    /// Returns the maximum admitted requests in the rolling window.
76    #[must_use]
77    pub const fn max_requests(self) -> usize {
78        self.max_requests
79    }
80
81    /// Returns the rolling-window duration.
82    #[must_use]
83    pub const fn window(self) -> Duration {
84        self.window
85    }
86}
87
88/// Rate limits applied to authenticated `ProjectX` REST requests.
89///
90/// [`Default`] uses the provider's documented limits: 50 history requests per
91/// 30 seconds and 200 other authenticated requests per 60 seconds.
92#[derive(Clone, Copy, Debug, Eq, PartialEq)]
93#[non_exhaustive]
94pub struct RateLimitConfig {
95    history: RateLimit,
96    general: RateLimit,
97}
98
99impl RateLimitConfig {
100    /// Creates configuration from dedicated history and general limits.
101    #[must_use]
102    pub const fn new(history: RateLimit, general: RateLimit) -> Self {
103        Self { history, general }
104    }
105
106    /// Returns the history endpoint limit.
107    #[must_use]
108    pub const fn history(self) -> RateLimit {
109        self.history
110    }
111
112    /// Returns the limit for all other authenticated endpoints.
113    #[must_use]
114    pub const fn general(self) -> RateLimit {
115        self.general
116    }
117
118    pub(crate) const fn limit(self, kind: RateLimitKind) -> RateLimit {
119        match kind {
120            RateLimitKind::History => self.history,
121            RateLimitKind::General => self.general,
122        }
123    }
124}
125
126impl Default for RateLimitConfig {
127    fn default() -> Self {
128        Self {
129            history: RateLimit::trusted(DEFAULT_HISTORY_REQUESTS, DEFAULT_HISTORY_WINDOW),
130            general: RateLimit::trusted(DEFAULT_GENERAL_REQUESTS, DEFAULT_GENERAL_WINDOW),
131        }
132    }
133}
134
135pub(crate) struct RateLimits {
136    config: Option<RateLimitConfig>,
137    history: Option<WindowLimiter>,
138    general: Option<WindowLimiter>,
139}
140
141impl RateLimits {
142    pub(crate) fn new(config: Option<RateLimitConfig>) -> Self {
143        let history = config.map(|limits| WindowLimiter::new(limits.history));
144        let general = config.map(|limits| WindowLimiter::new(limits.general));
145        Self {
146            config,
147            history,
148            general,
149        }
150    }
151
152    pub(crate) const fn config(&self) -> Option<RateLimitConfig> {
153        self.config
154    }
155
156    pub(crate) fn limit(&self, kind: RateLimitKind) -> RateLimit {
157        self.config.unwrap_or_default().limit(kind)
158    }
159
160    pub(crate) async fn wait(&self, kind: RateLimitKind) {
161        let Some(limiter) = self.limiter(kind) else {
162            return;
163        };
164        loop {
165            match limiter.try_acquire(Instant::now()) {
166                Ok(()) => return,
167                Err(retry_after) => sleep(retry_after).await,
168            }
169        }
170    }
171
172    pub(crate) fn try_acquire(&self, kind: RateLimitKind) -> Result<(), Duration> {
173        self.limiter(kind)
174            .map_or(Ok(()), |limiter| limiter.try_acquire(Instant::now()))
175    }
176
177    pub(crate) fn cool_down(&self, kind: RateLimitKind, retry_after: Duration) {
178        if let Some(limiter) = self.limiter(kind) {
179            limiter.cool_down(Instant::now(), retry_after);
180        }
181    }
182
183    fn limiter(&self, kind: RateLimitKind) -> Option<&WindowLimiter> {
184        match kind {
185            RateLimitKind::History => self.history.as_ref(),
186            RateLimitKind::General => self.general.as_ref(),
187        }
188    }
189}
190
191struct WindowLimiter {
192    limit: RateLimit,
193    state: Mutex<WindowState>,
194}
195
196impl WindowLimiter {
197    fn new(limit: RateLimit) -> Self {
198        Self {
199            limit,
200            state: Mutex::new(WindowState::with_capacity(
201                limit.max_requests.min(MAX_INITIAL_CAPACITY),
202            )),
203        }
204    }
205
206    fn try_acquire(&self, now: Instant) -> Result<(), Duration> {
207        let mut state = self.state.lock();
208        state.expire(now, self.limit.window);
209        let retry_after = state.retry_after(now, self.limit);
210        if retry_after.is_zero() {
211            state.admitted.push_back(now);
212            Ok(())
213        } else {
214            Err(retry_after)
215        }
216    }
217
218    fn cool_down(&self, now: Instant, retry_after: Duration) {
219        let Some(deadline) = now.checked_add(retry_after) else {
220            return;
221        };
222        let mut state = self.state.lock();
223        if state
224            .cooldown_until
225            .is_none_or(|current| deadline > current)
226        {
227            state.cooldown_until = Some(deadline);
228        }
229    }
230}
231
232struct WindowState {
233    admitted: VecDeque<Instant>,
234    cooldown_until: Option<Instant>,
235}
236
237impl WindowState {
238    fn with_capacity(capacity: usize) -> Self {
239        Self {
240            admitted: VecDeque::with_capacity(capacity),
241            cooldown_until: None,
242        }
243    }
244
245    fn expire(&mut self, now: Instant, window: Duration) {
246        while self
247            .admitted
248            .front()
249            .is_some_and(|admitted| now.saturating_duration_since(*admitted) >= window)
250        {
251            self.admitted.pop_front();
252        }
253        if self.cooldown_until.is_some_and(|deadline| deadline <= now) {
254            self.cooldown_until = None;
255        }
256    }
257
258    fn retry_after(&self, now: Instant, limit: RateLimit) -> Duration {
259        let cooldown = self.cooldown_until.map_or(Duration::ZERO, |deadline| {
260            deadline.saturating_duration_since(now)
261        });
262        let capacity = if self.admitted.len() < limit.max_requests {
263            Duration::ZERO
264        } else {
265            self.admitted.front().map_or(limit.window, |oldest| {
266                oldest
267                    .checked_add(limit.window)
268                    .map_or(limit.window, |deadline| {
269                        deadline.saturating_duration_since(now)
270                    })
271            })
272        };
273        cooldown.max(capacity)
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use std::sync::Arc;
280
281    use tokio::time::{Instant, advance};
282
283    use super::*;
284
285    fn limits(requests: usize, window: Duration) -> Arc<RateLimits> {
286        let limit = RateLimit::new(requests, window)
287            .unwrap_or_else(|error| panic!("test rate limit must be valid: {error}"));
288        Arc::new(RateLimits::new(Some(RateLimitConfig::new(limit, limit))))
289    }
290
291    #[tokio::test(start_paused = true)]
292    async fn wait_resumes_when_the_rolling_window_releases_capacity() {
293        let window = Duration::from_secs(10);
294        let limits = limits(1, window);
295        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
296
297        let started = Instant::now();
298        let waiting_limits = Arc::clone(&limits);
299        let waiter = tokio::spawn(async move {
300            waiting_limits.wait(RateLimitKind::General).await;
301            Instant::now()
302        });
303        tokio::task::yield_now().await;
304        assert!(!waiter.is_finished());
305
306        advance(window).await;
307        let admitted = waiter
308            .await
309            .unwrap_or_else(|error| panic!("waiting task must finish: {error}"));
310        assert_eq!(admitted.saturating_duration_since(started), window);
311    }
312
313    #[tokio::test(start_paused = true)]
314    async fn provider_cooldown_blocks_admission_until_its_deadline() {
315        let cooldown = Duration::from_secs(15);
316        let limits = limits(2, Duration::from_secs(30));
317        limits.cool_down(RateLimitKind::General, cooldown);
318
319        assert_eq!(limits.try_acquire(RateLimitKind::General), Err(cooldown));
320        advance(cooldown).await;
321        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
322    }
323
324    #[tokio::test(start_paused = true)]
325    async fn rolling_window_expires_staggered_attempts_individually() {
326        let window = Duration::from_secs(10);
327        let limits = limits(2, window);
328        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
329        advance(Duration::from_secs(5)).await;
330        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
331
332        advance(Duration::from_secs(5)).await;
333        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
334        assert_eq!(
335            limits.try_acquire(RateLimitKind::General),
336            Err(Duration::from_secs(5))
337        );
338    }
339
340    #[tokio::test(start_paused = true)]
341    async fn cancelled_waiter_leaves_no_ghost_reservation() {
342        let window = Duration::from_secs(10);
343        let limits = limits(1, window);
344        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
345        let waiting_limits = Arc::clone(&limits);
346        let waiter = tokio::spawn(async move {
347            waiting_limits.wait(RateLimitKind::General).await;
348        });
349        tokio::task::yield_now().await;
350        waiter.abort();
351
352        advance(window).await;
353        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
354    }
355
356    #[tokio::test(start_paused = true)]
357    async fn default_history_budget_reaches_its_limit_before_throttling() {
358        let limits = Arc::new(RateLimits::new(Some(RateLimitConfig::default())));
359        for request_number in 1..=50 {
360            assert_eq!(
361                limits.try_acquire(RateLimitKind::History),
362                Ok(()),
363                "history request {request_number} must be admitted"
364            );
365        }
366        assert_eq!(
367            limits.try_acquire(RateLimitKind::History),
368            Err(Duration::from_secs(30))
369        );
370    }
371
372    #[tokio::test(start_paused = true)]
373    async fn default_general_budget_reaches_its_limit_before_throttling() {
374        let limits = Arc::new(RateLimits::new(Some(RateLimitConfig::default())));
375        for request_number in 1..=200 {
376            assert_eq!(
377                limits.try_acquire(RateLimitKind::General),
378                Ok(()),
379                "general request {request_number} must be admitted"
380            );
381        }
382        assert_eq!(
383            limits.try_acquire(RateLimitKind::General),
384            Err(Duration::from_mins(1))
385        );
386    }
387
388    #[tokio::test(start_paused = true)]
389    async fn later_provider_cooldown_cannot_be_shortened() {
390        let limits = limits(200, Duration::from_mins(1));
391        limits.cool_down(RateLimitKind::General, Duration::from_secs(30));
392        limits.cool_down(RateLimitKind::General, Duration::from_secs(5));
393        assert_eq!(
394            limits.try_acquire(RateLimitKind::General),
395            Err(Duration::from_secs(30))
396        );
397    }
398
399    #[tokio::test(start_paused = true)]
400    async fn admission_uses_the_later_of_window_and_provider_cooldown() {
401        let limits = limits(1, Duration::from_secs(20));
402        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
403        limits.cool_down(RateLimitKind::General, Duration::from_secs(5));
404        assert_eq!(
405            limits.try_acquire(RateLimitKind::General),
406            Err(Duration::from_secs(20))
407        );
408    }
409
410    #[test]
411    fn history_and_general_budgets_are_independent() {
412        let limits = limits(1, Duration::from_secs(30));
413        assert_eq!(limits.try_acquire(RateLimitKind::General), Ok(()));
414        assert_eq!(limits.try_acquire(RateLimitKind::History), Ok(()));
415        assert!(limits.try_acquire(RateLimitKind::General).is_err());
416        assert!(limits.try_acquire(RateLimitKind::History).is_err());
417    }
418
419    #[test]
420    fn public_defaults_match_the_provider_contract() {
421        let config = RateLimitConfig::default();
422        assert_eq!(config.history().max_requests(), 50);
423        assert_eq!(config.history().window(), Duration::from_secs(30));
424        assert_eq!(config.general().max_requests(), 200);
425        assert_eq!(config.general().window(), Duration::from_mins(1));
426    }
427
428    #[test]
429    fn public_rate_limit_rejects_zero_values() {
430        assert!(RateLimit::new(0, Duration::from_secs(1)).is_err());
431        assert!(RateLimit::new(1, Duration::ZERO).is_err());
432    }
433}