Skip to main content

krabka_bench_driver/
rate.rs

1//! Token-bucket pacer for `FixedRate` scenarios. `saturate` mode never
2//! calls this. It only `await`s the producer's natural backpressure.
3
4use {
5    krabka_units::prelude::*,
6    tokio::time::{Instant, Interval, MissedTickBehavior, interval_at},
7};
8
9/// The slowest rate the pacer runs at. A scenario that asks for less, or for
10/// nothing at all, clamps to this rate and does not sleep for an unbounded
11/// period.
12///
13/// Public because [`Pacer::new`] documents the clamp and a reader of that
14/// documentation has to be able to follow it to the number.
15pub const MIN_RATE: Frequency = per_sec(1);
16
17/// A simple steady-rate pacer.
18///
19/// Every `await_token().await` returns after the bucket has accumulated one
20/// whole token. The schedule is pinned to the pacer's creation time, so
21/// per-call drift cannot accumulate.
22pub struct Pacer {
23    inner: Interval,
24}
25
26impl Pacer {
27    /// Paces at `rate`. This clamps up to [`MIN_RATE`] when the scenario asks
28    /// for less than one message a second.
29    #[must_use]
30    pub fn new(rate: Frequency) -> Self {
31        let period = if rate < MIN_RATE { MIN_RATE } else { rate }.period();
32        let period = period.to_std();
33        let mut inner = interval_at(Instant::now() + period, period);
34        // If a slow consumer falls behind, deliver one tick per call
35        // rather than burst-catching-up to "now" — keeps the per-message
36        // latency comparable across runs.
37        inner.set_missed_tick_behavior(MissedTickBehavior::Delay);
38        Self { inner }
39    }
40
41    /// Sleeps until the next token is available.
42    pub async fn await_token(&mut self) { self.inner.tick().await; }
43}
44
45#[cfg(test)]
46mod tests {
47    use {super::*, assert2::check};
48
49    #[tokio::test(start_paused = true)]
50    async fn pacer_spaces_ticks() {
51        let mut p = Pacer::new(per_sec(1000)); // 1 ms period
52        let start = Instant::now();
53        for _ in 0..5 {
54            p.await_token().await;
55        }
56        check!(start.elapsed().as_time() >= millis(5));
57    }
58
59    #[tokio::test(start_paused = true)]
60    async fn sub_minimum_rates_are_clamped_to_one_per_second() {
61        // Constructed inside a runtime — clamps to `MIN_RATE`.
62        let mut p = Pacer::new(Frequency::ZERO);
63        let start = Instant::now();
64        p.await_token().await;
65        check!(start.elapsed().as_time() == secs(1));
66    }
67}