tracing_betterstack/
layer.rs

1use std::sync::Arc;
2use chrono::Utc;
3use tracing::{span, Event, Subscriber};
4use tracing_subscriber::{
5    layer::Context,
6    registry::LookupSpan,
7    Layer,
8};
9
10use crate::{
11    client::BetterstackClient,
12    dispatch::{BetterstackDispatcher, Dispatcher, LogEvent, NoopDispatcher},
13    export::ExportConfig,
14};
15
16#[derive(Default)]
17struct MessageVisitor(String);
18
19impl tracing::field::Visit for MessageVisitor {
20    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
21        if field.name() == "message" {
22            self.0 = value.to_string();
23        }
24    }
25
26    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
27        if field.name() == "message" {
28            self.0 = format!("{:?}", value).trim_matches('"').to_string();
29        }
30    }
31}
32
33/// A Better Stack propagation layer.
34pub struct BetterstackLayer<S> {
35    dispatcher: Arc<dyn Dispatcher>,
36    _subscriber: std::marker::PhantomData<S>,
37}
38
39/// Construct BetterstackLayer to compose with tracing subscriber.
40pub fn layer<S>() -> BetterstackLayer<S>
41where
42    S: Subscriber + for<'span> LookupSpan<'span>,
43{
44    BetterstackLayer::default()
45}
46
47impl<S> Default for BetterstackLayer<S>
48where
49    S: Subscriber + for<'span> LookupSpan<'span>,
50{
51    fn default() -> Self {
52        let dispatcher = Arc::new(NoopDispatcher::new());
53        BetterstackLayer::new(dispatcher)
54    }
55}
56
57impl<S> BetterstackLayer<S>
58where
59    S: Subscriber + for<'span> LookupSpan<'span>,
60{
61    fn new(dispatcher: Arc<dyn Dispatcher>) -> Self {
62        Self {
63            dispatcher,
64            _subscriber: std::marker::PhantomData,
65        }
66    }
67
68    /// Set the client using a source token and export configuration.
69    pub fn with_client(
70        self,
71        source_token: impl Into<String>,
72        ingestion_url: impl Into<String>,
73        export_config: ExportConfig,
74    ) -> BetterstackLayer<S>
75    where
76        BetterstackDispatcher: Dispatcher,
77    {
78        let client = BetterstackClient::new(source_token, ingestion_url);
79        let dispatcher = Arc::new(BetterstackDispatcher::new(client, export_config));
80        BetterstackLayer::new(dispatcher)
81    }
82}
83
84impl<S> Layer<S> for BetterstackLayer<S>
85where
86    S: Subscriber + for<'span> LookupSpan<'span>,
87{
88    fn on_event(&self, event: &Event<'_>, _: Context<'_, S>) {
89        // Extract message using visitor
90        let mut visitor = MessageVisitor::default();
91        event.record(&mut visitor);
92
93        // Create structured log event
94        let metadata = event.metadata();
95        let log_event = LogEvent {
96            message: visitor.0,
97            timestamp: Utc::now(),
98            level: Some(metadata.level().to_string()),
99            target: Some(metadata.target().to_string()),
100            thread_id: Some(format!(
101                "{:?}",
102                std::thread::current().id()
103            )),
104            file: metadata.file().map(String::from),
105            line: metadata.line(),
106        };
107
108        // Dispatch the event
109        self.dispatcher.dispatch(log_event);
110    }
111
112    fn on_enter(&self, _: &span::Id, _: Context<'_, S>) {}
113    fn on_exit(&self, _: &span::Id, _: Context<'_, S>) {}
114    fn on_close(&self, _: span::Id, _: Context<'_, S>) {}
115    fn on_new_span(&self, _: &span::Attributes<'_>, _: &span::Id, _: Context<'_, S>) {}
116    fn on_record(&self, _: &span::Id, _: &span::Record<'_>, _: Context<'_, S>) {}
117    
118    fn enabled(&self, metadata: &tracing::Metadata<'_>, _: Context<'_, S>) -> bool {
119        metadata.is_event()
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use std::sync::Mutex;
127    use tracing::{span, Level};
128    use tracing_subscriber::layer::SubscriberExt;
129
130    struct TestDispatcher {
131        events: Arc<Mutex<Vec<LogEvent>>>,
132    }
133
134    impl TestDispatcher {
135        fn new() -> Self {
136            Self {
137                events: Arc::new(Mutex::new(Vec::new())),
138            }
139        }
140
141        fn events(&self) -> Vec<LogEvent> {
142            self.events.lock().unwrap().clone()
143        }
144    }
145
146    impl Dispatcher for TestDispatcher {
147        fn dispatch(&self, input: LogEvent) {
148            self.events.lock().unwrap().push(input);
149        }
150    }
151
152    #[test]
153    fn test_layer_basic_logging() {
154        let dispatcher = Arc::new(TestDispatcher::new());
155        let subscriber = tracing_subscriber::registry().with(BetterstackLayer::new(dispatcher.clone()));
156
157        tracing::subscriber::with_default(subscriber, || {
158            tracing::info!("test message");
159        });
160
161        let events = dispatcher.events();
162        assert_eq!(events.len(), 1);
163        assert_eq!(events[0].message, "test message");
164        assert_eq!(events[0].level.as_deref(), Some("INFO"));
165    }
166
167    #[test]
168    fn test_layer_with_spans() {
169        let dispatcher = Arc::new(TestDispatcher::new());
170        let subscriber = tracing_subscriber::registry().with(BetterstackLayer::new(dispatcher.clone()));
171
172        tracing::subscriber::with_default(subscriber, || {
173            let span = span!(Level::INFO, "test_span", field = "value");
174            let _guard = span.enter();
175            tracing::info!("test message in span");
176        });
177
178        let events = dispatcher.events();
179        assert_eq!(events.len(), 1);
180        assert_eq!(events[0].message, "test message in span");
181        assert_eq!(events[0].level.as_deref(), Some("INFO"));
182    }
183}
tracing_betterstack/layer.rs

tracing_betterstack/
layer.rs
