ig_client/utils/
rate_limiter.rs

1// Rate limiter for integration tests
2// This module provides utilities to prevent hitting API rate limits
3
4use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::{Duration, Instant};
7use tokio::time::sleep;
8use tracing::info;
9
10/// Rate limiter type for different API endpoints
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum RateLimitType {
13    /// Non-trading requests (per-account): 30 per minute
14    NonTradingAccount,
15    /// Trading requests (per-account): 100 per minute
16    TradingAccount,
17    /// Non-trading requests (per-app): 60 per minute
18    NonTradingApp,
19    /// Historical price data: 10,000 points per week
20    HistoricalPrice,
21}
22
23impl RateLimitType {
24    /// Gets the minimum interval in milliseconds for this rate limit type
25    pub fn min_interval_ms(&self) -> u64 {
26        match self {
27            // 30 requests per minute = 1 request per 2 seconds
28            // Using a more conservative 1 request per 4 seconds to avoid hitting limits
29            Self::NonTradingAccount => 4000,
30            // 100 requests per minute = 1 request per 600ms
31            // Using a more conservative 1 request per 2 seconds
32            Self::TradingAccount => 2000,
33            // 60 requests per minute = 1 request per second
34            // Using a more conservative 1 request per 3 seconds
35            Self::NonTradingApp => 3000,
36            // For historical price data, we'll use a very conservative limit
37            // 10,000 points per week = ~1 request per 60 seconds
38            Self::HistoricalPrice => 120000, // 2 minutes between requests
39        }
40    }
41}
42
43/// Singleton rate limiter for API calls
44pub struct RateLimiter {
45    last_call: AtomicU64,
46    limit_type: RateLimitType,
47}
48
49impl RateLimiter {
50    /// Creates a new rate limiter with the specified rate limit type
51    pub fn new(limit_type: RateLimitType) -> Self {
52        Self {
53            last_call: AtomicU64::new(0),
54            limit_type,
55        }
56    }
57
58    /// Waits if necessary to respect the rate limit
59    pub async fn wait(&self) {
60        let now = Instant::now().elapsed().as_millis() as u64;
61        let last = self.last_call.load(Ordering::Acquire);
62        let min_interval_ms = self.limit_type.min_interval_ms();
63
64        if last > 0 && now - last < min_interval_ms {
65            let wait_time = min_interval_ms - (now - last);
66            info!(
67                "Rate limiter ({:?}): waiting for {}ms",
68                self.limit_type, wait_time
69            );
70            sleep(Duration::from_millis(wait_time)).await;
71        }
72
73        self.last_call.store(
74            Instant::now().elapsed().as_millis() as u64,
75            Ordering::Release,
76        );
77    }
78}
79
80/// Global rate limiter for non-trading account requests (30 per minute)
81pub fn account_non_trading_limiter() -> Arc<RateLimiter> {
82    static INSTANCE: once_cell::sync::Lazy<Arc<RateLimiter>> =
83        once_cell::sync::Lazy::new(|| Arc::new(RateLimiter::new(RateLimitType::NonTradingAccount)));
84
85    INSTANCE.clone()
86}
87
88/// Global rate limiter for trading account requests (100 per minute)
89pub fn account_trading_limiter() -> Arc<RateLimiter> {
90    static INSTANCE: once_cell::sync::Lazy<Arc<RateLimiter>> =
91        once_cell::sync::Lazy::new(|| Arc::new(RateLimiter::new(RateLimitType::TradingAccount)));
92
93    INSTANCE.clone()
94}
95
96/// Global rate limiter for non-trading app requests (60 per minute)
97pub fn app_non_trading_limiter() -> Arc<RateLimiter> {
98    static INSTANCE: once_cell::sync::Lazy<Arc<RateLimiter>> =
99        once_cell::sync::Lazy::new(|| Arc::new(RateLimiter::new(RateLimitType::NonTradingApp)));
100
101    INSTANCE.clone()
102}
103
104/// Global rate limiter for historical price data requests (10,000 points per week)
105pub fn historical_price_limiter() -> Arc<RateLimiter> {
106    static INSTANCE: once_cell::sync::Lazy<Arc<RateLimiter>> =
107        once_cell::sync::Lazy::new(|| Arc::new(RateLimiter::new(RateLimitType::HistoricalPrice)));
108
109    INSTANCE.clone()
110}
111
112/// Default global rate limiter (uses the most conservative limit: non-trading account)
113pub fn global_rate_limiter() -> Arc<RateLimiter> {
114    account_non_trading_limiter()
115}
116
117/// Macro to mark tests that should be run individually to avoid rate limiting
118#[macro_export]
119macro_rules! rate_limited_test {
120    (fn $name:ident() $body:block) => {
121        #[test]
122        #[ignore]
123        fn $name() $body
124    };
125}