Skip to main content

brokapi_bcs_protocol/utils/
rate_limiter.rs

1//! A small token-bucket rate limiter for the broker's request limits.
2
3use std::sync::Arc;
4use tokio::sync::Mutex;
5use tokio::time::{Duration, Instant};
6
7/// A token-bucket rate limiter: it lets at most `max_rps` calls per second pass, with bursts up to
8/// `max_rps`.
9///
10/// The BCS API limits REST calls per account, about 10 per second for most services and 3 per second
11/// for non-trade operations (broker docs). The HTTP client does **not** throttle on its own. Put one
12/// limiter in front of the calls of a service and call [`RateLimiter::acquire`] before each call.
13/// A limiter is cheap to clone, and clones share the same bucket, so one limiter can guard calls from
14/// many tasks.
15///
16/// This limiter is for the protocol layer. It does not know the limits of the broker: you give it the
17/// number. The connector crate has its own per-service limits.
18///
19/// # Examples
20///
21/// ```
22/// # #[tokio::main(flavor = "current_thread")]
23/// # async fn main() {
24/// use brokapi_bcs_protocol::utils::rate_limiter::RateLimiter;
25///
26/// let limiter = RateLimiter::new(10); // 10 calls per second
27/// for _ in 0..3 {
28///     limiter.acquire().await; // returns at once while the bucket has tokens
29///     // client.get_order_book(...).await
30/// }
31/// # }
32/// ```
33#[derive(Debug, Clone)]
34pub struct RateLimiter {
35    inner: Arc<Mutex<RateLimiterInner>>,
36}
37
38#[derive(Debug)]
39struct RateLimiterInner {
40    max_rps: u32,
41    tokens: f64,
42    last_refill: Instant,
43}
44
45impl RateLimiter {
46    /// Creates a limiter that allows `max_rps` calls per second. The bucket starts full, so the first
47    /// `max_rps` calls pass without waiting.
48    ///
49    /// `max_rps` must be above 0. With 0 the limiter would wait for a token that never comes, and
50    /// [`RateLimiter::acquire`] would not return.
51    ///
52    /// # Examples
53    ///
54    /// ```
55    /// use brokapi_bcs_protocol::utils::rate_limiter::RateLimiter;
56    ///
57    /// let limiter = RateLimiter::new(3);
58    /// let _same_bucket = limiter.clone();
59    /// ```
60    pub fn new(max_rps: u32) -> Self {
61        Self {
62            inner: Arc::new(Mutex::new(RateLimiterInner {
63                max_rps,
64                tokens: max_rps as f64,
65                last_refill: Instant::now(),
66            })),
67        }
68    }
69
70    /// Waits until one call is allowed, then takes a token.
71    ///
72    /// Returns at once if the bucket has a token, otherwise sleeps for the time one token needs to
73    /// grow back. Tokens grow back at `max_rps` per second. There is no timeout, and no ordering
74    /// guarantee between waiting tasks.
75    ///
76    /// # Examples
77    ///
78    /// ```
79    /// # #[tokio::main(flavor = "current_thread")]
80    /// # async fn main() {
81    /// use brokapi_bcs_protocol::utils::rate_limiter::RateLimiter;
82    /// use std::time::{Duration, Instant};
83    ///
84    /// let limiter = RateLimiter::new(100);
85    /// let started = Instant::now();
86    /// for _ in 0..100 {
87    ///     limiter.acquire().await; // a full bucket: no waiting
88    /// }
89    /// assert!(started.elapsed() < Duration::from_millis(500));
90    /// # }
91    /// ```
92    pub async fn acquire(&self) {
93        loop {
94            let sleep_dur = {
95                let mut inner = self.inner.lock().await;
96                let now = Instant::now();
97                let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
98                inner.tokens += elapsed * inner.max_rps as f64;
99                if inner.tokens > inner.max_rps as f64 {
100                    inner.tokens = inner.max_rps as f64;
101                }
102                inner.last_refill = now;
103
104                if inner.tokens >= 1.0 {
105                    inner.tokens -= 1.0;
106                    return;
107                }
108                let wait = (1.0 - inner.tokens) / inner.max_rps as f64;
109                Duration::from_secs_f64(wait)
110            };
111            tokio::time::sleep(sleep_dur).await;
112        }
113    }
114}