tide_server_timing/
timing_layer.rs

1use tracing_core::span::{Attributes, Id};
2use tracing_core::{Dispatch, Subscriber};
3
4use tracing_subscriber::layer::Context;
5use tracing_subscriber::registry::LookupSpan;
6use tracing_subscriber::Layer;
7
8use std::iter::Extend;
9use std::{any::TypeId, marker::PhantomData, time::Instant};
10
11// this function "remembers" the types of the subscriber and the formatter,
12// so that we can downcast to something aware of them without knowing those
13// types at the callsite.
14pub(crate) struct WithContext {
15    set: fn(&Dispatch, &Id, SpanRootTiming),
16    take: fn(&Dispatch, &Id, f: &mut dyn FnMut(SpanTiming)),
17}
18
19impl WithContext {
20    pub(crate) fn set<'a>(&self, dispatch: &'a Dispatch, id: &Id, value: SpanRootTiming) {
21        (self.set)(dispatch, id, value)
22    }
23
24    pub(crate) fn take<'a>(&self, dispatch: &'a Dispatch, id: &Id, f: &mut dyn FnMut(SpanTiming)) {
25        (self.take)(dispatch, id, f)
26    }
27}
28
29/// A subscriber for `tracing` that encodes metadata for the Tide middleware to read out.
30#[allow(missing_debug_implementations)]
31pub struct TimingLayer<S> {
32    ctx: WithContext,
33    _subscriber: PhantomData<fn(S)>,
34    _priv: (),
35}
36
37impl<S> TimingLayer<S>
38where
39    S: Subscriber + for<'span> LookupSpan<'span>,
40{
41    /// Create a new instance.
42    pub fn new() -> Self {
43        let ctx = WithContext {
44            set: Self::set_timing,
45            take: Self::take_timing,
46        };
47
48        Self {
49            ctx,
50            _subscriber: PhantomData,
51            _priv: (),
52        }
53    }
54
55    fn set_timing(dispatch: &Dispatch, id: &Id, timing: SpanRootTiming) {
56        let subscriber = dispatch
57            .downcast_ref::<S>()
58            .expect("subscriber should downcast to expected type; this is a bug!");
59        let span = subscriber
60            .span(id)
61            .expect("registry should have a span for the current ID");
62
63        let mut extensions = span.extensions_mut();
64        extensions.insert(timing);
65    }
66
67    fn take_timing(dispatch: &Dispatch, id: &Id, f: &mut dyn FnMut(SpanTiming)) {
68        let subscriber = dispatch
69            .downcast_ref::<S>()
70            .expect("subscriber should downcast to expected type; this is a bug!");
71        let span = subscriber
72            .span(id)
73            .expect("registry should have a span for the current ID");
74
75        let mut extensions = span.extensions_mut();
76        if let Some(value) = extensions.remove::<SpanTiming>() {
77            f(value);
78        }
79    }
80}
81
82impl<S> Layer<S> for TimingLayer<S>
83where
84    S: Subscriber + for<'a> LookupSpan<'a>,
85{
86    fn new_span(&self, _attrs: &Attributes<'_>, id: &Id, cx: Context<'_, S>) {
87        let span = cx.span(id).unwrap();
88        let meta = span.metadata();
89        let id = span.id();
90        let timing = SpanTiming::new(id, meta.name(), meta.target());
91        span.extensions_mut().insert(timing);
92    }
93
94    /// On exit fold the current span's timing into its parent timing.
95    fn on_exit(&self, id: &Id, cx: Context<'_, S>) {
96        let span = cx.span(id).unwrap();
97
98        // Don't fold the root timing into its parent.
99        if span.extensions().get::<SpanRootTiming>().is_some() {
100            return;
101        };
102
103        // let name = span.metadata().name();
104        let mut timing = match span.extensions_mut().remove::<SpanTiming>() {
105            Some(timing) => timing,
106            None => return,
107        };
108
109        // Snapshot the end time of the span.
110        timing.end_timing();
111
112        // fold the current timing into its parent's timing.
113        let span = cx.span(id).unwrap();
114        if let Some(parent_id) = span.parent_id() {
115            let parent = cx.span(parent_id).expect("parent not found");
116            if let Some(parent_timing) = parent.extensions_mut().get_mut::<SpanTiming>() {
117                parent_timing.children.extend(timing.flatten());
118            };
119        }
120    }
121
122    unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
123        match id {
124            id if id == TypeId::of::<Self>() => Some(self as *const _ as *const ()),
125            id if id == TypeId::of::<WithContext>() => Some(&self.ctx as *const _ as *const ()),
126            _ => None,
127        }
128    }
129}
130
131/// Indicated the current struct is the root struct.
132#[derive(Debug)]
133pub(crate) struct SpanRootTiming;
134
135/// A timing that represent a span beneath the root span.
136#[derive(Debug)]
137pub(crate) struct SpanTiming {
138    pub(crate) target: &'static str,
139    pub(crate) id: Id,
140    pub(crate) span_name: &'static str,
141    pub(crate) start_time: Instant,
142    pub(crate) end_time: Option<Instant>,
143    children: Vec<Self>,
144}
145
146// /// An empty type, denoting the current span is a root.
147// struct RootSpanTiming {}
148
149impl SpanTiming {
150    fn new(id: Id, span_name: &'static str, target: &'static str) -> Self {
151        Self {
152            start_time: Instant::now(),
153            end_time: None,
154            id,
155            span_name,
156            target,
157            children: vec![],
158        }
159    }
160
161    fn end_timing(&mut self) {
162        self.end_time = Some(Instant::now());
163    }
164
165    pub(crate) fn flatten(mut self) -> Vec<Self> {
166        let mut children = vec![];
167        std::mem::swap(&mut self.children, &mut children);
168        children.push(self);
169        children
170    }
171}
tide_server_timing/timing_layer.rs

tide_server_timing/
timing_layer.rs
