tracing-stacks 0.1.0

Basic integration of usdt probes with tracing crate.
Documentation
use probe::probe;
use tracing::{field::Visit, Subscriber};
use tracing_core::span;
use tracing_subscriber::{prelude::*, Registry};

pub fn init() {
    let registry = Registry::default().with(StacksSubscriber {});
    tracing::dispatcher::set_global_default(registry.into()).expect("failed to set global subscriber");
}

#[derive(Debug, Clone)]
struct SpanInfo {
    span_id: u64,
    parent_span_id: u64,
    id: u64,
    amount: u64,
    name: [u8; 16],

    exit_stack: bool,
}

const ID: &str = "id";
const THROUGHPUT_METRIC: &str = "amount";
const EXIT_STACK: &str = "exit_stack";

impl Visit for SpanInfo {
    fn record_u64(&mut self, field: &tracing_core::Field, value: u64) {
        match field.name() {
            ID => self.id = value,
            THROUGHPUT_METRIC => self.amount = value,
            _ => {}
        }
    }

    fn record_bool(&mut self, field: &tracing_core::Field, value: bool) {
        if field.name() == EXIT_STACK {
            self.exit_stack = value;
        }
    }

    fn record_debug(&mut self, _: &tracing_core::Field, _: &dyn std::fmt::Debug) {}
}

pub struct StacksSubscriber {}

impl<S> tracing_subscriber::Layer<S> for StacksSubscriber
where
    S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
{
    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        if let Some(span) = ctx.span(id) {
            let name = attrs.metadata().name();
            let buf = {
                let mut buf = [0u8; 16];
                let lth = name.len().min(buf.len());
                buf[..lth].copy_from_slice(name.as_bytes()[..lth].as_ref());
                buf
            };
            let mut info = SpanInfo {
                name: buf,
                parent_span_id: attrs.parent().or(ctx.current_span().id()).map_or(0, |id| id.into_u64()),
                span_id: id.into_u64(),
                id: 0,
                amount: 0,
                exit_stack: false,
            };
            attrs.record(&mut info);
            span.extensions_mut().insert(info);
        }
    }

    fn on_enter(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        if let Some(span) = ctx.span(id) {
            let extensions = span.extensions();
            let span_info = extensions.get::<SpanInfo>().unwrap();
            probe!(
                stacks_tracing,
                enter,
                span_info.span_id as *const u64,
                span_info.parent_span_id as *const u64,
                span_info.id as *const u64,
                span_info.amount as *const u64,
                span_info.name.as_ptr()
            );
        }
    }

    fn on_exit(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        if let Some(span) = ctx.span(id) {
            let extensions = span.extensions();
            let span_info = extensions.get::<SpanInfo>().unwrap();
            if span_info.exit_stack {
                probe!(stacks_tracing, exit_stack, span_info.span_id as *const u64);
            } else {
                probe!(stacks_tracing, exit, span_info.span_id as *const u64);
            }
        }
    }

    fn on_close(&self, id: span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        if let Some(span) = ctx.span(&id) {
            let extensions = span.extensions();
            let span_info = extensions.get::<SpanInfo>().unwrap();
            probe!(stacks_tracing, close, span_info.span_id as *const u64,);
        }
    }
}