tracing_span_capture/
lib.rs

1//! This crate allows testing code that should emit logs.
2//! It do that by capturing and recording logs for a given tracing span id.
3//!
4//! # Examples
5//!
6//! ```no_run
7//! use tracing::{error, span, Level};
8//! use tracing_span_capture::{RecordedLogs, TracingSpanCaptureLayer};
9//! use tracing_subscriber::layer::SubscriberExt;
10//! use tracing_subscriber::util::SubscriberInitExt;
11//!
12//! tracing_subscriber::fmt()
13//!     .finish()
14//!     .with(TracingSpanCaptureLayer)
15//!     .init();
16//!
17//! let span = span!(Level::INFO, "");
18//! let record = RecordedLogs::new(&span);
19//! {
20//!     let _enter = span.enter();
21//!     error!("try capture this");
22//! }
23//!
24//! let logs = record.into_logs();
25//! let last_log = logs.into_iter().rev().next().unwrap();
26//! assert_eq!(last_log.message, "try capture this");
27//! ```
28//!
29#![warn(missing_docs)]
30use once_cell::sync::Lazy;
31use std::collections::HashMap;
32use std::ops::DerefMut;
33use std::sync::{Arc, Mutex};
34use std::{fmt, mem};
35use tracing::field::Field;
36use tracing::{Id, Level, Span};
37use tracing_subscriber::field::Visit;
38use tracing_subscriber::{layer, Layer};
39
40type Storage = Arc<Mutex<Vec<EventLog>>>;
41
42static GLOBAL_DATA: Lazy<Mutex<HashMap<Id, Storage>>> = Lazy::new(Default::default);
43
44/// Tracing Subscriber layer which has to be registered globally
45///
46/// # Examples
47///
48/// ```no_run
49/// use tracing_subscriber::layer::SubscriberExt;
50/// use tracing_span_capture::TracingSpanCaptureLayer;
51/// use tracing_subscriber::util::SubscriberInitExt;
52///
53/// tracing_subscriber::fmt()
54///     .finish()
55///     .with(TracingSpanCaptureLayer)
56///     .init();
57/// ```
58pub struct TracingSpanCaptureLayer;
59
60impl<S> Layer<S> for TracingSpanCaptureLayer
61where
62    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
63{
64    fn on_event(&self, event: &tracing::Event<'_>, ctx: layer::Context<'_, S>) {
65        if let Some(scope) = ctx.event_scope(event) {
66            let data = GLOBAL_DATA.lock().unwrap();
67
68            for span in scope {
69                if let Some(logs) = data.get(&span.id()) {
70                    let mut fields = FieldsVisitor::default();
71                    event.record(&mut fields);
72
73                    let e = EventLog {
74                        level: *event.metadata().level(),
75                        message: fields.fields.remove("message").unwrap_or(String::new()),
76                        fields: fields.fields,
77                    };
78                    logs.lock().unwrap().push(e);
79
80                    return;
81                }
82            }
83        }
84    }
85}
86
87#[derive(Default)]
88struct FieldsVisitor {
89    fields: HashMap<&'static str, String>,
90}
91
92impl Visit for FieldsVisitor {
93    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
94        self.fields.insert(field.name(), format!("{value:?}"));
95    }
96}
97
98/// Captured log event
99#[derive(Clone, Debug)]
100pub struct EventLog {
101    /// Log Level
102    pub level: Level,
103
104    /// Emitted message from log event
105    pub message: String,
106
107    /// Emitted fields from log event
108    pub fields: HashMap<&'static str, String>,
109}
110
111/// Handler for captured logs storage
112///
113/// ```no_run
114/// use tracing::{span, Level};
115/// use tracing_span_capture::RecordedLogs;
116///
117/// let span = span!(Level::INFO, "");
118/// let logs = RecordedLogs::new(&span);
119/// {
120///     let _enter = span.enter();
121/// }
122/// let _logs_list = logs.into_logs();
123/// ```
124pub struct RecordedLogs {
125    span_id: Id,
126    logs: Storage,
127}
128
129impl RecordedLogs {
130    /// Initialize logs storage for given span id
131    ///
132    /// # Panics
133    ///
134    /// Panics If span has been either closed or was never enabled
135    pub fn new(span: &Span) -> Self {
136        let span_id = span.id().expect("span not enabled, missing id");
137        let logs: Arc<Mutex<Vec<EventLog>>> = Default::default();
138
139        GLOBAL_DATA
140            .lock()
141            .unwrap()
142            .insert(span_id.clone(), Arc::clone(&logs));
143
144        RecordedLogs { span_id, logs }
145    }
146
147    /// Take logs from storage
148    pub fn into_logs(self) -> Vec<EventLog> {
149        let mut storage = self.logs.lock().unwrap();
150        let logs = mem::take(storage.deref_mut());
151        logs
152    }
153}
154
155impl Drop for RecordedLogs {
156    fn drop(&mut self) {
157        GLOBAL_DATA.lock().unwrap().remove(&self.span_id);
158    }
159}