tracing_stacks/
lib.rs

1use probe::probe;
2use tracing::{field::Visit, Subscriber};
3use tracing_core::span;
4use tracing_subscriber::{prelude::*, Registry};
5
6pub fn init() {
7    let registry = Registry::default().with(StacksSubscriber {});
8    tracing::dispatcher::set_global_default(registry.into()).expect("failed to set global subscriber");
9}
10
11#[derive(Debug, Clone)]
12struct SpanInfo {
13    span_id: u64,
14    parent_span_id: u64,
15    id: u64,
16    amount: u64,
17    name: [u8; 16],
18
19    exit_stack: bool,
20}
21
22const ID: &str = "id";
23const THROUGHPUT_METRIC: &str = "amount";
24const EXIT_STACK: &str = "exit_stack";
25
26impl Visit for SpanInfo {
27    fn record_u64(&mut self, field: &tracing_core::Field, value: u64) {
28        match field.name() {
29            ID => self.id = value,
30            THROUGHPUT_METRIC => self.amount = value,
31            _ => {}
32        }
33    }
34
35    fn record_bool(&mut self, field: &tracing_core::Field, value: bool) {
36        if field.name() == EXIT_STACK {
37            self.exit_stack = value;
38        }
39    }
40
41    fn record_debug(&mut self, _: &tracing_core::Field, _: &dyn std::fmt::Debug) {}
42}
43
44pub struct StacksSubscriber {}
45
46impl<S> tracing_subscriber::Layer<S> for StacksSubscriber
47where
48    S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
49{
50    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
51        if let Some(span) = ctx.span(id) {
52            let name = attrs.metadata().name();
53            let buf = {
54                let mut buf = [0u8; 16];
55                let lth = name.len().min(buf.len());
56                buf[..lth].copy_from_slice(name.as_bytes()[..lth].as_ref());
57                buf
58            };
59            let mut info = SpanInfo {
60                name: buf,
61                parent_span_id: attrs.parent().or(ctx.current_span().id()).map_or(0, |id| id.into_u64()),
62                span_id: id.into_u64(),
63                id: 0,
64                amount: 0,
65                exit_stack: false,
66            };
67            attrs.record(&mut info);
68            span.extensions_mut().insert(info);
69        }
70    }
71
72    fn on_enter(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
73        if let Some(span) = ctx.span(id) {
74            let extensions = span.extensions();
75            let span_info = extensions.get::<SpanInfo>().unwrap();
76            probe!(
77                stacks_tracing,
78                enter,
79                span_info.span_id as *const u64,
80                span_info.parent_span_id as *const u64,
81                span_info.id as *const u64,
82                span_info.amount as *const u64,
83                span_info.name.as_ptr()
84            );
85        }
86    }
87
88    fn on_exit(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
89        if let Some(span) = ctx.span(id) {
90            let extensions = span.extensions();
91            let span_info = extensions.get::<SpanInfo>().unwrap();
92            if span_info.exit_stack {
93                probe!(stacks_tracing, exit_stack, span_info.span_id as *const u64);
94            } else {
95                probe!(stacks_tracing, exit, span_info.span_id as *const u64);
96            }
97        }
98    }
99
100    fn on_close(&self, id: span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
101        if let Some(span) = ctx.span(&id) {
102            let extensions = span.extensions();
103            let span_info = extensions.get::<SpanInfo>().unwrap();
104            probe!(stacks_tracing, close, span_info.span_id as *const u64,);
105        }
106    }
107}
tracing_stacks/lib.rs

tracing_stacks/
lib.rs
