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}