brokapi_runtime/limits.rs
1//! Rate limits and connection quotas
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5use std::time::{Duration, Instant};
6
7use brokapi_core::{ConfigError, Error};
8use serde::Deserialize;
9
10/// Name of a token bucket in a limits file, for example `"global"` or `"orders"`.
11pub type BucketId = String;
12
13/// Name of an operation in the `route` table, for example `"PlaceOrder"`.
14///
15/// The names are chosen by the connector that calls [`BrokerLimiter::acquire`]. A name that has no
16/// entry in the route table falls back to [`DEFAULT_ROUTE`].
17pub type MethodId = String;
18
19/// Route key for "any method that has no route of its own".
20///
21/// The limits files set it to the broker's general bucket, so a call made through `raw()` for a
22/// method the library never named still spends quota instead of going out for free. If the route
23/// table has no `"*"` entry, an unnamed method spends nothing.
24pub const DEFAULT_ROUTE: &str = "*";
25
26/// Token-bucket configuration for one limit bucket.
27///
28/// A bucket starts full with `burst` tokens. It gains `rate` tokens every `per`, continuously (not
29/// in steps), and never holds more than `burst`. One call spends one token. So a quiet client may
30/// send `burst` calls at once, and a busy one is held to `rate` calls per `per`.
31///
32/// In a limits file the three keys are written as `rate = 10`, `per = "1s"`, `burst = 3`. The
33/// `per` value is an integer followed by `ms`, `s`, `m`, `h` or `d` (see [`BrokerLimits::from_toml_str`]).
34///
35/// # Examples
36///
37/// ```
38/// use std::time::Duration;
39/// use brokapi_runtime::BucketCfg;
40///
41/// let cfg = BucketCfg { rate: 10, per: Duration::from_secs(1), burst: 3 };
42/// assert_eq!(cfg.burst, 3);
43/// ```
44#[derive(Debug, Clone, PartialEq, Deserialize)]
45pub struct BucketCfg {
46 /// Tokens added per `per`. Must be greater than zero for the bucket to ever refill.
47 pub rate: u32,
48 /// The period in which `rate` tokens are added.
49 #[serde(with = "humantime_serde")]
50 pub per: Duration,
51 /// Bucket capacity, and the number of tokens it holds at start.
52 pub burst: u32,
53}
54
55/// Connection and subscription quotas of one broker user.
56///
57/// [`BrokerLimiter`] only stores and returns these numbers (see [`BrokerLimiter::conn_quota`]).
58/// It does not open or close connections. A connector decides what to do with them. For example
59/// the Finam connector reads `max_subs_per_conn` and refuses a new stream with
60/// [`brokapi_core::Error::Rejected`] once the quota is used. The other fields are carried for the
61/// connectors that need them and are not enforced by the runtime.
62///
63/// # Examples
64///
65/// ```
66/// use brokapi_runtime::ConnQuota;
67///
68/// let q = ConnQuota::default();
69/// assert_eq!((q.max_conns, q.max_subs_per_conn), (2, 50));
70/// assert!(!q.supports_multiplex);
71/// ```
72#[derive(Debug, Clone, PartialEq, Deserialize)]
73pub struct ConnQuota {
74 /// How many connections the broker lets one user hold at the same time.
75 pub max_conns: usize,
76 /// How many subscriptions (streams) one connection may carry.
77 pub max_subs_per_conn: usize,
78 /// Whether one connection can carry several subscriptions at once.
79 pub supports_multiplex: bool,
80}
81
82impl Default for ConnQuota {
83 /// Cautious values: 2 connections, 50 subscriptions per connection, no multiplexing.
84 fn default() -> Self {
85 Self {
86 max_conns: 2,
87 max_subs_per_conn: 50,
88 supports_multiplex: false,
89 }
90 }
91}
92
93/// The full limit policy of one broker user: token buckets, routes and connection quotas.
94///
95/// Normally read from the `[broker]` section of a limits file by [`from_toml_str`](Self::from_toml_str)
96/// or [`from_toml_section`](Self::from_toml_section), then handed to [`BrokerLimiter::new`].
97///
98/// Limits belong to the broker user (the token or account), not to a connector. Build one
99/// [`BrokerLimiter`] per user and share it between all connectors of that user. Two limiters for
100/// one broker limit would each allow the full rate, so together they would exceed the quota the
101/// broker gave.
102///
103/// # Examples
104///
105/// ```
106/// use brokapi_runtime::BrokerLimits;
107///
108/// let raw = r#"
109/// [demo.buckets.orders]
110/// rate = 10
111/// per = "1s"
112/// burst = 3
113///
114/// [demo.route]
115/// PlaceOrder = ["orders"]
116///
117/// [demo.conns]
118/// max_conns = 2
119/// max_subs_per_conn = 50
120/// supports_multiplex = false
121/// "#;
122/// let limits = BrokerLimits::from_toml_str(raw)?;
123/// assert_eq!(limits.buckets["orders"].burst, 3);
124/// assert_eq!(limits.buckets_for("PlaceOrder"), vec!["orders".to_string()]);
125/// # Ok::<(), brokapi_core::ConfigError>(())
126/// ```
127#[derive(Debug, Clone, PartialEq, Deserialize)]
128pub struct BrokerLimits {
129 /// The token buckets by name.
130 pub buckets: HashMap<BucketId, BucketCfg>,
131 /// For each method name, the buckets one call spends from (one token from each). A bucket
132 /// named here that is missing from `buckets` is a configuration error found at
133 /// [`BrokerLimiter::acquire`].
134 pub route: HashMap<MethodId, Vec<BucketId>>,
135 /// Connection and subscription quotas.
136 pub conns: ConnQuota,
137}
138
139/// Top-level TOML envelope: one section per broker (`[bcs]`, `[tbank]`, `[finam]`).
140type LimitsFile = HashMap<String, BrokerLimits>;
141
142impl BrokerLimits {
143 /// Reads a limits file that holds exactly one broker section.
144 ///
145 /// The file layout is one top-level table per broker, with `buckets`, `route` and `conns`
146 /// below it: `[bcs.buckets.global]`, `[bcs.route]`, `[bcs.conns]`. A `per` value is an integer
147 /// followed by `ms`, `s`, `m`, `h` or `d` (`"250ms"`, `"1s"`, `"24h"`). Zero, fractions (`"1.5s"`), other
148 /// units (`"1w"`) and spaces inside (`"1 s"`) are refused. `buckets` and `route` must be
149 /// present. Keys the parser does not know are ignored (an older file with
150 /// `dedicated_command_conns` still reads), so a misspelled optional key is silently lost.
151 ///
152 /// Use [`from_toml_section`](Self::from_toml_section) for a file that holds several brokers.
153 ///
154 /// # Errors
155 ///
156 /// Returns [`brokapi_core::ConfigError::Message`] when the text is not valid TOML, when a
157 /// required key is missing or has the wrong type, when `per` cannot be parsed, or when the file
158 /// has zero or several broker sections.
159 ///
160 /// # Examples
161 ///
162 /// ```
163 /// use brokapi_runtime::BrokerLimits;
164 ///
165 /// let raw = r#"
166 /// [demo.buckets.global]
167 /// rate = 100
168 /// per = "1s"
169 /// burst = 20
170 /// [demo.route]
171 /// "*" = ["global"]
172 /// [demo.conns]
173 /// max_conns = 4
174 /// max_subs_per_conn = 50
175 /// supports_multiplex = false
176 /// "#;
177 /// let limits = BrokerLimits::from_toml_str(raw)?;
178 /// assert_eq!(limits.conns.max_conns, 4);
179 ///
180 /// // An unsupported duration unit is a configuration error, not a silent default.
181 /// let bad = raw.replace("\"1s\"", "\"1w\"");
182 /// assert!(BrokerLimits::from_toml_str(&bad).is_err());
183 ///
184 /// // Two broker sections: pick one with `from_toml_section`.
185 /// let two = format!("{raw}\n{}", raw.replace("demo", "other"));
186 /// assert!(BrokerLimits::from_toml_str(&two).is_err());
187 /// assert!(BrokerLimits::from_toml_section(&two, "other").is_ok());
188 /// # Ok::<(), brokapi_core::ConfigError>(())
189 /// ```
190 pub fn from_toml_str(raw: &str) -> Result<Self, ConfigError> {
191 let mut file: LimitsFile =
192 toml::from_str(raw).map_err(|e| ConfigError::Message(e.to_string()))?;
193 if file.len() != 1 {
194 return Err(ConfigError::Message(format!(
195 "limits file must have exactly one broker section, found {}; use from_toml_section",
196 file.len()
197 )));
198 }
199 Ok(file.drain().next().expect("one section").1)
200 }
201
202 /// Reads one named broker section (for example `"tbank"`) of a limits file.
203 ///
204 /// Other sections are parsed too, so they must be valid, but they are not returned. The
205 /// format is the same as in [`from_toml_str`](Self::from_toml_str).
206 ///
207 /// # Errors
208 ///
209 /// Returns [`brokapi_core::ConfigError::Message`] for invalid TOML or an invalid section, and
210 /// when the file has no `[section]`.
211 ///
212 /// # Examples
213 ///
214 /// ```
215 /// use brokapi_runtime::BrokerLimits;
216 ///
217 /// let raw = r#"
218 /// [a.buckets.g]
219 /// rate = 1
220 /// per = "1s"
221 /// burst = 1
222 /// [a.route]
223 /// "*" = ["g"]
224 /// [a.conns]
225 /// max_conns = 1
226 /// max_subs_per_conn = 1
227 /// supports_multiplex = false
228 /// "#;
229 /// assert!(BrokerLimits::from_toml_section(raw, "a").is_ok());
230 /// assert!(BrokerLimits::from_toml_section(raw, "b").is_err());
231 /// ```
232 pub fn from_toml_section(raw: &str, section: &str) -> Result<Self, ConfigError> {
233 let mut file: LimitsFile =
234 toml::from_str(raw).map_err(|e| ConfigError::Message(e.to_string()))?;
235 file.remove(section)
236 .ok_or_else(|| ConfigError::Message(format!("no [{section}] section in limits file")))
237 }
238
239 /// Reads a limits file from disk. The file must hold exactly one broker section.
240 ///
241 /// Same as [`from_toml_str`](Self::from_toml_str) after reading the file.
242 ///
243 /// # Errors
244 ///
245 /// Returns [`brokapi_core::ConfigError::Message`] when the file cannot be read (the message
246 /// names the path), and for every reason [`from_toml_str`](Self::from_toml_str) lists.
247 ///
248 /// # Examples
249 ///
250 /// ```
251 /// use brokapi_runtime::BrokerLimits;
252 ///
253 /// let err = BrokerLimits::from_toml_path("/definitely/not/here.toml").unwrap_err();
254 /// assert!(err.to_string().contains("/definitely/not/here.toml"));
255 /// ```
256 pub fn from_toml_path(path: &str) -> Result<Self, ConfigError> {
257 let raw = std::fs::read_to_string(path)
258 .map_err(|e| ConfigError::Message(format!("read {path}: {e}")))?;
259 Self::from_toml_str(&raw)
260 }
261
262 /// The buckets a call to `method` spends from: its own route, else the [`DEFAULT_ROUTE`]
263 /// route, else none.
264 ///
265 /// The ids come back sorted and without repeats. This is the order in which
266 /// [`BrokerLimiter::acquire`] takes the tokens. The ids are not checked against
267 /// [`buckets`](Self::buckets).
268 ///
269 /// # Examples
270 ///
271 /// ```
272 /// use std::collections::HashMap;
273 /// use brokapi_runtime::{BrokerLimits, ConnQuota, DEFAULT_ROUTE};
274 ///
275 /// let limits = BrokerLimits {
276 /// buckets: HashMap::new(),
277 /// route: HashMap::from([
278 /// ("PlaceOrder".to_string(), vec!["orders".to_string(), "global".to_string()]),
279 /// (DEFAULT_ROUTE.to_string(), vec!["global".to_string()]),
280 /// ]),
281 /// conns: ConnQuota::default(),
282 /// };
283 /// assert_eq!(limits.buckets_for("PlaceOrder"), ["global", "orders"]);
284 /// assert_eq!(limits.buckets_for("SomethingNew"), ["global"]);
285 /// ```
286 pub fn buckets_for(&self, method: &str) -> Vec<BucketId> {
287 let mut ids = self
288 .route
289 .get(method)
290 .or_else(|| self.route.get(DEFAULT_ROUTE))
291 .cloned()
292 .unwrap_or_default();
293 ids.sort();
294 ids.dedup();
295 ids
296 }
297}
298
299#[derive(Debug)]
300struct TokenBucket {
301 cfg: BucketCfg,
302 tokens: f64,
303 last_refill: Instant,
304}
305
306impl TokenBucket {
307 fn new(cfg: BucketCfg, now: Instant) -> Self {
308 Self {
309 tokens: cfg.burst as f64,
310 last_refill: now,
311 cfg,
312 }
313 }
314
315 fn refill(&mut self, now: Instant) {
316 if now <= self.last_refill {
317 return;
318 }
319 let elapsed = now.duration_since(self.last_refill);
320 let rate_per_sec = self.cfg.rate as f64 / self.cfg.per.as_secs_f64().max(1e-9);
321 self.tokens =
322 (self.tokens + elapsed.as_secs_f64() * rate_per_sec).min(self.cfg.burst as f64);
323 self.last_refill = now;
324 }
325
326 fn try_acquire(&mut self, now: Instant) -> Result<(), Duration> {
327 self.refill(now);
328 if self.tokens >= 1.0 {
329 self.tokens -= 1.0;
330 Ok(())
331 } else {
332 let deficit = 1.0 - self.tokens;
333 let rate_per_sec = self.cfg.rate as f64 / self.cfg.per.as_secs_f64().max(1e-9);
334 let wait_secs = deficit / rate_per_sec.max(1e-9);
335 Err(Duration::from_secs_f64(wait_secs))
336 }
337 }
338
339 fn release(&mut self) {
340 self.tokens = (self.tokens + 1.0).min(self.cfg.burst as f64);
341 }
342}
343
344/// Shared multi-bucket rate limiter: one instance per broker user.
345///
346/// Every request a connector sends first calls [`acquire`](Self::acquire) with the name of the
347/// method. The limiter spends one token from each bucket on the method's route, or spends nothing
348/// and refuses. It never waits and never sends anything. The caller decides whether to wait for
349/// `retry_after` and try again.
350///
351/// Give the same `Arc<BrokerLimiter>` to every connector of one broker user (market data and user
352/// data together). The broker counts calls per user, so two limiters for one user would each let
353/// through the full rate and together exceed it. A different user (another token) needs its own
354/// limiter.
355///
356/// The limiter is thread-safe. Its internal lock is held only for a short, non-blocking update.
357///
358/// # Examples
359///
360/// ```
361/// use std::collections::HashMap;
362/// use std::time::{Duration, Instant};
363/// use brokapi_core::Error;
364/// use brokapi_runtime::{BrokerLimiter, BrokerLimits, BucketCfg, ConnQuota};
365///
366/// // One bucket of 2 tokens that refills 1 token per second. PlaceOrder spends from it.
367/// let limits = BrokerLimits {
368/// buckets: HashMap::from([(
369/// "orders".to_string(),
370/// BucketCfg { rate: 1, per: Duration::from_secs(1), burst: 2 },
371/// )]),
372/// route: HashMap::from([("PlaceOrder".to_string(), vec!["orders".to_string()])]),
373/// conns: ConnQuota::default(),
374/// };
375/// let limiter = BrokerLimiter::new(&limits);
376///
377/// let t0 = Instant::now();
378/// assert!(limiter.acquire("PlaceOrder", t0).is_ok());
379/// assert!(limiter.acquire("PlaceOrder", t0).is_ok());
380/// // The burst is spent: the third call is refused, nothing was sent.
381/// match limiter.acquire("PlaceOrder", t0) {
382/// Err(Error::RateLimited { retry_after }) => assert_eq!(retry_after, Duration::from_secs(1)),
383/// other => panic!("expected RateLimited, got {other:?}"),
384/// }
385/// // A second later one token is back.
386/// assert!(limiter.acquire("PlaceOrder", t0 + Duration::from_secs(1)).is_ok());
387/// ```
388pub struct BrokerLimiter {
389 inner: Mutex<LimiterInner>,
390}
391
392struct LimiterInner {
393 buckets: HashMap<BucketId, TokenBucket>,
394 route: HashMap<MethodId, Vec<BucketId>>,
395 conns: ConnQuota,
396}
397
398impl BrokerLimiter {
399 /// Creates a limiter from a policy. Every bucket starts full.
400 ///
401 /// The limiter copies the policy, so later changes to `cfg` have no effect on it. Returns an
402 /// [`Arc`] because the limiter is meant to be shared between connectors.
403 ///
404 /// # Examples
405 ///
406 /// ```
407 /// use brokapi_runtime::{BrokerLimiter, BrokerLimits};
408 ///
409 /// # let raw = r#"
410 /// # [demo.buckets.global]
411 /// # rate = 100
412 /// # per = "1s"
413 /// # burst = 20
414 /// # [demo.route]
415 /// # "*" = ["global"]
416 /// # [demo.conns]
417 /// # max_conns = 4
418 /// # max_subs_per_conn = 50
419 /// # supports_multiplex = false
420 /// # "#;
421 /// let limits = BrokerLimits::from_toml_str(raw)?;
422 /// let limiter = BrokerLimiter::new(&limits);
423 /// // Clone the Arc for each connector of the same broker user.
424 /// let for_market_data = std::sync::Arc::clone(&limiter);
425 /// assert_eq!(for_market_data.conn_quota().max_conns, 4);
426 /// # Ok::<(), brokapi_core::ConfigError>(())
427 /// ```
428 pub fn new(cfg: &BrokerLimits) -> Arc<Self> {
429 let now = Instant::now();
430 let buckets = cfg
431 .buckets
432 .iter()
433 .map(|(id, bucket_cfg)| (id.clone(), TokenBucket::new(bucket_cfg.clone(), now)))
434 .collect();
435 Arc::new(Self {
436 inner: Mutex::new(LimiterInner {
437 buckets,
438 route: cfg.route.clone(),
439 conns: cfg.conns.clone(),
440 }),
441 })
442 }
443
444 /// The connection quota from the policy this limiter was built with.
445 ///
446 /// The limiter does not enforce it, see [`ConnQuota`].
447 ///
448 /// # Examples
449 ///
450 /// ```
451 /// use std::collections::HashMap;
452 /// use brokapi_runtime::{BrokerLimiter, BrokerLimits, ConnQuota};
453 ///
454 /// let limits = BrokerLimits {
455 /// buckets: HashMap::new(),
456 /// route: HashMap::new(),
457 /// conns: ConnQuota { max_subs_per_conn: 7, ..ConnQuota::default() },
458 /// };
459 /// assert_eq!(BrokerLimiter::new(&limits).conn_quota().max_subs_per_conn, 7);
460 /// ```
461 pub fn conn_quota(&self) -> ConnQuota {
462 self.inner.lock().expect("limiter poisoned").conns.clone()
463 }
464
465 /// How much of each bucket is left at `now`: `(bucket id, tokens left, capacity)`, sorted by id. Nothing is spent. For
466 /// a program that shows how close a client is to its limits.
467 ///
468 /// # Examples
469 ///
470 /// ```
471 /// use std::collections::HashMap;
472 /// use std::time::{Duration, Instant};
473 /// use brokapi_runtime::{BrokerLimiter, BrokerLimits, BucketCfg, ConnQuota};
474 ///
475 /// let limits = BrokerLimits {
476 /// buckets: HashMap::from([(
477 /// "orders".to_string(),
478 /// BucketCfg { rate: 1, per: Duration::from_secs(1), burst: 5 },
479 /// )]),
480 /// route: HashMap::from([("*".to_string(), vec!["orders".to_string()])]),
481 /// conns: ConnQuota::default(),
482 /// };
483 /// let limiter = BrokerLimiter::new(&limits);
484 /// let t0 = Instant::now();
485 /// limiter.acquire("PlaceOrder", t0).unwrap();
486 /// let usage = limiter.usage(t0);
487 /// assert_eq!(usage[0].0, "orders");
488 /// assert_eq!((usage[0].1, usage[0].2), (4, 5));
489 /// ```
490 pub fn usage(&self, now: Instant) -> Vec<(BucketId, u32, u32)> {
491 let mut inner = self.inner.lock().expect("limiter poisoned");
492 let mut out: Vec<_> = inner
493 .buckets
494 .iter_mut()
495 .map(|(id, b)| {
496 b.refill(now);
497 (id.clone(), b.tokens.floor() as u32, b.cfg.burst)
498 })
499 .collect();
500 out.sort();
501 out
502 }
503
504 /// Spends one token from every bucket on the route of `method`, all or nothing.
505 ///
506 /// The route is the method's own entry, else the [`DEFAULT_ROUTE`] entry, else empty (then the
507 /// call is free). Buckets are taken in sorted id order. If any bucket is empty, the tokens
508 /// already taken from the earlier buckets are given back (up to each bucket's `burst`) and the
509 /// call is refused. A refused call therefore costs nothing.
510 ///
511 /// `now` is the current time. Pass `Instant::now()` in production. Tests pass a made-up
512 /// instant to move time without sleeping. Time must not go backwards: an instant earlier than
513 /// the last one seen refills nothing.
514 ///
515 /// The `retry_after` of a refusal is the time the first empty bucket (in sorted order) needs
516 /// to get one token back. Another bucket on the route may need longer, so a retry after
517 /// `retry_after` can be refused again.
518 ///
519 /// A refusal is local: nothing was sent to the broker, so repeating the call later cannot
520 /// create a duplicate order.
521 ///
522 /// # Errors
523 ///
524 /// - [`brokapi_core::Error::RateLimited`] when a bucket on the route has less than one token.
525 /// - [`brokapi_core::Error::Config`] when the route names a bucket that is not defined. No
526 /// token is spent in this case.
527 ///
528 /// # Examples
529 ///
530 /// ```
531 /// use std::collections::HashMap;
532 /// use std::time::{Duration, Instant};
533 /// use brokapi_core::Error;
534 /// use brokapi_runtime::{BrokerLimiter, BrokerLimits, BucketCfg, ConnQuota};
535 ///
536 /// let one_per_sec = |burst| BucketCfg { rate: 1, per: Duration::from_secs(1), burst };
537 /// let limits = BrokerLimits {
538 /// buckets: HashMap::from([
539 /// ("global".to_string(), one_per_sec(10)),
540 /// ("orders".to_string(), one_per_sec(1)),
541 /// ]),
542 /// route: HashMap::from([
543 /// ("PlaceOrder".to_string(), vec!["global".to_string(), "orders".to_string()]),
544 /// ("GetOrder".to_string(), vec!["global".to_string()]),
545 /// ("Broken".to_string(), vec!["nope".to_string()]),
546 /// ]),
547 /// conns: ConnQuota::default(),
548 /// };
549 /// let limiter = BrokerLimiter::new(&limits);
550 /// let t0 = Instant::now();
551 ///
552 /// assert!(limiter.acquire("PlaceOrder", t0).is_ok());
553 /// // "orders" is empty now. The refusal gives back the token taken from "global".
554 /// assert!(matches!(limiter.acquire("PlaceOrder", t0), Err(Error::RateLimited { .. })));
555 /// // A method with no route is free.
556 /// assert!(limiter.acquire("Unlisted", t0).is_ok());
557 /// // A route to a bucket that does not exist is a configuration error.
558 /// assert!(matches!(limiter.acquire("Broken", t0), Err(Error::Config(_))));
559 /// ```
560 pub fn acquire(&self, method: &str, now: Instant) -> Result<(), Error> {
561 let mut inner = self.inner.lock().expect("limiter poisoned");
562 let bucket_ids = {
563 let mut ids = inner
564 .route
565 .get(method)
566 .or_else(|| inner.route.get(DEFAULT_ROUTE))
567 .cloned()
568 .unwrap_or_default();
569 ids.sort();
570 ids.dedup();
571 ids
572 };
573
574 // A route to a bucket that does not exist is refused before any token is taken, so the
575 // buckets that sort ahead of the missing one do not lose a token for nothing.
576 if let Some(missing) = bucket_ids
577 .iter()
578 .find(|id| !inner.buckets.contains_key(*id))
579 {
580 return Err(Error::Config(ConfigError::Message(format!(
581 "unknown bucket {missing}"
582 ))));
583 }
584
585 let mut acquired = Vec::new();
586 for id in bucket_ids {
587 let bucket = inner.buckets.get_mut(&id).ok_or_else(|| {
588 Error::Config(ConfigError::Message(format!("unknown bucket {id}")))
589 })?;
590 match bucket.try_acquire(now) {
591 Ok(()) => acquired.push(id),
592 Err(retry_after) => {
593 for rollback in acquired.iter().rev() {
594 if let Some(b) = inner.buckets.get_mut(rollback) {
595 b.release();
596 }
597 }
598 return Err(Error::RateLimited { retry_after });
599 }
600 }
601 }
602 Ok(())
603 }
604}
605
606mod humantime_serde {
607 use serde::{Deserialize, Deserializer};
608 use std::time::Duration;
609
610 pub fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
611 where
612 D: Deserializer<'de>,
613 {
614 let s = String::deserialize(deserializer)?;
615 parse_duration(&s).map_err(serde::de::Error::custom)
616 }
617
618 fn parse_duration(s: &str) -> Result<Duration, String> {
619 let s = s.trim();
620 // "ms" must be tried before "s" and "m".
621 let (num, unit_ms) = if let Some(n) = s.strip_suffix("ms") {
622 (n, 1u64)
623 } else if let Some(n) = s.strip_suffix('s') {
624 (n, 1_000)
625 } else if let Some(n) = s.strip_suffix('m') {
626 (n, 60_000)
627 } else if let Some(n) = s.strip_suffix('h') {
628 (n, 3_600_000)
629 } else if let Some(n) = s.strip_suffix('d') {
630 (n, 86_400_000)
631 } else {
632 return Err(format!("unsupported duration: {s}"));
633 };
634 let v: u64 = num
635 .parse()
636 .map_err(|e: std::num::ParseIntError| e.to_string())?;
637 if v == 0 {
638 return Err(format!("a period must be longer than zero: {s}"));
639 }
640 let ms = v
641 .checked_mul(unit_ms)
642 .ok_or_else(|| format!("period too long: {s}"))?;
643 Ok(Duration::from_millis(ms))
644 }
645}
646
647#[cfg(test)]
648mod tests {
649 use super::*;
650
651 fn sample_limits() -> BrokerLimits {
652 BrokerLimits {
653 buckets: HashMap::from([
654 (
655 "global".into(),
656 BucketCfg {
657 rate: 10,
658 per: Duration::from_secs(1),
659 burst: 10,
660 },
661 ),
662 (
663 "orders".into(),
664 BucketCfg {
665 rate: 2,
666 per: Duration::from_secs(1),
667 burst: 2,
668 },
669 ),
670 ]),
671 route: HashMap::from([
672 ("PlaceOrder".into(), vec!["global".into(), "orders".into()]),
673 ("CancelOrder".into(), vec!["global".into()]),
674 ]),
675 conns: ConnQuota::default(),
676 }
677 }
678
679 #[test]
680 fn sequential_acquire_rolls_back_on_failure() {
681 let limiter = BrokerLimiter::new(&sample_limits());
682 let t0 = Instant::now();
683 assert!(limiter.acquire("PlaceOrder", t0).is_ok());
684 assert!(limiter.acquire("PlaceOrder", t0).is_ok());
685 assert!(matches!(
686 limiter.acquire("PlaceOrder", t0),
687 Err(Error::RateLimited { .. })
688 ));
689 assert!(limiter.acquire("CancelOrder", t0).is_ok());
690 }
691
692 #[test]
693 fn unknown_method_falls_back_to_the_default_route() {
694 let mut limits = sample_limits();
695 // Without a default route an unnamed method is free (the old behavior).
696 let free = BrokerLimiter::new(&limits);
697 let t0 = Instant::now();
698 for _ in 0..50 {
699 assert!(free.acquire("SomethingNew", t0).is_ok());
700 }
701 limits
702 .route
703 .insert(DEFAULT_ROUTE.into(), vec!["global".into(), "orders".into()]);
704 assert_eq!(
705 limits.buckets_for("SomethingNew"),
706 vec!["global".to_string(), "orders".to_string()]
707 );
708 // A method with its own route keeps it.
709 assert_eq!(
710 limits.buckets_for("CancelOrder"),
711 vec!["global".to_string()]
712 );
713 let limiter = BrokerLimiter::new(&limits);
714 assert!(limiter.acquire("SomethingNew", t0).is_ok());
715 assert!(limiter.acquire("SomethingNew", t0).is_ok());
716 assert!(matches!(
717 limiter.acquire("SomethingNew", t0),
718 Err(Error::RateLimited { .. })
719 ));
720 // The shared bucket is the same one the named methods spend.
721 assert!(matches!(
722 limiter.acquire("PlaceOrder", t0),
723 Err(Error::RateLimited { .. })
724 ));
725 }
726
727 #[test]
728 fn a_file_with_the_removed_key_still_reads() {
729 // `dedicated_command_conns` was in the files of older versions and was never read.
730 let raw = r#"
731[demo.buckets.global]
732rate = 100
733per = "1s"
734burst = 20
735[demo.route]
736"*" = ["global"]
737[demo.conns]
738max_conns = 4
739max_subs_per_conn = 50
740supports_multiplex = false
741dedicated_command_conns = 2
742"#;
743 let limits = BrokerLimits::from_toml_str(raw).expect("an old file reads");
744 assert_eq!(limits.conns.max_conns, 4);
745 }
746
747 #[test]
748 fn unknown_bucket_in_route_spends_nothing() {
749 let mut limits = sample_limits();
750 // "global" sorts ahead of "zzz-missing": it must not lose a token to the failed call.
751 limits
752 .route
753 .insert("Broken".into(), vec!["global".into(), "zzz-missing".into()]);
754 let limiter = BrokerLimiter::new(&limits);
755 let t0 = Instant::now();
756 for _ in 0..50 {
757 assert!(matches!(
758 limiter.acquire("Broken", t0),
759 Err(Error::Config(_))
760 ));
761 }
762 // The global bucket (burst 10) is still full.
763 for _ in 0..10 {
764 assert!(limiter.acquire("CancelOrder", t0).is_ok());
765 }
766 assert!(matches!(
767 limiter.acquire("CancelOrder", t0),
768 Err(Error::RateLimited { .. })
769 ));
770 }
771}