reth_tracing/
throttle.rs

1//! Throttling utilities for rate-limiting expression execution.
2
3use std::{
4    sync::{
5        atomic::{AtomicU64, Ordering},
6        LazyLock,
7    },
8    time::Instant,
9};
10
11/// Sentinel value indicating the throttle has never run.
12#[doc(hidden)]
13pub const NOT_YET_RUN: u64 = u64::MAX;
14
15/// Checks if enough time has passed since the last run. Implementation detail for [`throttle!`].
16#[doc(hidden)]
17pub fn should_run(start: &LazyLock<Instant>, last: &AtomicU64, duration_millis: u64) -> bool {
18    let now = start.elapsed().as_millis() as u64;
19    let last_val = last.load(Ordering::Relaxed);
20
21    if last_val == NOT_YET_RUN {
22        return last
23            .compare_exchange(NOT_YET_RUN, now, Ordering::Relaxed, Ordering::Relaxed)
24            .is_ok();
25    }
26
27    if now.saturating_sub(last_val) >= duration_millis {
28        last.compare_exchange(last_val, now, Ordering::Relaxed, Ordering::Relaxed).is_ok()
29    } else {
30        false
31    }
32}
33
34/// Throttles the execution of an expression to run at most once per specified duration.
35///
36/// Uses static variables with lazy initialization to track the last execution time.
37/// Thread-safe via atomic operations.
38///
39/// # Examples
40///
41/// ```ignore
42/// use std::time::Duration;
43/// use reth_tracing::throttle;
44///
45/// // Log at most once per second.
46/// throttle!(Duration::from_secs(1), || {
47///     tracing::info!("This message is throttled");
48/// });
49/// ```
50#[macro_export]
51macro_rules! throttle {
52    ($duration:expr, || $expr:expr) => {{
53        static START: ::std::sync::LazyLock<::std::time::Instant> =
54            ::std::sync::LazyLock::new(::std::time::Instant::now);
55        static LAST: ::core::sync::atomic::AtomicU64 =
56            ::core::sync::atomic::AtomicU64::new($crate::__private::NOT_YET_RUN);
57
58        if $crate::__private::should_run(&START, &LAST, $duration.as_millis() as u64) {
59            $expr
60        }
61    }};
62}
63
64#[cfg(test)]
65mod tests {
66    use std::sync::atomic::{AtomicUsize, Ordering};
67
68    #[test]
69    fn test_throttle_runs_once_initially() {
70        static COUNTER: AtomicUsize = AtomicUsize::new(0);
71
72        throttle!(std::time::Duration::from_secs(10), || {
73            COUNTER.fetch_add(1, Ordering::SeqCst);
74        });
75
76        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
77    }
78}