Skip to main content

tracing_limit/
lib.rs

1#![deny(warnings)]
2#![deny(clippy::unwrap_used)]
3//! Rate limiting for tracing events.
4//!
5//! This crate provides a tracing-subscriber layer that rate limits log events to prevent
6//! log flooding. Events are grouped by their callsite and contextual fields, with each
7//! unique combination rate limited independently.
8//!
9//! # How it works
10//!
11//! Within each rate limit window (default 10 seconds):
12//! - **1st occurrence**: Event is emitted normally
13//! - **2nd occurrence**: Emits a "suppressing" warning
14//! - **3rd+ occurrences**: Silent until window expires
15//! - **After window**: Emits a summary of suppressed count, then next event normally
16//!
17//! Note: the suppressed-count summary and the resumption of normal emission are both
18//! triggered by the *next arriving event* after the window has elapsed, not by the
19//! window expiry itself. If the event stops firing, no summary is ever emitted.
20//!
21//! # Rate limit grouping
22//!
23//! Events are rate limited independently based on a combination of:
24//! - **Callsite**: The code location where the log statement appears
25//! - **Contextual fields**: Any fields attached to the event or its parent spans
26//!
27//! ## How fields contribute to grouping
28//!
29//! **Only these fields create distinct rate limit groups:**
30//! - `component_id` - Different components are rate limited independently
31//!
32//! **All other fields are ignored for grouping**, including:
33//! - `fanout_id`, `input_id`, `output_id` - Not used for grouping to avoid resource/cost implications from high-cardinality tags
34//! - `message` - The log message itself doesn't differentiate groups
35//! - `internal_log_rate_limit` - Control field for enabling/disabling rate limiting
36//! - `internal_log_rate_secs` - Control field for customizing the rate limit window
37//! - Any custom fields you add
38//!
39//! ## Examples
40//!
41//! ```rust,ignore
42//! // Example 1: Different component_id values create separate rate limit groups
43//! info!(component_id = "transform_1", "Processing event");  // Group A
44//! info!(component_id = "transform_2", "Processing event");  // Group B
45//! // Even though the message is identical, these are rate limited independently
46//!
47//! // Example 2: Only component_id matters for grouping
48//! info!(component_id = "router", fanout_id = "output_1", "Routing event");  // Group C
49//! info!(component_id = "router", fanout_id = "output_2", "Routing event");  // Group C (same group!)
50//! info!(component_id = "router", fanout_id = "output_1", "Routing event");  // Group C (same group!)
51//! info!(component_id = "router", fanout_id = "output_1", input_id = "kafka", "Routing event");  // Group C (same!)
52//! // All of these share the same group because they have the same component_id
53//! // The fanout_id and input_id fields are ignored to avoid resource/cost implications
54//!
55//! // Example 3: Span fields contribute to grouping
56//! let span = info_span!("process", component_id = "transform_1");
57//! let _enter = span.enter();
58//! info!("Processing event");  // Group E: callsite + component_id from span
59//! drop(_enter);
60//!
61//! let span = info_span!("process", component_id = "transform_2");
62//! let _enter = span.enter();
63//! info!("Processing event");  // Group F: same callsite but different component_id
64//!
65//! // Example 4: Nested spans - child span fields take precedence
66//! let outer = info_span!("outer", component_id = "parent");
67//! let _outer_guard = outer.enter();
68//! let inner = info_span!("inner", component_id = "child");
69//! let _inner_guard = inner.enter();
70//! info!("Nested event");  // Grouped by component_id = "child"
71//!
72//! // Example 5: Same callsite with no fields = single rate limit group
73//! info!("Simple message");  // Group G
74//! info!("Simple message");  // Group G
75//! info!("Simple message");  // Group G
76//!
77//! // Example 6: Custom fields are ignored for grouping
78//! info!(component_id = "source", input_id = "in_1", "Received data");  // Group H
79//! info!(component_id = "source", input_id = "in_2", "Received data");  // Group H (same group!)
80//! // The input_id field is ignored - only component_id matters
81//!
82//! // Example 7: Disabling rate limiting for specific logs
83//! // Rate limiting is ON by default - explicitly disable for important logs
84//! warn!(
85//!     component_id = "critical_component",
86//!     message = "Fatal error occurred",
87//!     internal_log_rate_limit = false
88//! );
89//! // This event will NEVER be rate limited, regardless of how often it fires
90//!
91//! // Example 8: Custom rate limit window for specific events
92//! info!(
93//!     component_id = "noisy_component",
94//!     message = "Frequent status update",
95//!     internal_log_rate_secs = 60  // Only log once per minute
96//! );
97//! // Override the default window for this specific log
98//! ```
99//!
100//! This ensures logs from different components are rate limited independently,
101//! while avoiding resource/cost implications from high-cardinality tags.
102
103use std::fmt;
104
105use dashmap::DashMap;
106use tracing_core::{
107    Event, Metadata, Subscriber,
108    callsite::Identifier,
109    field::{Field, Value, Visit, display},
110    span,
111    subscriber::Interest,
112};
113use tracing_subscriber::layer::{Context, Layer};
114
115#[cfg(test)]
116#[macro_use]
117extern crate tracing;
118
119#[cfg(not(test))]
120use std::time::Instant;
121
122#[cfg(test)]
123use mock_instant::global::Instant;
124
125const RATE_LIMIT_FIELD: &str = "internal_log_rate_limit";
126const RATE_LIMIT_SECS_FIELD: &str = "internal_log_rate_secs";
127const MESSAGE_FIELD: &str = "message";
128
129// These fields will cause events to be independently rate limited by the values
130// for these keys
131const COMPONENT_ID_FIELD: &str = "component_id";
132
133#[derive(Eq, PartialEq, Hash, Clone)]
134struct RateKeyIdentifier {
135    callsite: Identifier,
136    rate_limit_key_values: RateLimitedSpanKeys,
137}
138
139pub struct RateLimitedLayer<S, L>
140where
141    L: Layer<S> + Sized,
142    S: Subscriber,
143{
144    events: DashMap<RateKeyIdentifier, State>,
145    inner: L,
146    internal_log_rate_limit: u64,
147    _subscriber: std::marker::PhantomData<S>,
148}
149
150impl<S, L> RateLimitedLayer<S, L>
151where
152    L: Layer<S> + Sized,
153    S: Subscriber,
154{
155    pub fn new(layer: L) -> Self {
156        RateLimitedLayer {
157            events: Default::default(),
158            internal_log_rate_limit: 10,
159            inner: layer,
160            _subscriber: std::marker::PhantomData,
161        }
162    }
163
164    /// Sets the default rate limit window in seconds.
165    ///
166    /// This controls how long logs are suppressed before they can be emitted again.
167    /// Within each window:
168    /// - 1st occurrence: Emitted normally
169    /// - 2nd occurrence: Shows "suppressing" warning
170    /// - 3rd+ occurrences: Silent until window expires
171    /// - After window: Summary and next event emitted on next arrival (see module-level note)
172    pub fn with_default_limit(mut self, internal_log_rate_limit: u64) -> Self {
173        self.internal_log_rate_limit = internal_log_rate_limit;
174        self
175    }
176}
177
178impl<S, L> Layer<S> for RateLimitedLayer<S, L>
179where
180    L: Layer<S>,
181    S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
182{
183    #[inline]
184    fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
185        self.inner.register_callsite(metadata)
186    }
187
188    #[inline]
189    fn enabled(&self, metadata: &Metadata<'_>, ctx: Context<'_, S>) -> bool {
190        self.inner.enabled(metadata, ctx)
191    }
192
193    // keep track of any span fields we use for grouping rate limiting
194    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
195        {
196            let span = ctx.span(id).expect("Span not found, this is a bug");
197            let mut extensions = span.extensions_mut();
198
199            if extensions.get_mut::<RateLimitedSpanKeys>().is_none() {
200                let mut fields = RateLimitedSpanKeys::default();
201                attrs.record(&mut fields);
202                extensions.insert(fields);
203            };
204        }
205        self.inner.on_new_span(attrs, id, ctx);
206    }
207
208    // keep track of any span fields we use for grouping rate limiting
209    fn on_record(&self, id: &span::Id, values: &span::Record<'_>, ctx: Context<'_, S>) {
210        {
211            let span = ctx.span(id).expect("Span not found, this is a bug");
212            let mut extensions = span.extensions_mut();
213
214            match extensions.get_mut::<RateLimitedSpanKeys>() {
215                Some(fields) => {
216                    values.record(fields);
217                }
218                None => {
219                    let mut fields = RateLimitedSpanKeys::default();
220                    values.record(&mut fields);
221                    extensions.insert(fields);
222                }
223            };
224        }
225        self.inner.on_record(id, values, ctx);
226    }
227
228    #[inline]
229    fn on_follows_from(&self, span: &span::Id, follows: &span::Id, ctx: Context<'_, S>) {
230        self.inner.on_follows_from(span, follows, ctx);
231    }
232
233    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
234        // Visit the event, grabbing the limit status if one is defined. Rate limiting is ON by default
235        // unless explicitly disabled by setting `internal_log_rate_limit = false`.
236        let mut limit_visitor = LimitVisitor::default();
237        event.record(&mut limit_visitor);
238
239        let limit_exists = limit_visitor.limit.unwrap_or(true);
240        if !limit_exists {
241            return self.inner.on_event(event, ctx);
242        }
243
244        let limit = match limit_visitor.limit_secs {
245            Some(limit_secs) => limit_secs, // override the cli limit
246            None => self.internal_log_rate_limit,
247        };
248
249        // Build a composite key from event fields and span context to determine the rate limit group.
250        // This multi-step process ensures we capture all relevant contextual information:
251        //
252        // 1. Start with event-level fields (e.g., fields directly on the log macro call)
253        // 2. Walk up the span hierarchy from root to current span
254        // 3. Merge in fields from each span, with child spans taking precedence
255        //
256        // This means an event's rate limit group is determined by the combination of:
257        // - Its callsite (handled separately via RateKeyIdentifier)
258        // - All contextual fields from both the event and its span ancestry
259        //
260        // Example: The same `info!("msg")` callsite in different component contexts becomes
261        // distinct rate limit groups, allowing fine-grained control over log flooding.
262        let rate_limit_key_values = {
263            let mut keys = RateLimitedSpanKeys::default();
264            // Capture fields directly on this event
265            event.record(&mut keys);
266
267            // Walk span hierarchy and merge in contextual fields
268            ctx.lookup_current()
269                .into_iter()
270                .flat_map(|span| span.scope().from_root())
271                .fold(keys, |mut keys, span| {
272                    let extensions = span.extensions();
273                    if let Some(span_keys) = extensions.get::<RateLimitedSpanKeys>() {
274                        keys.merge(span_keys);
275                    }
276                    keys
277                })
278        };
279
280        // Build the key to represent this event, given its span fields, and see if we're already rate limiting it. If
281        // not, we'll initialize an entry for it.
282        let metadata = event.metadata();
283        let id = RateKeyIdentifier {
284            callsite: metadata.callsite(),
285            rate_limit_key_values,
286        };
287
288        let mut state = self.events.entry(id).or_insert_with(|| {
289            let mut message_visitor = MessageVisitor::default();
290            event.record(&mut message_visitor);
291
292            let message = message_visitor
293                .message
294                .unwrap_or_else(|| metadata.name().into());
295
296            State::new(message, limit)
297        });
298
299        // Update our suppressed state for this event, and see if we should still be suppressing it.
300        //
301        // When this is the first time seeing the event, we emit it like we normally would. The second time we see it in
302        // the limit period, we emit a new event to indicate that the original event is being actively suppressed.
303        // Otherwise, we don't emit anything.
304        let previous_count = state.increment_count();
305        if state.should_limit() {
306            match previous_count {
307                0 => self.inner.on_event(event, ctx),
308                1 => {
309                    let message = format!(
310                        "Internal log [{}] is being suppressed to avoid flooding.",
311                        state.message
312                    );
313                    self.create_event(&ctx, metadata, message, state.limit);
314                }
315                _ => {}
316            }
317        } else {
318            // If we saw this event 3 or more times total, emit an event that indicates the total number of times we
319            // suppressed the event in the limit period.
320            if previous_count > 1 {
321                let message = format!(
322                    "Internal log [{}] has been suppressed {} times.",
323                    state.message,
324                    previous_count - 1
325                );
326
327                self.create_event(&ctx, metadata, message, state.limit);
328            }
329
330            // We're not suppressing anymore, so we also emit the current event as normal.. but we update our rate
331            // limiting state since this is effectively equivalent to seeing the event again for the first time.
332            self.inner.on_event(event, ctx);
333
334            state.reset();
335        }
336    }
337
338    #[inline]
339    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
340        self.inner.on_enter(id, ctx);
341    }
342
343    #[inline]
344    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
345        self.inner.on_exit(id, ctx);
346    }
347
348    #[inline]
349    fn on_close(&self, id: span::Id, ctx: Context<'_, S>) {
350        self.inner.on_close(id, ctx);
351    }
352
353    #[inline]
354    fn on_id_change(&self, old: &span::Id, new: &span::Id, ctx: Context<'_, S>) {
355        self.inner.on_id_change(old, new, ctx);
356    }
357
358    #[inline]
359    fn on_layer(&mut self, subscriber: &mut S) {
360        self.inner.on_layer(subscriber);
361    }
362}
363
364impl<S, L> RateLimitedLayer<S, L>
365where
366    S: Subscriber,
367    L: Layer<S>,
368{
369    fn create_event(
370        &self,
371        ctx: &Context<S>,
372        metadata: &'static Metadata<'static>,
373        message: String,
374        rate_limit: u64,
375    ) {
376        let fields = metadata.fields();
377
378        let message = display(message);
379
380        if let Some(message_field) = fields.field("message") {
381            let values = [(&message_field, Some(&message as &dyn Value))];
382
383            let valueset = fields.value_set(&values);
384            let event = Event::new(metadata, &valueset);
385            self.inner.on_event(&event, ctx.clone());
386        } else if let Some(rate_limit_field) = fields.field(RATE_LIMIT_FIELD) {
387            let values = [(&rate_limit_field, Some(&rate_limit as &dyn Value))];
388
389            let valueset = fields.value_set(&values);
390            let event = Event::new(metadata, &valueset);
391            self.inner.on_event(&event, ctx.clone());
392        } else {
393            // If the event metadata has neither a "message" nor "internal_log_rate_limit" field,
394            // we cannot create a proper synthetic event. This can happen with custom debug events
395            // that have their own field structure. In this case, we simply skip emitting the
396            // rate limit notification rather than panicking.
397        }
398    }
399}
400
401#[derive(Debug)]
402struct State {
403    start: Instant,
404    count: u64,
405    limit: u64,
406    message: String,
407}
408
409impl State {
410    fn new(message: String, limit: u64) -> Self {
411        Self {
412            start: Instant::now(),
413            count: 0,
414            limit,
415            message,
416        }
417    }
418
419    fn reset(&mut self) {
420        self.start = Instant::now();
421        self.count = 1;
422    }
423
424    fn increment_count(&mut self) -> u64 {
425        let prev = self.count;
426        self.count += 1;
427        prev
428    }
429
430    fn should_limit(&self) -> bool {
431        self.start.elapsed().as_secs() < self.limit
432    }
433}
434
435#[derive(PartialEq, Eq, Clone, Hash)]
436enum TraceValue {
437    String(String),
438    Int(i64),
439    Uint(u64),
440    Bool(bool),
441}
442
443#[cfg(test)]
444impl fmt::Display for TraceValue {
445    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
446        match self {
447            TraceValue::String(s) => write!(f, "{}", s),
448            TraceValue::Int(i) => write!(f, "{}", i),
449            TraceValue::Uint(u) => write!(f, "{}", u),
450            TraceValue::Bool(b) => write!(f, "{}", b),
451        }
452    }
453}
454
455impl From<bool> for TraceValue {
456    fn from(b: bool) -> Self {
457        TraceValue::Bool(b)
458    }
459}
460
461impl From<i64> for TraceValue {
462    fn from(i: i64) -> Self {
463        TraceValue::Int(i)
464    }
465}
466
467impl From<u64> for TraceValue {
468    fn from(u: u64) -> Self {
469        TraceValue::Uint(u)
470    }
471}
472
473impl From<String> for TraceValue {
474    fn from(s: String) -> Self {
475        TraceValue::String(s)
476    }
477}
478
479/// RateLimitedSpanKeys records span and event fields that differentiate rate limit groups.
480///
481/// This struct is used to build a composite key that uniquely identifies a rate limit bucket.
482/// Events with different field values will be rate limited independently, even if they come
483/// from the same callsite.
484///
485/// ## Field categories:
486///
487/// **Tracked fields** (only these create distinct rate limit groups):
488/// - `component_id` - Different components are rate limited independently
489///
490/// **Ignored fields**: All other fields are ignored for grouping purposes. This avoids resource/cost implications from high-cardinality tags.
491/// ```
492#[derive(Default, Eq, PartialEq, Hash, Clone)]
493struct RateLimitedSpanKeys {
494    component_id: Option<TraceValue>,
495}
496
497impl RateLimitedSpanKeys {
498    fn record(&mut self, field: &Field, value: TraceValue) {
499        if field.name() == COMPONENT_ID_FIELD {
500            self.component_id = Some(value);
501        }
502    }
503
504    fn merge(&mut self, other: &Self) {
505        if let Some(component_id) = &other.component_id {
506            self.component_id = Some(component_id.clone());
507        }
508    }
509}
510
511impl Visit for RateLimitedSpanKeys {
512    fn record_i64(&mut self, field: &Field, value: i64) {
513        self.record(field, value.into());
514    }
515
516    fn record_u64(&mut self, field: &Field, value: u64) {
517        self.record(field, value.into());
518    }
519
520    fn record_bool(&mut self, field: &Field, value: bool) {
521        self.record(field, value.into());
522    }
523
524    fn record_str(&mut self, field: &Field, value: &str) {
525        self.record(field, value.to_owned().into());
526    }
527
528    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
529        self.record(field, format!("{value:?}").into());
530    }
531}
532
533#[derive(Default)]
534struct LimitVisitor {
535    pub limit: Option<bool>,
536    pub limit_secs: Option<u64>,
537}
538
539impl Visit for LimitVisitor {
540    fn record_bool(&mut self, field: &Field, value: bool) {
541        if field.name() == RATE_LIMIT_FIELD {
542            self.limit = Some(value);
543        }
544    }
545
546    fn record_i64(&mut self, field: &Field, value: i64) {
547        if field.name() == RATE_LIMIT_SECS_FIELD {
548            self.limit = Some(true); // limit if we have this field
549            self.limit_secs = Some(u64::try_from(value).unwrap_or_default()); // override the cli passed limit
550        }
551    }
552
553    fn record_u64(&mut self, field: &Field, value: u64) {
554        if field.name() == RATE_LIMIT_SECS_FIELD {
555            self.limit = Some(true); // limit if we have this field
556            self.limit_secs = Some(value); // override the cli passed limit
557        }
558    }
559
560    fn record_debug(&mut self, _field: &Field, _value: &dyn fmt::Debug) {}
561}
562
563#[derive(Default)]
564struct MessageVisitor {
565    pub message: Option<String>,
566}
567
568impl Visit for MessageVisitor {
569    fn record_str(&mut self, field: &Field, value: &str) {
570        if self.message.is_none() && field.name() == MESSAGE_FIELD {
571            self.message = Some(value.to_string());
572        }
573    }
574
575    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
576        if self.message.is_none() && field.name() == MESSAGE_FIELD {
577            self.message = Some(format!("{value:?}"));
578        }
579    }
580}
581
582#[cfg(test)]
583mod test {
584    use std::{
585        collections::BTreeMap,
586        sync::{Arc, Mutex},
587        time::Duration,
588    };
589
590    use mock_instant::global::MockClock;
591    use serial_test::serial;
592    use tracing_subscriber::layer::SubscriberExt;
593
594    use super::*;
595
596    #[derive(Debug, Clone, PartialEq, Eq)]
597    struct RecordedEvent {
598        message: String,
599        fields: BTreeMap<String, String>,
600    }
601
602    impl RecordedEvent {
603        fn new(message: impl Into<String>) -> Self {
604            Self {
605                message: message.into(),
606                fields: BTreeMap::new(),
607            }
608        }
609
610        fn with_field(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
611            self.fields.insert(key.into(), value.into());
612            self
613        }
614    }
615
616    /// Macro to create RecordedEvent with optional fields
617    /// Usage:
618    /// - `event!("message")` - just message
619    /// - `event!("message", key1: "value1")` - message with one field
620    /// - `event!("message", key1: "value1", key2: "value2")` - message with multiple fields
621    macro_rules! event {
622        ($msg:expr) => {
623            RecordedEvent::new($msg)
624        };
625        ($msg:expr, $($key:ident: $value:expr),+ $(,)?) => {
626            RecordedEvent::new($msg)
627                $(.with_field(stringify!($key), $value))+
628        };
629    }
630
631    #[derive(Default)]
632    struct AllFieldsVisitor {
633        fields: BTreeMap<String, String>,
634    }
635
636    impl Visit for AllFieldsVisitor {
637        fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
638            self.fields
639                .insert(field.name().to_string(), format!("{value:?}"));
640        }
641
642        fn record_str(&mut self, field: &Field, value: &str) {
643            self.fields
644                .insert(field.name().to_string(), value.to_string());
645        }
646
647        fn record_i64(&mut self, field: &Field, value: i64) {
648            self.fields
649                .insert(field.name().to_string(), value.to_string());
650        }
651
652        fn record_u64(&mut self, field: &Field, value: u64) {
653            self.fields
654                .insert(field.name().to_string(), value.to_string());
655        }
656
657        fn record_bool(&mut self, field: &Field, value: bool) {
658            self.fields
659                .insert(field.name().to_string(), value.to_string());
660        }
661    }
662
663    impl AllFieldsVisitor {
664        fn into_event(self) -> RecordedEvent {
665            let message = self
666                .fields
667                .get("message")
668                .cloned()
669                .unwrap_or_else(|| String::from(""));
670
671            let mut fields = BTreeMap::new();
672            for (key, value) in self.fields {
673                if key != "message"
674                    && key != "internal_log_rate_limit"
675                    && key != "internal_log_rate_secs"
676                {
677                    fields.insert(key, value);
678                }
679            }
680
681            RecordedEvent { message, fields }
682        }
683    }
684
685    #[derive(Default)]
686    struct RecordingLayer<S> {
687        events: Arc<Mutex<Vec<RecordedEvent>>>,
688
689        _subscriber: std::marker::PhantomData<S>,
690    }
691
692    impl<S> RecordingLayer<S> {
693        fn new(events: Arc<Mutex<Vec<RecordedEvent>>>) -> Self {
694            RecordingLayer {
695                events,
696
697                _subscriber: std::marker::PhantomData,
698            }
699        }
700    }
701
702    impl<S> Layer<S> for RecordingLayer<S>
703    where
704        S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
705    {
706        fn register_callsite(&self, _metadata: &'static Metadata<'static>) -> Interest {
707            Interest::always()
708        }
709
710        fn enabled(&self, _metadata: &Metadata<'_>, _ctx: Context<'_, S>) -> bool {
711            true
712        }
713
714        fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
715            let mut visitor = AllFieldsVisitor::default();
716            event.record(&mut visitor);
717
718            // Also capture fields from span context
719            if let Some(span) = ctx.lookup_current() {
720                for span_ref in span.scope().from_root() {
721                    let extensions = span_ref.extensions();
722                    if let Some(span_keys) = extensions.get::<RateLimitedSpanKeys>() {
723                        // Add component_id
724                        if let Some(TraceValue::String(ref s)) = span_keys.component_id {
725                            visitor.fields.insert("component_id".to_string(), s.clone());
726                        }
727                    }
728                }
729            }
730
731            let mut events = self.events.lock().unwrap();
732            events.push(visitor.into_event());
733        }
734    }
735
736    /// Helper function to set up a test with a rate-limited subscriber.
737    /// Returns the events Arc for asserting on collected events.
738    fn setup_test(
739        default_limit: u64,
740    ) -> (
741        Arc<Mutex<Vec<RecordedEvent>>>,
742        impl Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
743    ) {
744        let events: Arc<Mutex<Vec<RecordedEvent>>> = Default::default();
745        let recorder = RecordingLayer::new(Arc::clone(&events));
746        let sub = tracing_subscriber::registry::Registry::default()
747            .with(RateLimitedLayer::new(recorder).with_default_limit(default_limit));
748        (events, sub)
749    }
750
751    #[test]
752    #[serial]
753    fn rate_limits() {
754        let (events, sub) = setup_test(1);
755        tracing::subscriber::with_default(sub, || {
756            for _ in 0..21 {
757                info!(message = "Hello world!");
758                MockClock::advance(Duration::from_millis(100));
759            }
760        });
761
762        let events = events.lock().unwrap();
763
764        assert_eq!(
765            *events,
766            vec![
767                event!("Hello world!"),
768                event!("Internal log [Hello world!] is being suppressed to avoid flooding."),
769                event!("Internal log [Hello world!] has been suppressed 9 times."),
770                event!("Hello world!"),
771                event!("Internal log [Hello world!] is being suppressed to avoid flooding."),
772                event!("Internal log [Hello world!] has been suppressed 9 times."),
773                event!("Hello world!"),
774            ]
775        );
776    }
777
778    #[test]
779    #[serial]
780    fn override_rate_limit_at_callsite() {
781        let (events, sub) = setup_test(100);
782        tracing::subscriber::with_default(sub, || {
783            for _ in 0..31 {
784                info!(message = "Hello world!", internal_log_rate_secs = 2);
785                MockClock::advance(Duration::from_millis(100));
786            }
787        });
788
789        let events = events.lock().unwrap();
790
791        // With a 2-second window and 100ms advances, we get:
792        // - Event every 20 iterations (2000ms / 100ms = 20)
793        // - First window: iteration 0-19 (suppressed 19 times after first 2)
794        // - Second window: iteration 20-39 (but we only go to 30)
795        assert_eq!(
796            *events,
797            vec![
798                event!("Hello world!"),
799                event!("Internal log [Hello world!] is being suppressed to avoid flooding."),
800                event!("Internal log [Hello world!] has been suppressed 19 times."),
801                event!("Hello world!"),
802                event!("Internal log [Hello world!] is being suppressed to avoid flooding."),
803            ]
804        );
805    }
806
807    #[test]
808    #[serial]
809    fn rate_limit_by_event_key() {
810        let (events, sub) = setup_test(1);
811        tracing::subscriber::with_default(sub, || {
812            for _ in 0..21 {
813                for key in &["foo", "bar"] {
814                    info!(
815                        message = format!("Hello {key}!").as_str(),
816                        component_id = &key
817                    );
818                }
819                MockClock::advance(Duration::from_millis(100));
820            }
821        });
822
823        let events = events.lock().unwrap();
824
825        // Events with different component_id values create separate rate limit groups
826        assert_eq!(
827            *events,
828            vec![
829                event!("Hello foo!", component_id: "foo"),
830                event!("Hello bar!", component_id: "bar"),
831                event!("Internal log [Hello foo!] is being suppressed to avoid flooding."),
832                event!("Internal log [Hello bar!] is being suppressed to avoid flooding."),
833                event!("Internal log [Hello foo!] has been suppressed 9 times."),
834                event!("Hello foo!", component_id: "foo"),
835                event!("Internal log [Hello bar!] has been suppressed 9 times."),
836                event!("Hello bar!", component_id: "bar"),
837                event!("Internal log [Hello foo!] is being suppressed to avoid flooding."),
838                event!("Internal log [Hello bar!] is being suppressed to avoid flooding."),
839                event!("Internal log [Hello foo!] has been suppressed 9 times."),
840                event!("Hello foo!", component_id: "foo"),
841                event!("Internal log [Hello bar!] has been suppressed 9 times."),
842                event!("Hello bar!", component_id: "bar"),
843            ]
844        );
845    }
846
847    #[test]
848    #[serial]
849    fn disabled_rate_limit() {
850        let (events, sub) = setup_test(1);
851        tracing::subscriber::with_default(sub, || {
852            for _ in 0..21 {
853                info!(message = "Hello world!", internal_log_rate_limit = false);
854                MockClock::advance(Duration::from_millis(100));
855            }
856        });
857
858        let events = events.lock().unwrap();
859
860        // All 21 events should be emitted since rate limiting is disabled
861        assert_eq!(events.len(), 21);
862        assert!(events.iter().all(|e| e == &event!("Hello world!")));
863    }
864
865    #[test]
866    #[serial]
867    fn rate_limit_ignores_non_special_fields() {
868        let (events, sub) = setup_test(1);
869        tracing::subscriber::with_default(sub, || {
870            for i in 0..21 {
871                // Call the SAME info! macro multiple times per iteration with varying fanout_id
872                // to verify that fanout_id doesn't create separate rate limit groups
873                for _ in 0..3 {
874                    let fanout = if i % 2 == 0 { "output_1" } else { "output_2" };
875                    info!(
876                        message = "Routing event",
877                        component_id = "router",
878                        fanout_id = fanout
879                    );
880                }
881                MockClock::advance(Duration::from_millis(100));
882            }
883        });
884
885        let events = events.lock().unwrap();
886
887        // All events share the same rate limit group (same callsite + component_id)
888        // First event emits normally, second shows suppression, third and beyond are silent
889        // until the window expires
890        assert_eq!(
891            *events,
892            vec![
893                // First iteration - first emits, second shows suppression, 3rd+ silent
894                event!("Routing event", component_id: "router", fanout_id: "output_1"),
895                event!("Internal log [Routing event] is being suppressed to avoid flooding."),
896                // After rate limit window (1 sec) - summary shows suppressions
897                event!("Internal log [Routing event] has been suppressed 29 times."),
898                event!("Routing event", component_id: "router", fanout_id: "output_1"),
899                event!("Internal log [Routing event] is being suppressed to avoid flooding."),
900                event!("Internal log [Routing event] has been suppressed 29 times."),
901                event!("Routing event", component_id: "router", fanout_id: "output_1"),
902                event!("Internal log [Routing event] is being suppressed to avoid flooding."),
903            ]
904        );
905    }
906
907    #[test]
908    #[serial]
909    fn nested_spans_child_takes_precedence() {
910        let (events, sub) = setup_test(1);
911        tracing::subscriber::with_default(sub, || {
912            // Create nested spans where child overrides parent's component_id
913            let outer = info_span!("outer", component_id = "parent");
914            let _outer_guard = outer.enter();
915
916            for _ in 0..21 {
917                // Inner span with different component_id should take precedence
918                let inner = info_span!("inner", component_id = "child");
919                let _inner_guard = inner.enter();
920                info!(message = "Nested event");
921                drop(_inner_guard);
922
923                MockClock::advance(Duration::from_millis(100));
924            }
925        });
926
927        let events = events.lock().unwrap();
928
929        // All events should be grouped by component_id = "child" (from inner span)
930        // not "parent" (from outer span), demonstrating child precedence
931        assert_eq!(
932            *events,
933            vec![
934                event!("Nested event", component_id: "child"),
935                event!("Internal log [Nested event] is being suppressed to avoid flooding.", component_id: "child"),
936                event!("Internal log [Nested event] has been suppressed 9 times.", component_id: "child"),
937                event!("Nested event", component_id: "child"),
938                event!("Internal log [Nested event] is being suppressed to avoid flooding.", component_id: "child"),
939                event!("Internal log [Nested event] has been suppressed 9 times.", component_id: "child"),
940                event!("Nested event", component_id: "child"),
941            ]
942        );
943    }
944
945    #[test]
946    #[serial]
947    fn nested_spans_ignores_untracked_fields() {
948        let (events, sub) = setup_test(1);
949        tracing::subscriber::with_default(sub, || {
950            // Parent has component_id, child has some_field - only component_id is tracked
951            let outer = info_span!("outer", component_id = "transform");
952            let _outer_guard = outer.enter();
953
954            for _ in 0..21 {
955                let inner = info_span!("inner", some_field = "value");
956                let _inner_guard = inner.enter();
957                info!(message = "Event message");
958                drop(_inner_guard);
959
960                MockClock::advance(Duration::from_millis(100));
961            }
962        });
963
964        let events = events.lock().unwrap();
965
966        // Events should have component_id from parent, some_field from child is ignored for grouping
967        // All events are in the same rate limit group
968        assert_eq!(
969            *events,
970            vec![
971                event!("Event message", component_id: "transform"),
972                event!(
973                    "Internal log [Event message] is being suppressed to avoid flooding.",
974                    component_id: "transform"
975                ),
976                event!(
977                    "Internal log [Event message] has been suppressed 9 times.",
978                    component_id: "transform"
979                ),
980                event!("Event message", component_id: "transform"),
981                event!(
982                    "Internal log [Event message] is being suppressed to avoid flooding.",
983                    component_id: "transform"
984                ),
985                event!(
986                    "Internal log [Event message] has been suppressed 9 times.",
987                    component_id: "transform"
988                ),
989                event!("Event message", component_id: "transform"),
990            ]
991        );
992    }
993
994    #[test]
995    #[serial]
996    fn rate_limit_same_message_different_component() {
997        let (events, sub) = setup_test(1);
998        tracing::subscriber::with_default(sub, || {
999            // Use a loop with the SAME callsite to demonstrate that identical messages
1000            // with different component_ids create separate rate limit groups
1001            for component in &["foo", "foo", "bar"] {
1002                info!(message = "Hello!", component_id = component);
1003                MockClock::advance(Duration::from_millis(100));
1004            }
1005        });
1006
1007        let events = events.lock().unwrap();
1008
1009        // The first "foo" event is emitted normally (count=0)
1010        // The second "foo" event triggers suppression warning (count=1)
1011        // The "bar" event is emitted normally (count=0 for its group)
1012        // This proves that even with identical message text, different component_ids
1013        // create separate rate limit groups
1014        assert_eq!(
1015            *events,
1016            vec![
1017                event!("Hello!", component_id: "foo"),
1018                event!("Internal log [Hello!] is being suppressed to avoid flooding."),
1019                event!("Hello!", component_id: "bar"),
1020            ]
1021        );
1022    }
1023
1024    #[test]
1025    #[serial]
1026    fn events_with_custom_fields_no_message_dont_panic() {
1027        // Verify events without "message" or "internal_log_rate_limit" fields don't panic
1028        // when rate limiting skips suppression notifications.
1029        let (events, sub) = setup_test(1);
1030        tracing::subscriber::with_default(sub, || {
1031            // Use closure to ensure all events share the same callsite
1032            let emit_event = || {
1033                debug!(component_id = "test_component", utilization = 0.85);
1034            };
1035
1036            // First window: emit 5 events, only the first one should be logged
1037            for _ in 0..5 {
1038                emit_event();
1039                MockClock::advance(Duration::from_millis(100));
1040            }
1041
1042            // Advance to the next window
1043            MockClock::advance(Duration::from_millis(1000));
1044
1045            // Second window: this event should be logged
1046            emit_event();
1047        });
1048
1049        let events = events.lock().unwrap();
1050
1051        // First event from window 1, first event from window 2
1052        // Suppression notifications are skipped (no message field)
1053        assert_eq!(
1054            *events,
1055            vec![
1056                event!("", component_id: "test_component", utilization: "0.85"),
1057                event!("", component_id: "test_component", utilization: "0.85"),
1058            ]
1059        );
1060    }
1061}