tracing_newrelic/
layer.rs

1use std::thread::JoinHandle;
2
3use tokio::sync::mpsc::UnboundedSender;
4use tracing_core::span::{Attributes, Id, Record};
5use tracing_core::{Event, Subscriber};
6use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
7
8use crate::types::{NewrAttributes, NewrCommon, NewrLog, NewrLogs, NewrSpan, NewrSpans, Value};
9use crate::utils::next_trace_id;
10
11/// A [`Layer`] that collects newrelic-compatible data from `tracing` span/event.
12///
13/// [`Layer`]: tracing_subscriber::layer::Layer
14pub struct NewRelicLayer {
15    pub(crate) channel: Option<UnboundedSender<(NewrLogs, NewrSpans)>>,
16    pub(crate) handle: Option<JoinHandle<()>>,
17}
18
19impl<S> Layer<S> for NewRelicLayer
20where
21    S: Subscriber + for<'span> LookupSpan<'span>,
22{
23    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
24        let span = ctx.span(id).expect("span not found");
25        let metadata = span.metadata();
26
27        // create a new span
28        let mut nr_span = NewrSpan::new(metadata.name().to_string());
29
30        nr_span.attributes.insert(
31            "source",
32            format!(
33                "{}:{}",
34                metadata.file().unwrap_or_default(),
35                metadata.line().unwrap_or_default()
36            ),
37        );
38
39        // record span attributes
40        attrs.record(&mut nr_span.attributes);
41
42        // insert into extensions
43        span.extensions_mut().insert(nr_span);
44    }
45
46    fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
47        let span = ctx.span(id).expect("span not found");
48        let mut extensions = span.extensions_mut();
49
50        if let Some(nr_span) = extensions.get_mut::<NewrSpan>() {
51            values.record(&mut nr_span.attributes);
52        }
53    }
54
55    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
56        // ignore event that is out of current span
57        if let Some(id) = ctx.current_span().id() {
58            let span = ctx.span(id).expect("span not found");
59            let mut extensions = span.extensions_mut();
60            let metadata = event.metadata();
61
62            // create a log
63            let mut nr_log = NewrLog::new(metadata.level());
64
65            if let Some(span_id) = extensions.get_mut::<NewrSpan>().map(|s| s.id.clone()) {
66                // add linking metadata
67                // https://github.com/newrelic/node-newrelic/blob/91967dd5cd997aa283b8aa0b2fdacc2a5f10a628/api.js#L132
68                nr_log.attributes.insert("span.id", span_id);
69            }
70
71            nr_log.attributes.insert(
72                "source",
73                format!(
74                    "{}:{}",
75                    metadata.file().unwrap_or_default(),
76                    metadata.line().unwrap_or_default()
77                ),
78            );
79
80            // record event attributes
81            event.record(&mut nr_log.attributes);
82
83            // insert into extensions
84            if let Some(nr_logs) = extensions.get_mut::<Vec<NewrLog>>() {
85                nr_logs.push(nr_log);
86            } else {
87                extensions.insert(vec![nr_log]);
88            }
89        }
90    }
91
92    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
93        let span = ctx.span(&id).expect("span not found");
94        let mut extensions = span.extensions_mut();
95
96        if let Some(mut nr_span) = extensions.remove::<NewrSpan>() {
97            // update duration
98            nr_span.update_duration();
99
100            let mut logs = extensions.remove::<Vec<NewrLog>>().unwrap_or_default();
101
102            let mut spans = vec![nr_span];
103
104            if let Some(mut children) = extensions.remove::<Vec<NewrSpan>>() {
105                spans.append(&mut children);
106            };
107
108            if let Some(parent) = span.parent() {
109                let mut parent_extensions = parent.extensions_mut();
110
111                if let Some(parent_span) = parent_extensions.get_mut::<NewrSpan>() {
112                    spans[0]
113                        .attributes
114                        .insert("parent.id", parent_span.id.clone());
115
116                    if let Some(siblings) = parent_extensions.get_mut::<Vec<NewrSpan>>() {
117                        siblings.append(&mut spans);
118                    } else {
119                        parent_extensions.insert(spans);
120                    }
121
122                    if !logs.is_empty() {
123                        if let Some(parent_logs) = parent_extensions.get_mut::<Vec<NewrLog>>() {
124                            parent_logs.append(&mut logs);
125                        } else {
126                            parent_extensions.insert(logs);
127                        }
128                    }
129                }
130
131                return;
132            }
133
134            let trace_id = next_trace_id();
135
136            for span in &mut spans {
137                span.trace_id = Some(trace_id.clone());
138            }
139
140            for log in &mut logs {
141                log.attributes.insert("trace.id", trace_id.clone());
142            }
143
144            if let Some(channel) = &self.channel {
145                let mut attributes = NewrAttributes::default();
146
147                if let Some(Value::String(service_name)) =
148                    &spans[0].attributes.0.get("service.name")
149                {
150                    attributes.insert("service.name", service_name.as_str());
151                }
152
153                if let Some(Value::String(hostname)) = &spans[0].attributes.0.get("hostname") {
154                    attributes.insert("hostname", hostname.as_str());
155                }
156
157                // TODO: error handling
158                let _ = channel.send((
159                    NewrLogs {
160                        logs,
161                        common: NewrCommon {
162                            attributes: attributes.clone(),
163                        },
164                    },
165                    NewrSpans {
166                        spans,
167                        common: NewrCommon { attributes },
168                    },
169                ));
170            }
171        }
172    }
173}
174
175impl Drop for NewRelicLayer {
176    fn drop(&mut self) {
177        if let Some(channel) = self.channel.take() {
178            drop(channel);
179        }
180
181        if let Some(handle) = self.handle.take() {
182            let _ = handle.join();
183        }
184    }
185}
tracing_newrelic/layer.rs

tracing_newrelic/
layer.rs
