dre_core/
rate.rs

1//! Traffic control for deliveries: a destination's `rate:` (`100/minute`, `5/second`,
2//! `20000/hour`), with an optional `burst:`. Deliveries are spaced evenly (`60/minute` sends one
3//! a second); `burst: N` lets up to N go back to back while keeping the same average. One pacer
4//! serves every worker of a run or release, so the rate holds per destination, whatever the
5//! concurrency.
6
7use std::collections::HashMap;
8use std::sync::Mutex;
9use std::time::{Duration, Instant};
10
11use serde::{Deserialize, Serialize};
12
13use crate::engine::CancelToken;
14
15/// A delivery rate: `count` deliveries per `per`, and how many may go back to back.
16#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
17pub struct Rate {
18    pub count: u64,
19    /// The window, in seconds.
20    pub per_seconds: u64,
21    #[serde(default = "one")]
22    pub burst: u64,
23}
24
25fn one() -> u64 {
26    1
27}
28
29impl Rate {
30    /// Parse `N/unit`: the unit is `second`, `minute` or `hour` (or `s`, `m`, `h`, plurals).
31    pub fn parse(s: &str) -> Result<Rate, String> {
32        let bad = || {
33            format!(
34                "`{s}` isn't a rate; write a count and a unit, like `100/minute`, `5/second` or `20000/hour`"
35            )
36        };
37        let (n, unit) = s.split_once('/').ok_or_else(bad)?;
38        let count: u64 = n.trim().replace('_', "").parse().map_err(|_| bad())?;
39        if count == 0 {
40            return Err(format!("`{s}`: the count must be 1 or more"));
41        }
42        let per_seconds = match unit.trim().to_ascii_lowercase().trim_end_matches('s') {
43            "second" | "sec" | "" => 1,
44            "minute" | "min" | "m" => 60,
45            "hour" | "hr" | "h" => 3600,
46            _ => return Err(bad()),
47        };
48        Ok(Rate {
49            count,
50            per_seconds,
51            burst: 1,
52        })
53    }
54
55    /// The even gap between two deliveries.
56    pub fn interval(&self) -> Duration {
57        Duration::from_secs_f64(self.per_seconds as f64 / self.count as f64)
58    }
59
60    /// As written: `100/minute`.
61    pub fn shown(&self) -> String {
62        let unit = match self.per_seconds {
63            1 => "second",
64            60 => "minute",
65            _ => "hour",
66        };
67        format!("{}/{unit}", self.count)
68    }
69}
70
71/// Paces deliveries per destination (by a key the caller chooses).
72#[derive(Default)]
73pub struct Pacer {
74    /// Each destination's theoretical arrival time of the next delivery.
75    next: Mutex<HashMap<String, Instant>>,
76}
77
78impl Pacer {
79    /// Wait for destination `key`'s next slot under `rate`. `false`: the run was cancelled meanwhile.
80    pub fn wait(&self, key: &str, rate: &Rate, cancel: &CancelToken) -> bool {
81        let gap = rate.interval();
82        // Up to `burst` deliveries may run ahead of the even schedule.
83        let slack = gap * (rate.burst.max(1) - 1) as u32;
84        let until = {
85            let mut next = self.next.lock().unwrap_or_else(|e| e.into_inner());
86            let now = Instant::now();
87            let tat = next.get(key).copied().unwrap_or(now).max(now);
88            let allowed = tat.checked_sub(slack).map_or(now, |t| t.max(now));
89            next.insert(key.to_string(), tat + gap);
90            allowed
91        };
92        while Instant::now() < until {
93            if cancel.is_cancelled() {
94                return false;
95            }
96            std::thread::sleep((until - Instant::now()).min(Duration::from_millis(100)));
97        }
98        !cancel.is_cancelled()
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn parses_rates() {
108        assert_eq!(
109            Rate::parse("100/minute").unwrap().interval(),
110            Duration::from_millis(600)
111        );
112        assert_eq!(
113            Rate::parse("5/second").unwrap().interval(),
114            Duration::from_millis(200)
115        );
116        assert_eq!(
117            Rate::parse("3600/hour").unwrap().interval(),
118            Duration::from_secs(1)
119        );
120        assert_eq!(Rate::parse("60 / Minutes").unwrap().shown(), "60/minute");
121        assert_eq!(Rate::parse("20_000/h").unwrap().count, 20_000);
122        for bad in ["100", "0/minute", "x/minute", "10/day"] {
123            assert!(Rate::parse(bad).is_err(), "{bad}");
124        }
125    }
126
127    #[test]
128    fn spaces_evenly_and_allows_a_burst() {
129        let cancel = CancelToken::new();
130        let p = Pacer::default();
131        let even = Rate::parse("20/second").unwrap();
132        let t = Instant::now();
133        for _ in 0..5 {
134            assert!(p.wait("a", &even, &cancel));
135        }
136        // Four gaps of 50 ms.
137        assert!(t.elapsed() >= Duration::from_millis(190), "{:?}", t.elapsed());
138        let burst = Rate {
139            burst: 5,
140            ..Rate::parse("2/second").unwrap()
141        };
142        let t = Instant::now();
143        for _ in 0..5 {
144            assert!(p.wait("b", &burst, &cancel));
145        }
146        // Five back to back, then the sixth waits for the average.
147        assert!(t.elapsed() < Duration::from_millis(100), "{:?}", t.elapsed());
148        assert!(p.wait("b", &burst, &cancel));
149        assert!(t.elapsed() >= Duration::from_millis(400), "{:?}", t.elapsed());
150    }
151}
dre_core/rate.rs

dre_core/
rate.rs
