tracing-calltree 0.1.2

use crate::builder::CallTreeInner;
use crate::sample::TimingSample;
use crate::tree::{CallTreeNode, NodeKey};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::Subscriber;
use tracing::span::{Attributes, Id, Record};
use tracing_subscriber::layer::{Context, Layer};
use tracing_subscriber::registry::LookupSpan;

#[derive(Clone)]
pub struct CallTreeLayer {
    inner: Arc<CallTreeInner>,
}

impl CallTreeLayer {
    pub(crate) fn new(inner: Arc<CallTreeInner>) -> Self {
        Self { inner }
    }
}

#[derive(Clone)]
struct SpanCallTreeState {
    nearest_profiled_node: Option<Arc<CallTreeNode>>,
    timing: Option<SpanTimingState>,
}

#[derive(Clone)]
struct SpanTimingState {
    node: Arc<CallTreeNode>,
    first_enter: Option<Instant>,
    active_started: Option<Instant>,
    active_elapsed: Duration,
    enter_depth: u32,
}

impl SpanTimingState {
    fn new(node: Arc<CallTreeNode>) -> Self {
        Self {
            node,
            first_enter: None,
            active_started: None,
            active_elapsed: Duration::ZERO,
            enter_depth: 0,
        }
    }
}

impl<S> Layer<S> for CallTreeLayer
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
        let Some(span) = ctx.span(id) else {
            return;
        };

        let parent_profiled = parent_profiled_node(attrs, &ctx);
        let metadata = span.metadata();
        let is_profiled = (self.inner.config.filter)(metadata);

        let mut nearest_profiled_node = parent_profiled.clone();
        let timing = if is_profiled {
            let depth = parent_profiled
                .as_ref()
                .map_or(1, |parent| parent.depth().saturating_add(1));

            if depth > self.inner.config.max_depth {
                None
            } else {
                let key = NodeKey::from_metadata(metadata);
                self.inner
                    .state
                    .get_or_create_node(
                        parent_profiled.as_ref(),
                        key,
                        depth,
                        self.inner.config.max_nodes,
                    )
                    .map(|node| {
                        nearest_profiled_node = Some(node.clone());
                        SpanTimingState::new(node)
                    })
            }
        } else {
            None
        };

        span.extensions_mut().insert(SpanCallTreeState {
            nearest_profiled_node,
            timing,
        });
    }

    fn on_record(&self, _span: &Id, _values: &Record<'_>, _ctx: Context<'_, S>) {}

    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
        let Some(span) = ctx.span(id) else {
            return;
        };

        let mut extensions = span.extensions_mut();
        let Some(state) = extensions.get_mut::<SpanCallTreeState>() else {
            return;
        };
        let Some(timing) = state.timing.as_mut() else {
            return;
        };

        if timing.enter_depth == 0 {
            let now = Instant::now();
            timing.first_enter.get_or_insert(now);
            timing.active_started = Some(now);
        }

        timing.enter_depth = timing.enter_depth.saturating_add(1);
    }

    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
        let Some(span) = ctx.span(id) else {
            return;
        };

        let mut extensions = span.extensions_mut();
        let Some(state) = extensions.get_mut::<SpanCallTreeState>() else {
            return;
        };
        let Some(timing) = state.timing.as_mut() else {
            return;
        };

        if timing.enter_depth == 0 {
            return;
        }

        timing.enter_depth -= 1;
        if timing.enter_depth == 0 {
            if let Some(started) = timing.active_started.take() {
                timing.active_elapsed += Instant::now().saturating_duration_since(started);
            }
        }
    }

    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
        let Some(span) = ctx.span(&id) else {
            return;
        };

        let state = span.extensions_mut().remove::<SpanCallTreeState>();
        let Some(mut timing) = state.and_then(|state| state.timing) else {
            return;
        };

        let Some(first_enter) = timing.first_enter else {
            return;
        };

        let now = Instant::now();
        if timing.enter_depth > 0 {
            if let Some(started) = timing.active_started.take() {
                timing.active_elapsed += now.saturating_duration_since(started);
            }
        }

        timing.node.record_sample(
            TimingSample::new(
                now.saturating_duration_since(first_enter),
                timing.active_elapsed,
            ),
            self.inner.config.window_size,
            self.inner.state.dropped_samples_counter(),
        );
    }
}

fn parent_profiled_node<S>(
    attrs: &Attributes<'_>,
    ctx: &Context<'_, S>,
) -> Option<Arc<CallTreeNode>>
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    let parent = attrs
        .parent()
        .and_then(|parent| ctx.span(parent))
        .or_else(|| {
            attrs
                .is_contextual()
                .then(|| ctx.lookup_current())
                .flatten()
        })?;

    let extensions = parent.extensions();
    extensions
        .get::<SpanCallTreeState>()
        .and_then(|state| state.nearest_profiled_node.clone())
}