tracelite/
tracer.rs

1use std::borrow::Cow;
2use std::collections::BTreeMap;
3use std::num::NonZeroU32;
4use std::sync::Arc;
5use std::num::{NonZeroU128, NonZeroU64};
6use crate::sampling::{DynamicTraceDetail, SamplingDecision, SamplingResult};
7use crate::{AttributeValue, Severity};
8
9#[derive(Debug, Clone, Copy)]
10pub struct NoTracerRegistered;
11
12impl std::fmt::Display for NoTracerRegistered {
13    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
14        f.write_str("operation requires tracer, but no tracer was registered")
15    }
16}
17
18impl std::error::Error for NoTracerRegistered {}
19
20#[derive(Debug, Clone, Copy)]
21pub struct NoCurrentSpan;
22
23impl std::fmt::Display for NoCurrentSpan {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.write_str("operation requires span, but no span has been opened in this context")
26    }
27}
28
29impl std::error::Error for NoCurrentSpan {}
30
31// #[derive(Debug, Clone, Copy)]
32// pub struct NonRecordingRootSpan;
33
34// impl std::fmt::Display for NonRecordingRootSpan {
35//     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36//         f.write_str("operation requires root span to be recording, but it was not sampled for recording")
37//     }
38// }
39
40// impl std::error::Error for NonRecordingRootSpan {}
41
42/// Can only be constructed by this crate - used to mark operation which
43/// must not be freely called by users, but only in one specific circumstance.
44#[derive(Debug)]
45pub struct PrivateMarker(());
46
47// ++++++++++++++++++++ Text ++++++++++++++++++++
48
49/// TODO with the addition of FormatArgs, the name does not fit anymore
50#[derive(Debug, Clone, Copy)]
51pub enum Text<'a> {
52    Static(&'static str),
53    Borrowed(&'a str),
54    FormatArgs(std::fmt::Arguments<'a>)
55}
56
57impl<'a> Text<'a> {
58    pub fn as_str(&self) -> Option<&'a str> {
59        match self {
60            Text::Static(s) => Some(s),
61            Text::Borrowed(s) => Some(s),
62            Text::FormatArgs(_) => None,
63        }
64    }
65}
66
67impl From<&'static str> for Text<'static> {
68    fn from(s: &'static str) -> Self { Self::Static(s) }
69}
70
71impl<'a> From<std::fmt::Arguments<'a>> for Text<'a> {
72    fn from(a: std::fmt::Arguments<'a>) -> Self { Self::FormatArgs(a) }
73}
74
75impl<'a> std::fmt::Display for Text<'a> {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        match self {
78            Text::Static(x) => x.fmt(f),
79            Text::Borrowed(x) => x.fmt(f),
80            Text::FormatArgs(x) => x.fmt(f),
81        }
82    }
83}
84
85// ++++++++++++++++++++ baggage ++++++++++++++++++++
86
87pub type BaggageKey = Cow<'static, str>;
88pub type BaggageValue = Cow<'static, str>;
89pub type BaggageMap = BTreeMap<BaggageKey, BaggageValue>;
90
91// ++++++++++++++++++++ ids ++++++++++++++++++++
92
93// NOTE uses big-endian to store 16 bytes as u128
94#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
95pub struct TraceId(pub NonZeroU128);
96
97impl std::fmt::Display for TraceId {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        write!(f, "{:032x}", self.0.get())
100    }
101}
102
103impl std::str::FromStr for TraceId {
104    type Err = &'static str;
105    fn from_str(hex: &str) -> Result<Self, Self::Err> {
106        let mut bytes = [0; 16];
107        if hex.len() != bytes.len() * 2 {
108            return Err("trace-id string has wrong length")
109        }
110        for (i, byte) in bytes.iter_mut().enumerate() {
111            *byte = match u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16) {
112                Ok(b) => b,
113                Err(_) => return Err("trace-id string contains non-hex letters")
114            }
115        }
116        let int = u128::from_be_bytes(bytes);
117        NonZeroU128::new(int).map(Self).ok_or("trace-id string has no non-zero bytes")
118    }
119}
120
121#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
122pub struct SpanId(pub NonZeroU64);
123
124impl std::fmt::Display for SpanId {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        write!(f, "{:016x}", self.0.get())
127    }
128}
129
130impl std::str::FromStr for SpanId {
131    type Err = &'static str;
132    fn from_str(hex: &str) -> Result<Self, Self::Err> {
133        let mut bytes = [0; 8];
134        if hex.len() != bytes.len() * 2 {
135            return Err("span-id string has wrong length")
136        }
137        for (i, byte) in bytes.iter_mut().enumerate() {
138            *byte = match u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16) {
139                Ok(b) => b,
140                Err(_) => return Err("span-id string contains non-hex letters")
141            }
142        }
143        let int = u64::from_be_bytes(bytes);
144        NonZeroU64::new(int).map(Self).ok_or("span-id string has no non-zero bytes")
145    }
146}
147
148// ++++++++++++++++++++ SpanContext ++++++++++++++++++++
149
150#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
151pub struct SpanContext {
152    pub trace_id: TraceId,
153    pub span_id: SpanId,
154    pub trace_flags: u8,
155}
156
157impl SpanContext {
158    pub fn sampled(&self) -> bool {
159        self.trace_flags & 0 != 0
160    }
161}
162
163impl std::fmt::Display for SpanContext {
164    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165        write!(f, "00-{}-{}-{:02x}", self.trace_id, self.span_id, self.trace_flags)
166    }
167}
168
169impl std::str::FromStr for SpanContext {
170    type Err = &'static str;
171
172    fn from_str(header: &str) -> Result<Self, Self::Err> {
173        let mut parts = header.split('-');
174
175        let Some(version) = parts.next() else { return Err("missing '-' delimiters") };
176        if version != "00" { return Err("unsupported version") }
177
178        let Some(trace_id_hex) = parts.next() else { return Err("missing trace-id field") };
179        let Some(span_id_hex) = parts.next() else { return Err("missing parent-id field") };
180        let Some(trace_flags_hex) = parts.next() else { return Err("missing trace-flags field") };
181
182        let trace_id = TraceId::from_str(trace_id_hex)?;
183        let span_id = SpanId::from_str(span_id_hex)?;
184
185        if trace_flags_hex.len() != 2 { return Err("trace-flags string has wrong length")}
186        let trace_flags = u8::from_str_radix(trace_flags_hex, 16)
187            .map_err(|_| "trace-flags string contains non-hex letters")?;
188
189        Ok(Self { trace_id, span_id, trace_flags })
190    }
191}
192
193// ++++++++++++++++++++ SpanParent ++++++++++++++++++++
194
195// TODO is NonZero worth it?
196#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
197pub struct SpanCollectionIndex(pub NonZeroU32, pub u32);
198
199#[derive(Default, Clone)]
200pub struct TracingContext {
201    pub baggage: BaggageMap,
202    pub on_start: Option<Arc<dyn Fn(&SpanRef) + Send + Sync>>,
203    pub on_ending: Option<Arc<dyn Fn(&SpanRef) + Send + Sync>>,
204}
205
206#[derive(Debug, Clone, Copy)]
207pub struct RecordingSpanContext {
208    pub span_context: SpanContext,
209    pub collect_idx: SpanCollectionIndex,
210}
211
212#[derive(Default, Clone)]
213pub struct SpanRef {
214    span_context: Option<SpanContext>,
215    collect_idx: Option<SpanCollectionIndex>,
216    references_ancestor: bool,
217    dyn_trace_detail: Option<DynamicTraceDetail>,
218    tracing_context: Option<Arc<TracingContext>>,
219}
220
221impl SpanRef {
222    pub fn disabled() -> Self { Self::default() }
223
224    pub fn remote(span_context: SpanContext) -> Self {
225        Self{ span_context: Some(span_context), ..Self::disabled() }
226    }
227
228    #[doc(hidden)]
229    pub fn recording(span_context: RecordingSpanContext) -> Self {
230        Self{ collect_idx: Some(span_context.collect_idx), ..Self::remote(span_context.span_context) }
231    }
232
233    pub fn as_ancestor(&self) -> Self {
234        let mut this = self.clone();
235        this.references_ancestor = true;
236        this
237    }
238
239    pub fn as_recording_self_or_ancestor(&self) -> Option<RecordingSpanContext> {
240        Some(RecordingSpanContext{ span_context: self.span_context?, collect_idx: self.collect_idx? })
241    }
242    pub fn as_recording_self(&self) -> Option<RecordingSpanContext> {
243        if self.references_ancestor { None } else { self.as_recording_self_or_ancestor() }
244    }
245
246    pub fn is_recording_self_or_ancestor(&self) -> bool {
247        self.collect_idx.is_some()
248    }
249    pub fn is_recording_self(&self) -> bool {
250        !self.references_ancestor && self.collect_idx.is_some()
251    }
252    pub fn span_context(&self) -> Option<SpanContext> {
253        self.span_context
254    }
255    pub fn tracing_context(&self) -> Option<&Arc<TracingContext>> {
256        self.tracing_context.as_ref()
257    }
258
259    // TODO remove these methods?
260    // pub fn is_recording(&self) -> bool {
261    //     self.collect_idx.is_some()
262    // }
263    // pub fn is_recording_current(&self) -> bool {
264    //     self.collect_idx.is_some() && !self.references_ancestor
265    // }
266    // pub fn collection_index(&self) -> Option<SpanCollectionIndex> {
267    //     self.collect_idx
268    // }
269
270    // pub fn with_dyn_trace_detail(self, dyn_trace_detail: ScopedSeverityFilter) -> Self {
271    //     Self{ scoped_severity_filter: f.or(self.scoped_severity_filter), ..self }
272    // }
273
274    // pub fn with_trace_context(self, t: Option<Arc<TracingContext>>) -> Self {
275    //     Self{ tracing_context: t, ..self }
276    // }
277}
278
279// ++++++++++++++++++++ spans ++++++++++++++++++++
280
281#[derive(Debug, Clone, Copy)]
282pub struct AttributeList<'a>(pub &'a [(Text<'a>, AttributeValue<'a>)]);
283
284impl<'a> AttributeList<'a> {
285    pub fn get(&self, key: &str) -> Option<AttributeValue<'a>> {
286        self.0.iter()
287            .find(|(k, _)| k.as_str() == Some(key))
288            .map(|(_, v)| *v)
289    }
290}
291
292#[derive(Debug, Clone)]
293pub enum SpanStatus<'a> {
294    Ok,
295    Error(Text<'a>),
296}
297
298impl<'a> SpanStatus<'a> {
299    pub fn error(s: impl Into<Text<'a>>) -> Self {
300        Self::Error(s.into())
301    }
302}
303
304#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305pub enum SpanKind {
306    Internal,
307    Client,
308    Server,
309    Producer,
310    Consumer,
311}
312
313#[non_exhaustive]
314#[derive(Debug)]
315pub struct SpanArgs<'a> {
316    pub name: Text<'a>, 
317    pub target: Option<&'static str>,
318    pub severity: Option<Severity>,
319
320    pub parent: Option<SpanContext>,
321    pub status: Option<SpanStatus<'a>>,
322    pub kind: Option<SpanKind>,
323
324    pub attributes: AttributeList<'a>,
325
326    _private: PrivateMarker,
327}
328
329pub struct SpanBuilder<'a> {
330    args: SpanArgs<'a>,
331    parent: Option<Option<&'a SpanRef>>,
332    on_start: Option<Arc<dyn Fn(&SpanRef) + Send + Sync>>,
333    on_ending: Option<Arc<dyn Fn(&SpanRef) + Send + Sync>>,
334    baggage_entries: Vec<(BaggageKey, BaggageValue)>,
335}
336
337impl<'a> SpanBuilder<'a> {
338    pub fn new(name: impl Into<Text<'a>>, target: Option<&'static str>, severity: Option<Severity>) -> Self {
339        let name = name.into();
340        Self{
341            args: SpanArgs{
342                name,
343                target,
344                severity,
345                parent: None,
346                kind: None,
347                status: None,
348                attributes: AttributeList(&[]),
349                _private: PrivateMarker(()),
350            },
351            parent: None,
352            on_start: None,
353            on_ending: None,
354            baggage_entries: vec![],
355        }
356    }
357
358    pub fn no_parent(self) -> Self {
359        Self{ parent: Some(None), ..self }
360    }
361    pub fn parent<'b>(self, parent: &'b SpanRef) -> SpanBuilder<'b>
362        where 'a: 'b
363    {
364        SpanBuilder{ parent: Some(Some(parent)), ..self }
365    }
366    pub fn name<'b>(self, name: impl Into<Text<'b>>) -> SpanBuilder<'b>
367        where 'a: 'b
368    {
369        SpanBuilder{ args: SpanArgs{ name: name.into(), ..self.args }, ..self }
370    }
371    pub fn status<'b>(self, status: SpanStatus<'a>) -> SpanBuilder<'b>
372        where 'a: 'b
373    {
374        SpanBuilder{ args: SpanArgs{ status: Some(status), ..self.args }, ..self }
375    }
376    pub fn kind(self, kind: SpanKind) -> Self {
377        Self{ args: SpanArgs{ kind: Some(kind), ..self.args }, ..self }
378    }
379
380    pub fn on_start(self, f: impl Fn(&SpanRef) + Send + Sync + 'static) -> Self {
381        Self{ on_start: Some(Arc::new(f)), ..self }
382    }
383    pub fn on_ending(self, f: impl Fn(&SpanRef) + Send + Sync + 'static) -> Self {
384        Self{ on_ending: Some(Arc::new(f)), ..self }
385    }
386    pub fn baggage_entry(mut self, key: impl Into<BaggageKey>, value: impl Into<BaggageValue>) -> Self {
387        self.baggage_entries.push((key.into(), value.into()));
388        self
389    }
390
391    fn _start<'b>(mut self, tracer: &dyn Tracer, parent: Option<&SpanRef>, attributes: AttributeList<'b>) -> OwnedSpanRef {
392        if let Some(parent) = parent {
393            match (parent.dyn_trace_detail, self.args.severity) {
394                (Some(min), Some(sev)) if sev < min.recording_level => {
395                    return OwnedSpanRef(SpanRef{ references_ancestor: true, ..parent.clone() })
396                }
397                _ => {}
398            }
399
400            // copy over parent to SpanArgs
401            self.args.parent = parent.span_context;
402        }
403
404        // set span status to Error if unset and severity>=Error
405        if self.args.status.is_none() && self.args.severity >= Some(Severity::Error) {
406            self.args.status = Some(SpanStatus::error(self.args.name))
407        }
408
409        // start a new span, copy the parent span, or create a disabled root span
410        let mut new_span = match tracer.start_span(SpanArgs{ attributes, ..self.args }, &mut PrivateMarker(())) {
411            Some((new_span, sampling)) => {
412                debug_assert_ne!(sampling.decision, SamplingDecision::Drop);
413                let dyn_trace_detail = sampling.dyn_trace_detail.or(parent.as_ref().and_then(|p| p.dyn_trace_detail));
414                SpanRef{ dyn_trace_detail, ..SpanRef::recording(new_span) }
415            }
416            None => match parent {
417                Some(parent) => SpanRef{ references_ancestor: true, ..parent.clone() },
418                None => SpanRef::disabled()
419            }
420        };
421
422        // update TracingContext, if necessary
423        if self.on_start.is_some() || self.on_ending.is_some() || !self.baggage_entries.is_empty() {
424            let mut tracing_ctx = new_span.tracing_context.as_ref()
425                .map(|t| (**t).clone())
426                .unwrap_or_default();
427            tracing_ctx.on_start = self.on_start.or(tracing_ctx.on_start);
428            tracing_ctx.on_ending = self.on_ending.or(tracing_ctx.on_ending);
429            tracing_ctx.baggage.extend(self.baggage_entries);
430            new_span.tracing_context = Some(Arc::new(tracing_ctx));
431        } else {
432            new_span.tracing_context = parent.as_ref().and_then(|p| p.tracing_context.clone());
433        }
434
435        // trigger on_start() 
436        if let Some(on_start) = new_span.tracing_context.as_ref().and_then(|t| t.on_start.as_ref()) {
437            on_start(&new_span);
438        }
439
440        OwnedSpanRef(new_span)
441    }
442
443    #[doc(hidden)]
444    pub fn start<'b>(mut self, tracer: &dyn Tracer, attributes: AttributeList<'b>) -> OwnedSpanRef {
445        match self.parent.take() {
446            Some(parent) => self._start(tracer, parent, attributes),
447            None => {
448                let mut this = Some(self);
449                globals::current_span(|parent| this.take().unwrap()._start(tracer, Some(parent), attributes))
450                    .unwrap_or_else(|_| this.take().unwrap()._start(tracer, None, attributes))
451            }
452        }
453    }
454}
455
456/// Like SpanBuilder, except everything but the parent span information gets discarded
457pub struct DisabledSpanBuilder<'a> {
458    parent: Option<Option<&'a SpanRef>>,
459}
460
461impl<'a> DisabledSpanBuilder<'a> {
462    pub fn new() -> DisabledSpanBuilder<'static> { DisabledSpanBuilder{ parent: None } }
463
464    pub fn no_parent(self) -> Self { 
465        DisabledSpanBuilder{ parent: Some(None) }
466    }
467    pub fn parent<'b>(self, parent: &'b SpanRef) -> DisabledSpanBuilder<'b>
468        where 'a: 'b
469    {
470        DisabledSpanBuilder{ parent: Some(Some(parent)) }
471    }
472
473    pub fn name<'b>(self, _name: impl Into<Text<'b>>) -> Self { self }
474    pub fn status(self, _status: SpanStatus) -> Self { self }
475    pub fn kind(self, _kind: SpanKind) -> Self { self }
476    pub fn on_start(self, _f: impl Fn(&SpanRef) + Send + Sync + 'static) -> Self { self }
477    pub fn on_ending(self, _f: impl Fn(&SpanRef) + Send + Sync + 'static) -> Self { self }
478    pub fn baggage_entry(self, _key: impl Into<BaggageKey>, _value: impl Into<BaggageValue>) -> Self { self }
479
480    fn _start(_tracer: &dyn Tracer, parent: Option<&SpanRef>) -> Option<OwnedSpanRef> {
481        if let Some(parent) = parent {
482            if parent.references_ancestor {
483                None
484            } else {
485                Some(OwnedSpanRef(SpanRef{ references_ancestor: true, ..parent.clone() }))
486            }
487        } else {
488            Some(OwnedSpanRef(SpanRef::disabled()))
489        }
490    }
491
492    #[doc(hidden)]
493    pub fn start<'b>(mut self, tracer: &dyn Tracer) -> Option<OwnedSpanRef> {
494        match self.parent.take() {
495            Some(parent) => Self::_start(tracer, parent),
496            None => {
497                globals::current_span(|parent| Self::_start(tracer, Some(parent)))
498                    .unwrap_or_else(|_| Self::_start(tracer, None))
499            }
500        }
501    }
502}
503
504// ++++++++++++++++++++ events ++++++++++++++++++++
505
506#[non_exhaustive]
507#[derive(Debug)]
508pub enum Exception<'a> {
509    Error{
510        object: &'a (dyn std::error::Error + 'a),
511        type_name: &'static str
512    },
513    Dbgfmt{
514        object: &'a (dyn std::fmt::Debug + 'a),
515        type_name: &'static str
516    }
517}
518
519#[non_exhaustive]
520#[derive(Debug)]
521pub struct EventArgs<'a> {
522    // NOTE is `exception` is set, this may be interpreted as the span status message
523    pub name: Text<'a>,
524    pub target: Option<&'a str>,
525    pub severity: Option<Severity>,
526    pub exception: Option<Exception<'a>>,
527    pub attributes: AttributeList<'a>,
528    _private: PrivateMarker,
529}
530
531pub struct EventBuilder<'a>(EventArgs<'a>);
532
533impl<'a> EventBuilder<'a> {
534    pub fn new(name: impl Into<Text<'a>>, target: Option<&'a str>, severity: Option<Severity>) -> Self {
535        Self(EventArgs{
536            name: name.into(),
537            target,
538            severity,
539            exception: None,
540            attributes: AttributeList(&[]),
541            _private: PrivateMarker(()),
542        })
543    }
544
545    pub fn exception<'b>(self, object: &'b (impl std::error::Error + 'b)) -> EventBuilder<'b>
546        where 'a: 'b
547    {
548        let type_name = std::any::type_name_of_val(object);
549        EventBuilder(EventArgs{ exception: Some(Exception::Error{ object, type_name }), ..self.0 })
550    }
551
552    pub fn exception_dbgfmt<'b>(self, object: &'b (impl std::fmt::Debug + 'b)) -> EventBuilder<'b>
553        where 'a: 'b
554    {
555        let type_name = std::any::type_name_of_val(object);
556        EventBuilder(EventArgs{ exception: Some(Exception::Dbgfmt{ object, type_name }), ..self.0 })
557    }
558    
559    #[doc(hidden)]
560    pub fn add_to_span<'b>(self, tracer: &dyn Tracer, span: &OwnedSpanRef, attributes: AttributeList<'b>){
561        match (span.dyn_trace_detail, self.0.severity) {
562            (Some(min), Some(sev)) if sev < min.recording_level => {
563                return // discard event
564            }
565            _ => {}
566        }
567
568        let Some(recording) = span.as_recording_self_or_ancestor() else {
569            if self.0.severity >= Some(Severity::Error) {
570                let err = InstrumentationError::FailedErrorEventEscalation(span, self.0);
571                tracer.instrumentation_error(err);
572            }
573            return
574        };
575
576        // set SpanStatus if severity >= Error
577        if self.0.severity >= Some(Severity::Error) {
578            tracer.set_status(recording, SpanStatus::Error(self.0.name));
579        }
580
581        // add event to span
582        tracer.add_event(recording, EventArgs{ attributes, ..self.0 }, &mut PrivateMarker(()));
583    }
584
585    #[doc(hidden)]
586    pub fn add_to_current_span<'b>(self, tracer: &dyn Tracer, attributes: AttributeList<'b>){
587        let mut this = Some(self);
588        globals::current_span(|span| {
589            let this = this.take().unwrap();
590            this.add_to_span(tracer, span, attributes);
591        }).unwrap_or_else(|_| {
592            let this = this.take().unwrap();
593            tracer.instrumentation_error(InstrumentationError::StrayEvent(this.0));
594        });
595    }
596
597
598}
599
600// ++++++++++++++++++++ Tracer ++++++++++++++++++++
601
602#[non_exhaustive]
603pub enum InstrumentationError<'a> {
604    StrayAttributes(AttributeList<'a>),
605    StrayEvent(EventArgs<'a>),
606    StrayStatus(SpanStatus<'a>),
607    FailedErrorEventEscalation(&'a SpanRef, EventArgs<'a>),
608}
609
610impl<'a> std::fmt::Display for InstrumentationError<'a> {
611    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
612        match self {
613            Self::StrayAttributes(attr) => write!(f, "stray attributes without target span: {:?}", attr),
614            Self::StrayEvent(event_args) => write!(f, "stray event without target span: {:?}", event_args),
615            Self::StrayStatus(span_status) => write!(f, "stray span status without target span: {:?}", span_status),
616            // TODO rework text to make it more understandable
617            Self::FailedErrorEventEscalation(_span, event_args) => write!(f, "event has escalating severity without any root span: {:?}", event_args),
618        }
619    }
620}
621
622pub trait Tracer: Send + Sync + 'static {
623    fn is_enabled(&self, target: Option<&'static str>, severity: Option<Severity>) -> bool;
624
625    fn start_span(&self, args: SpanArgs, _: &mut PrivateMarker) -> Option<(RecordingSpanContext, SamplingResult)>;
626    fn set_attributes(&self, span: RecordingSpanContext, attrs: AttributeList);
627    fn add_event(&self, span: RecordingSpanContext, args: EventArgs, _: &mut PrivateMarker);
628    fn set_status(&self, span: RecordingSpanContext, status: SpanStatus);
629    fn drop_span(&self, span: RecordingSpanContext, _: &mut PrivateMarker);
630    fn flush(&self);
631
632    fn instrumentation_error(&self, err: InstrumentationError);
633}
634
635pub struct OwnedSpanRef(SpanRef);
636
637impl std::ops::Deref for OwnedSpanRef {
638    type Target = SpanRef;
639    fn deref(&self) -> &Self::Target { &self.0 }
640}
641
642impl Drop for OwnedSpanRef {
643    fn drop(&mut self) {
644        if let Some(recording) = self.0.as_recording_self() {
645            // trigger on_ending
646            self.0.tracing_context.as_ref()
647                .and_then(|ctx| ctx.on_ending.as_ref())
648                .map(|on_ending| (on_ending)(&self.0));
649
650            let tracer = globals::tracer().unwrap();
651            tracer.drop_span(recording, &mut PrivateMarker(()));
652        }
653    }
654}
655
656impl OwnedSpanRef {
657    pub fn set_attributes<'a>(&self, tracer: &dyn Tracer, attrs: AttributeList<'a>){
658        if attrs.0.is_empty() { return }
659        if let Some(recording) = self.as_recording_self_or_ancestor() {
660            tracer.set_attributes(recording, attrs);
661        }
662    }
663
664    pub fn set_status(&self, tracer: &dyn Tracer, status: SpanStatus){
665        if let Some(recording) = self.as_recording_self_or_ancestor() {
666            tracer.set_status(recording, status);
667        }
668    }
669}
670
671pub mod globals {
672    use super::{AttributeList, InstrumentationError, NoCurrentSpan, NoTracerRegistered, OwnedSpanRef, SpanStatus, Text, Tracer};
673    use tokio::task::futures::TaskLocalFuture;
674    use std::future::Future;
675
676    static TRACER: std::sync::OnceLock<Box<dyn Tracer>> = std::sync::OnceLock::new();
677
678    tokio::task_local! {
679        static CURRENT_SPAN: OwnedSpanRef;
680    }
681
682    pub fn set_tracer(tracer: Box<dyn Tracer>){
683        TRACER.set(tracer).ok()
684            .expect("[FATAL] tracelite: tracer already set");
685    }
686
687    pub fn tracer() -> Result<&'static dyn Tracer, NoTracerRegistered> {
688        TRACER.get().map(|f| &**f).ok_or(NoTracerRegistered)
689    }
690
691    pub fn current_span<T>(f: impl FnOnce(&OwnedSpanRef) -> T) -> Result<T, NoCurrentSpan> {
692        CURRENT_SPAN.try_with(f).map_err(|_| NoCurrentSpan)
693    }
694
695    pub enum MaybeInSpan<F>{
696        InSpan(TaskLocalFuture<OwnedSpanRef, F>),
697        NotInSpan(F),
698    }
699
700    impl<F> std::future::Future for MaybeInSpan<F>
701        where F: std::future::Future
702    {
703        type Output = F::Output;
704        fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
705            // TODO can we get rid of unsafe?
706            match unsafe { std::pin::Pin::get_unchecked_mut(self) } {
707                MaybeInSpan::InSpan(f) => unsafe { std::pin::Pin::new_unchecked(f).poll(cx) },
708                MaybeInSpan::NotInSpan(f) => unsafe { std::pin::Pin::new_unchecked(f).poll(cx) },
709            }
710        }
711    }
712
713    pub trait InSpan: Future + Sized {
714        fn in_span(self, span: impl Into<Option<OwnedSpanRef>>) -> MaybeInSpan<Self> {
715            if let Some(span) = span.into() {
716                MaybeInSpan::InSpan(CURRENT_SPAN.scope(span, self))
717            } else {
718                MaybeInSpan::NotInSpan(self)
719            }
720        }
721    }
722
723    impl<F: Future + Sized> InSpan for F {}
724
725    pub fn sync_in_span<R>(span: impl Into<Option<OwnedSpanRef>>, f: impl FnOnce() -> R) -> R {
726        if let Some(span) = span.into() {
727            CURRENT_SPAN.sync_scope(span, f)
728        } else {
729            f()
730        }
731    }
732
733    /// NOTE you will likely want to use span_attributes!() instead
734    pub fn set_attributes<'a>(attrs: AttributeList<'a>) {
735        if attrs.0.is_empty() { return }
736        let Some(tracer) = tracer().ok() else { return };
737        let mut attrs = Some(attrs);
738
739        current_span(|span| {
740            span.set_attributes(tracer, attrs.take().unwrap());
741        }).ok().unwrap_or_else(|| {
742            let attrs = attrs.take().unwrap();
743            let err = InstrumentationError::StrayAttributes(attrs);
744            tracer.instrumentation_error(err);
745        });
746    }
747
748    pub fn set_status(status: SpanStatus<'_>){
749        let Some(tracer) = tracer().ok() else { return };
750        let mut status = Some(status);
751
752        current_span(|span| {
753            span.set_status(tracer, status.take().unwrap());
754        }).ok().unwrap_or_else(|| {
755            let status = status.take().unwrap();
756            let err = InstrumentationError::StrayStatus(status);
757            tracer.instrumentation_error(err);
758        });
759    }
760
761    pub fn set_ok_status(){
762        set_status(SpanStatus::Ok);
763    }
764
765    pub fn set_error_status<'a>(msg: impl Into<Text<'a>>){
766        set_status(SpanStatus::error(msg));
767    }
768
769}
tracelite/tracer.rs

tracelite/
tracer.rs
