ig_client/utils/
rate_limiter.rs1use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::{Duration, Instant};
7use tokio::time::sleep;
8use tracing::info;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum RateLimitType {
13 NonTradingAccount,
15 TradingAccount,
17 NonTradingApp,
19 HistoricalPrice,
21}
22
23impl RateLimitType {
24 pub fn min_interval_ms(&self) -> u64 {
26 match self {
27 Self::NonTradingAccount => 4000,
30 Self::TradingAccount => 2000,
33 Self::NonTradingApp => 3000,
36 Self::HistoricalPrice => 120000, }
40 }
41}
42
43pub struct RateLimiter {
45 last_call: AtomicU64,
46 limit_type: RateLimitType,
47}
48
49impl RateLimiter {
50 pub fn new(limit_type: RateLimitType) -> Self {
52 Self {
53 last_call: AtomicU64::new(0),
54 limit_type,
55 }
56 }
57
58 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
80pub 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
88pub 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
96pub 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
104pub 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
112pub fn global_rate_limiter() -> Arc<RateLimiter> {
114 account_non_trading_limiter()
115}
116
117#[macro_export]
119macro_rules! rate_limited_test {
120 (fn $name:ident() $body:block) => {
121 #[test]
122 #[ignore]
123 fn $name() $body
124 };
125}