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    state: Mutex<RateLimiterState>,
42}
43
44struct RateLimiterState {
45    period_ms: u32,
46    events_per_period: u32,
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            state: Mutex::new(RateLimiterState {
63                period_ms,
64                events_per_period,
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        self.event_with_config(None, None)
78    }
79
80    /// Returns `Ok(missed_events)` if this event should be logged.
81    /// Optionally reconfigures the rate limiter if new parameters are provided.
82    ///
83    /// `missed_events` is `Some(n)` if there were any missed events or if this
84    /// event is the last one before rate limiting kicks in.
85    #[cold]
86    pub fn event_with_config(
87        &self,
88        period_ms: Option<u32>,
89        events_per_period: Option<u32>,
90    ) -> Result<Option<u64>, RateLimited> {
91        if DISABLE_RATE_LIMITING.load(Ordering::Relaxed) {
92            return Ok(None);
93        }
94
95        let mut state = self.state.try_lock().ok_or(RateLimited)?;
96
97        // Reconfigure if new parameters are provided
98        let mut reset_state = false;
99        if let Some(new_period) = period_ms {
100            if state.period_ms != new_period {
101                state.period_ms = new_period;
102                reset_state = true;
103            }
104        }
105        if let Some(new_events_per_period) = events_per_period {
106            if state.events_per_period != new_events_per_period {
107                state.events_per_period = new_events_per_period;
108                reset_state = true;
109            }
110        }
111
112        // Reset state when parameters change
113        if reset_state {
114            state.start = None;
115            state.events = 0;
116            state.missed = 0;
117        }
118
119        let now = Instant::now();
120        let period_ms = state.period_ms;
121        let start = state.start.get_or_insert(now);
122        let elapsed = now.duration_since(*start);
123        if elapsed.as_millis() > period_ms as u128 {
124            *start = now;
125            state.events = 0;
126        }
127        if state.events >= state.events_per_period {
128            state.missed += 1;
129            return Err(RateLimited);
130        }
131        state.events += 1;
132        let missed = std::mem::take(&mut state.missed);
133        let missed = (missed != 0 || state.events == state.events_per_period).then_some(missed);
134        Ok(missed)
135    }
136}
137
138/// As [`tracing::error!`], but rate limited.
139///
140/// Can be called with optional parameters to customize rate limiting:
141/// - `period: <ms>` - rate limiting period in milliseconds
142/// - `limit: <count>` - maximum events per period
143///
144/// Examples:
145/// ```
146/// use tracelimit::error_ratelimited;
147/// error_ratelimited!("simple message");
148/// error_ratelimited!(period: 1000, limit: 5, "custom rate limit");
149/// error_ratelimited!(period: 10000, "custom period only");
150/// error_ratelimited!(limit: 50, "custom limit only");
151/// ```
152#[macro_export]
153macro_rules! error_ratelimited {
154    // With both period and limit
155    (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => {
156        {
157            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
158            if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) {
159                $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*);
160            }
161        }
162    };
163    // With period only
164    (period: $period:expr, $($rest:tt)*) => {
165        {
166            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
167            if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) {
168                $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*);
169            }
170        }
171    };
172    // With limit only
173    (limit: $limit:expr, $($rest:tt)*) => {
174        {
175            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
176            if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) {
177                $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*);
178            }
179        }
180    };
181    // Default case (no custom parameters)
182    ($($rest:tt)*) => {
183        {
184            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
185            if let Ok(missed_events) = RATE_LIMITER.event() {
186                $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*);
187            }
188        }
189    };
190}
191
192/// As [`tracing::warn!`], but rate limited.
193///
194/// Can be called with optional parameters to customize rate limiting:
195/// - `period: <ms>` - rate limiting period in milliseconds
196/// - `limit: <count>` - maximum events per period
197///
198/// Examples:
199/// ```
200/// use tracelimit::warn_ratelimited;
201/// warn_ratelimited!("simple message");
202/// warn_ratelimited!(period: 1000, limit: 5, "custom rate limit");
203/// warn_ratelimited!(period: 10000, "custom period only");
204/// warn_ratelimited!(limit: 50, "custom limit only");
205/// ```
206#[macro_export]
207macro_rules! warn_ratelimited {
208    // With both period and limit
209    (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => {
210        {
211            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
212            if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) {
213                $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*);
214            }
215        }
216    };
217    // With period only
218    (period: $period:expr, $($rest:tt)*) => {
219        {
220            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
221            if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) {
222                $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*);
223            }
224        }
225    };
226    // With limit only
227    (limit: $limit:expr, $($rest:tt)*) => {
228        {
229            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
230            if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) {
231                $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*);
232            }
233        }
234    };
235    // Default case (no custom parameters)
236    ($($rest:tt)*) => {
237        {
238            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
239            if let Ok(missed_events) = RATE_LIMITER.event() {
240                $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*);
241            }
242        }
243    };
244}
245
246/// As [`tracing::info!`], but rate limited.
247///
248/// Can be called with optional parameters to customize rate limiting:
249/// - `period: <ms>` - rate limiting period in milliseconds
250/// - `limit: <count>` - maximum events per period
251///
252/// Examples:
253/// ```
254/// use tracelimit::info_ratelimited;
255/// info_ratelimited!("simple message");
256/// info_ratelimited!(period: 1000, limit: 5, "custom rate limit");
257/// info_ratelimited!(period: 10000, "custom period only");
258/// info_ratelimited!(limit: 50, "custom limit only");
259/// ```
260#[macro_export]
261macro_rules! info_ratelimited {
262    // With both period and limit
263    (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => {
264        {
265            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
266            if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) {
267                $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*);
268            }
269        }
270    };
271    // With period only
272    (period: $period:expr, $($rest:tt)*) => {
273        {
274            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
275            if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) {
276                $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*);
277            }
278        }
279    };
280    // With limit only
281    (limit: $limit:expr, $($rest:tt)*) => {
282        {
283            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
284            if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) {
285                $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*);
286            }
287        }
288    };
289    // Default case (no custom parameters)
290    ($($rest:tt)*) => {
291        {
292            static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
293            if let Ok(missed_events) = RATE_LIMITER.event() {
294                $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*);
295            }
296        }
297    };
298}
299
300/// As [`tracing::event!`], but rate limited.
301///
302/// Can be called with optional parameters to customize rate limiting:
303/// - `period: <ms>` - rate limiting period in milliseconds
304/// - `limit: <count>` - maximum events per period
305///
306/// `level` is required and must be a compile-time literal identifier (ERROR, WARN, INFO, DEBUG, TRACE).
307#[macro_export]
308macro_rules! event_ratelimited_static {
309    // With both period and limit and level
310    (level: $level:ident, period: $period:expr, limit: $limit:expr, $($rest:tt)*) => {{
311        static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
312        if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) {
313            $crate::tracing::event!(
314                $crate::tracing::Level::$level,
315                dropped_ratelimited = missed_events,
316                $($rest)*
317            );
318        }
319    }};
320    // With only period and level
321    (level: $level:ident, period: $period:expr, $($rest:tt)*) => {{
322        static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
323        if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) {
324            $crate::tracing::event!(
325                $crate::tracing::Level::$level,
326                dropped_ratelimited = missed_events,
327                $($rest)*
328            );
329        }
330    }};
331    // With only limit and level
332    (level: $level:ident, limit: $limit:expr, $($rest:tt)*) => {{
333        static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
334        if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) {
335            $crate::tracing::event!(
336                $crate::tracing::Level::$level,
337                dropped_ratelimited = missed_events,
338                $($rest)*
339            );
340        }
341    }};
342    // Default case (only level provided)
343    (level: $level:ident, $($rest:tt)*) => {{
344        static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default();
345        if let Ok(missed_events) = RATE_LIMITER.event() {
346            $crate::tracing::event!(
347                $crate::tracing::Level::$level,
348                dropped_ratelimited = missed_events,
349                $($rest)*
350            );
351        }
352    }};
353}
354
355/// Helper macro for dynamically dispatching to [`event_ratelimited_static!`] based on a runtime level.
356///
357/// This macro accepts a runtime `tracing::Level` expression and dispatches to the appropriate
358/// compile-time level identifier. Allows the log level to be determined at runtime.
359///
360/// Examples:
361/// ```
362/// use tracing::Level;
363/// use tracelimit::event_ratelimited;
364/// event_ratelimited!(Level::ERROR, period: 1000, limit: 5, "custome period and limit");
365/// event_ratelimited!(Level::WARN, period: 10000, "custom period only");
366/// event_ratelimited!(Level::INFO, limit: 50, "custom limit only");
367/// event_ratelimited!(Level::TRACE, "simple message");
368/// ```
369#[macro_export]
370macro_rules! event_ratelimited {
371    // With period and limit and level
372    ($level:expr, period: $period:expr, limit: $limit:expr, $($rest:tt)*) => {
373        match $level {
374            $crate::tracing::Level::ERROR => {
375                $crate::event_ratelimited_static!(level: ERROR, period: $period, limit: $limit, $($rest)*);
376            }
377            $crate::tracing::Level::WARN => {
378                $crate::event_ratelimited_static!(level: WARN, period: $period, limit: $limit, $($rest)*);
379            }
380            $crate::tracing::Level::INFO => {
381                $crate::event_ratelimited_static!(level: INFO, period: $period, limit: $limit, $($rest)*);
382            }
383            $crate::tracing::Level::DEBUG => {
384                $crate::event_ratelimited_static!(level: DEBUG, period: $period, limit: $limit, $($rest)*);
385            }
386            $crate::tracing::Level::TRACE => {
387                $crate::event_ratelimited_static!(level: TRACE, period: $period, limit: $limit, $($rest)*);
388            }
389        }
390    };
391    // With period and level
392    ($level:expr, period: $period:expr, $($rest:tt)*) => {
393        match $level {
394            $crate::tracing::Level::ERROR => {
395                $crate::event_ratelimited_static!(level: ERROR, period: $period, $($rest)*);
396            }
397            $crate::tracing::Level::WARN => {
398                $crate::event_ratelimited_static!(level: WARN, period: $period, $($rest)*);
399            }
400            $crate::tracing::Level::INFO => {
401                $crate::event_ratelimited_static!(level: INFO, period: $period, $($rest)*);
402            }
403            $crate::tracing::Level::DEBUG => {
404                $crate::event_ratelimited_static!(level: DEBUG, period: $period, $($rest)*);
405            }
406            $crate::tracing::Level::TRACE => {
407                $crate::event_ratelimited_static!(level: TRACE, period: $period, $($rest)*);
408            }
409        }
410    };
411    // With limit and level
412    ($level:expr, limit: $limit:expr, $($rest:tt)*) => {
413        match $level {
414            $crate::tracing::Level::ERROR => {
415                $crate::event_ratelimited_static!(level: ERROR, limit: $limit, $($rest)*);
416            }
417            $crate::tracing::Level::WARN => {
418                $crate::event_ratelimited_static!(level: WARN, limit: $limit, $($rest)*);
419            }
420            $crate::tracing::Level::INFO => {
421                $crate::event_ratelimited_static!(level: INFO, limit: $limit, $($rest)*);
422            }
423            $crate::tracing::Level::DEBUG => {
424                $crate::event_ratelimited_static!(level: DEBUG, limit: $limit, $($rest)*);
425            }
426            $crate::tracing::Level::TRACE => {
427                $crate::event_ratelimited_static!(level: TRACE, limit: $limit, $($rest)*);
428            }
429        }
430    };
431    // Default case (only level provided)
432    ($level:expr, $($rest:tt)*) => {
433        match $level {
434            $crate::tracing::Level::ERROR => {
435                $crate::event_ratelimited_static!(level: ERROR, $($rest)*);
436            }
437            $crate::tracing::Level::WARN => {
438                $crate::event_ratelimited_static!(level: WARN, $($rest)*);
439            }
440            $crate::tracing::Level::INFO => {
441                $crate::event_ratelimited_static!(level: INFO, $($rest)*);
442            }
443            $crate::tracing::Level::DEBUG => {
444                $crate::event_ratelimited_static!(level: DEBUG, $($rest)*);
445            }
446            $crate::tracing::Level::TRACE => {
447                $crate::event_ratelimited_static!(level: TRACE, $($rest)*);
448            }
449        }
450    };
451}
tracelimit/lib.rs

tracelimit/
lib.rs
