xapi_okx/common/
ratelimiter.rs

1use crate::data::instrument_type::OkxInstrumentType;
2use governor::{DefaultDirectRateLimiter, InsufficientCapacity, Quota};
3use serde::{Deserialize, Serialize};
4use std::{collections::HashMap, num::NonZero, time::Duration};
5use xapi_shared::ratelimiter::SharedRatelimiterTrait;
6
7pub struct OkxRatelimiter {
8    limiters: HashMap<OkxRateLimitKey, DefaultDirectRateLimiter>,
9}
10
11#[async_trait::async_trait]
12impl SharedRatelimiterTrait<OkxRateLimitKey> for OkxRatelimiter {
13    async fn limit_on(
14        &self,
15        key: &OkxRateLimitKey,
16        value: NonZero<u32>,
17    ) -> Result<(), InsufficientCapacity> {
18        if let Some(limiter) = self.limiters.get(key) {
19            limiter
20                .until_n_ready(value)
21                .await
22                .inspect_err(|err| tracing::error!("okx rate limiter error for {:?}: {err}", key))?
23        }
24
25        Ok(())
26    }
27}
28
29impl Default for OkxRatelimiter {
30    fn default() -> Self {
31        let rules = vec![
32            /* public */
33            OkxRateLimitRule {
34                key: OkxRateLimitKey::new_ip_based("/api/v5/public/instruments"),
35                requests_per_2_seconds: 20,
36            },
37            OkxRateLimitRule {
38                key: OkxRateLimitKey::new_ip_based("/api/v5/public/time"),
39                requests_per_2_seconds: 10,
40            },
41            /* trading */
42            OkxRateLimitRule {
43                key: OkxRateLimitKey::new_user_based("/api/v5/account/balance"),
44                requests_per_2_seconds: 10,
45            },
46        ];
47
48        Self::new(&rules)
49    }
50}
51
52impl OkxRatelimiter {
53    pub fn new(rules: &[OkxRateLimitRule]) -> Self {
54        let mut limiters = HashMap::new();
55
56        for rule in rules {
57            if rule.requests_per_2_seconds == 0 {
58                tracing::warn!("skipping rate limit rule with zero limit: {:?}", rule);
59                continue;
60            }
61
62            let limit = NonZero::new(rule.requests_per_2_seconds).unwrap();
63
64            // Create quota for 2-second intervals
65            let quota = Quota::with_period(Duration::from_secs(2))
66                .unwrap()
67                .allow_burst(limit);
68
69            let limiter = DefaultDirectRateLimiter::direct(quota);
70
71            limiters.insert(rule.key.clone(), limiter);
72        }
73
74        Self { limiters }
75    }
76}
77
78#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
79pub enum OkxRateLimitScope {
80    /// Public endpoints limited by IP address
81    IpBased,
82    /// Private endpoints limited by User ID
83    UserBased,
84    /// Private endpoints limited by User ID + Instrument Type
85    UserAndInstrumentType(OkxInstrumentType),
86    /// WebSocket connections limited by connection
87    Connection,
88}
89
90#[derive(Debug, Clone, PartialEq, Eq, Hash)]
91pub struct OkxRateLimitKey {
92    pub endpoint: String,
93    pub scope: OkxRateLimitScope,
94}
95
96impl OkxRateLimitKey {
97    pub fn new_ip_based(endpoint: &str) -> Self {
98        Self {
99            endpoint: endpoint.to_string(),
100            scope: OkxRateLimitScope::IpBased,
101        }
102    }
103
104    pub fn new_user_based(endpoint: &str) -> Self {
105        Self {
106            endpoint: endpoint.to_string(),
107            scope: OkxRateLimitScope::UserBased,
108        }
109    }
110
111    pub fn new_user_and_instrument_type(
112        endpoint: &str,
113        instrument_type: OkxInstrumentType,
114    ) -> Self {
115        Self {
116            endpoint: endpoint.to_string(),
117            scope: OkxRateLimitScope::UserAndInstrumentType(instrument_type),
118        }
119    }
120}
121
122#[derive(Debug, Clone)]
123pub struct OkxRateLimitRule {
124    pub key: OkxRateLimitKey,
125    pub requests_per_2_seconds: u32,
126}
xapi_okx/common/ratelimiter.rs

xapi_okx/common/
ratelimiter.rs
