tracelimit/
lib.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Implementations of [`tracing`] macros that are rate limited.
5//!
6//! These are useful when an untrusted entity can arbitrarily trigger an event
7//! to be logged. In that case, rate limiting the events can prevent other
8//! important events from being lost.
9//!
10//! Note that there currently are no macros for rate limiting debug or trace
11//! level events. This is due to an implementation limitation--the macros in
12//! this crate check the rate limit before evaluating whether the event should
13//! be logged, and so this would add an extra check to every debug/trace event.
14//! This could be fixed by using some of the hidden but `pub` machinery of the
15//! `tracing` crate, which is probably not worth it for now.
16
17#![forbid(unsafe_code)]
18
19use parking_lot::Mutex;
20use std::sync::atomic::AtomicBool;
21use std::sync::atomic::Ordering;
22use std::time::Instant;
23#[doc(hidden)]
24pub use tracing;
25
26const PERIOD_MS: u32 = 5000;
27const EVENTS_PER_PERIOD: u32 = 10;
28
29static DISABLE_RATE_LIMITING: AtomicBool = AtomicBool::new(false);
30
31/// Disables or reenables rate limiting globally.
32///
33/// Rate limiting defaults to enabled. You might want to disable it during local
34/// development or in tests.
35pub fn disable_rate_limiting(disabled: bool) {
36    DISABLE_RATE_LIMITING.store(disabled, Ordering::Relaxed);
37}
38
39#[doc(hidden)]
40pub struct RateLimiter {
41    period_ms: u32,
42    events_per_period: u32,
43    state: Mutex<RateLimiterState>,
44}
45
46struct RateLimiterState {
47    start: Option<Instant>,
48    events: u32,
49    missed: u64,
50}
51
52#[doc(hidden)]
53pub struct RateLimited;
54
55impl RateLimiter {
56    pub const fn new_default() -> Self {
57        Self::new(PERIOD_MS, EVENTS_PER_PERIOD)
58    }
59
60    pub const fn new(period_ms: u32, events_per_period: u32) -> Self {
61        Self {
62            period_ms,
63            events_per_period,
64            state: Mutex::new(RateLimiterState {
65                start: None,
66                events: 0,
67                missed: 0,
68            }),
69        }
70    }
71
72    /// Returns `Ok(missed_events)` if this event should be logged.
73    ///
74    /// `missed_events` is `Some(n)` if there were any missed events or if this
75    /// event is the last one before rate limiting kicks in.
76    pub fn event(&self) -> Result<Option<u64>, RateLimited> {
77        if DISABLE_RATE_LIMITING.load(Ordering::Relaxed) {
78            return Ok(None);
79        }
80        let mut state = self.state.try_lock().ok_or(RateLimited)?;
81        let now = Instant::now();
82        let start = state.start.get_or_insert(now);
83        let elapsed = now.duration_since(*start);
84        if elapsed.as_millis() > self.period_ms as u128 {
85            *start = now;
86            state.events = 0;
87        }
88        if state.events >= self.events_per_period {
89            state.missed += 1;
90            return Err(RateLimited);
91        }
92        state.events += 1;
93        let missed = std::mem::take(&mut state.missed);
94        let missed = (missed != 0 || state.events == self.events_per_period).then_some(missed);
95        Ok(missed)
96    }
97}
98
99/// As [`tracing::error!`], but rate limited.
100#[macro_export]
101macro_rules! error_ratelimited {
102    ($($rest:tt)*) => {
103        {
104            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
105            if let Ok(missed_events) = RATE_LIMITER.event() {
106                $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*);
107            }
108        }
109    };
110}
111
112/// As [`tracing::warn!`], but rate limited.
113#[macro_export]
114macro_rules! warn_ratelimited {
115    ($($rest:tt)*) => {
116        {
117            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
118            if let Ok(missed_events) = RATE_LIMITER.event() {
119                $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*);
120            }
121        }
122    };
123}
124
125/// As [`tracing::info!`], but rate limited.
126#[macro_export]
127macro_rules! info_ratelimited {
128    ($($rest:tt)*) => {
129        {
130            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
131            if let Ok(missed_events) = RATE_LIMITER.event() {
132                $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*);
133            }
134        }
135    };
136}
tracelimit/lib.rs

tracelimit/
lib.rs
