onesignal_tracing_tail_sample/
lib.rs

1//! Provide extensions to the [`tracing`] libraries to buffer complete traces
2//! and enable tail sampling.
3//!
4//! Some functionality provided by [`tracing`] is reimplemented in terms of this package, such as
5//! the opentelemetry integration. The upstream integration forwards span data as they close rather
6//! than buffer, but this prevents sophisticated tail sampling.
7//!
8//! [`tracing`]: https://github.com/tokio-rs/tracing
9
10use std::sync::{Arc, RwLock};
11use uuid::Uuid;
12
13use tracing::span;
14use tracing::subscriber::Subscriber;
15use tracing_subscriber::layer::Context;
16use tracing_subscriber::registry::LookupSpan;
17
18mod extensions;
19use extensions::{Extensions, ExtensionsInner, ExtensionsMut};
20
21pub mod opentelemetry;
22
23/// Buffers data and builds complete traces prior to exporting
24#[derive(Default, Debug)]
25pub struct TraceContextLayer<S> {
26    _registry: std::marker::PhantomData<S>,
27}
28
29pub struct TraceContext {
30    pub span_id: Uuid,
31    pub parent_id: Option<Uuid>,
32    pub trace: Trace,
33    _hidden: (),
34}
35
36#[derive(Debug)]
37pub struct TraceInner {
38    id: Uuid,
39    ext: RwLock<ExtensionsInner>,
40}
41
42#[derive(Clone)]
43pub struct Trace {
44    inner: Arc<TraceInner>,
45}
46
47impl Trace {
48    fn new() -> Self {
49        Trace {
50            inner: Arc::new(TraceInner {
51                id: Uuid::new_v4(),
52                ext: RwLock::new(ExtensionsInner::new()),
53            }),
54        }
55    }
56
57    pub fn id(&self) -> &Uuid {
58        &self.inner.id
59    }
60
61    pub fn extensions(&self) -> Extensions<'_> {
62        Extensions::new(self.inner.ext.read().expect("Mutex poisoned"))
63    }
64
65    pub fn extensions_mut(&self) -> ExtensionsMut<'_> {
66        ExtensionsMut::new(self.inner.ext.write().expect("Mutex poisoned"))
67    }
68}
69
70pub struct SampleDecision {
71    pub record_trace: bool,
72}
73
74impl TraceContext {
75    fn new() -> Self {
76        Self {
77            span_id: Uuid::new_v4(),
78            parent_id: None,
79            trace: Trace::new(),
80            _hidden: (),
81        }
82    }
83
84    fn child(&self) -> Self {
85        TraceContext {
86            span_id: Uuid::new_v4(),
87            parent_id: Some(self.span_id),
88            trace: self.trace.clone(),
89            _hidden: (),
90        }
91    }
92}
93
94impl<S> tracing_subscriber::Layer<S> for TraceContextLayer<S>
95where
96    S: Subscriber + for<'span> LookupSpan<'span>,
97    Self: 'static,
98{
99    /// Notifies this layer that a new span was constructed with the given
100    /// `Attributes` and `Id`.
101    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
102        let span = ctx.span(id).expect("Span not found, this is a bug");
103        let mut extensions = span.extensions_mut();
104        let parent = self.parent_span(attrs, &ctx);
105
106        let trace_context = parent
107            .and_then(|parent_id| {
108                let parent = ctx.span(&parent_id).expect("Span not found, this is a bug");
109                let parent_ext = parent.extensions();
110                parent_ext.get::<TraceContext>().map(|p| p.child())
111            })
112            .unwrap_or_else(TraceContext::new);
113
114        extensions.insert(trace_context);
115    }
116}
117
118impl<S> TraceContextLayer<S>
119where
120    S: Subscriber + for<'span> LookupSpan<'span>,
121    Self: 'static,
122{
123    fn parent_span(&self, attrs: &span::Attributes<'_>, ctx: &Context<'_, S>) -> Option<span::Id> {
124        if let Some(parent) = attrs.parent() {
125            Some(parent.clone())
126        } else if attrs.is_contextual() {
127            ctx.lookup_current().map(|s| s.id())
128        } else {
129            None
130        }
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use tracing::dispatcher;
137    use tracing_subscriber::prelude::*;
138    use tracing_subscriber::registry::LookupSpan;
139    use tracing_subscriber::Registry;
140
141    use super::*;
142
143    #[test]
144    fn it_works() {
145        let subscriber = Registry::default().with(TraceContextLayer::default());
146
147        tracing::subscriber::set_global_default(subscriber).unwrap();
148
149        tracing::info_span!("base_span").in_scope(|| {
150            let mut root = Uuid::nil();
151            dispatcher::get_default(|d| {
152                let registry = d.downcast_ref::<Registry>().unwrap();
153                let span = d.current_span();
154                let spanref = registry.span(&span.id().unwrap()).unwrap();
155                let extensions = spanref.extensions();
156                let trace_context = extensions.get::<TraceContext>().unwrap();
157                assert!(trace_context.parent_id.is_none());
158                root = trace_context.span_id;
159                trace_context.trace.extensions_mut().insert(42usize);
160            });
161
162            tracing::info_span!("nested_span").in_scope(|| {
163                dispatcher::get_default(|d| {
164                    let registry = d.downcast_ref::<Registry>().unwrap();
165                    let span = d.current_span();
166                    let spanref = registry.span(&span.id().unwrap()).unwrap();
167                    let extensions = spanref.extensions();
168                    let trace_context = extensions.get::<TraceContext>().unwrap();
169                    assert_eq!(trace_context.parent_id, Some(root));
170                    let trace_ext = trace_context.trace.extensions();
171                    let trace_data = trace_ext.get::<usize>().unwrap();
172                    assert_eq!(*trace_data, 42usize);
173                });
174            })
175        });
176    }
177}
onesignal_tracing_tail_sample/lib.rs

onesignal_tracing_tail_sample/
lib.rs
