trace_weft/
context.rs

1//! Ambient span context.
2//!
3//! A task-local "current span" so child spans link to their parent without the
4//! caller threading `(TraceId, RunId, SpanId)` through every signature.
5//! `SpanBuilder::run` and the instrumentation macros set this for the duration
6//! of the instrumented body; spans created inside that body read it
7//! automatically. `SpanBuilder::with_parent` remains the explicit override for
8//! cross-task / cross-thread handoffs.
9
10use std::future::Future;
11use trace_weft_core::{RunId, SpanId, SpanRecord, TraceId};
12
13/// The trace/run/span identity threaded down to child spans.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub struct SpanContext {
16    pub trace_id: TraceId,
17    pub run_id: RunId,
18    pub span_id: SpanId,
19}
20
21tokio::task_local! {
22    static CURRENT_SPAN: SpanContext;
23}
24
25/// The ambient span context for the current task, if one has been set.
26pub fn current_span_context() -> Option<SpanContext> {
27    CURRENT_SPAN.try_with(|ctx| *ctx).ok()
28}
29
30/// Run `future` with `ctx` installed as the ambient span context. Spans created
31/// inside the future link to `ctx` as their parent unless they set one
32/// explicitly.
33pub fn scope_current<F>(ctx: SpanContext, future: F) -> impl Future<Output = F::Output>
34where
35    F: Future,
36{
37    CURRENT_SPAN.scope(ctx, future)
38}
39
40/// Link a span to the ambient parent when it has no explicit parent. A span
41/// that already carries a `parent_span_id` (e.g. via `with_parent`) is left
42/// untouched.
43pub(crate) fn link_to_ambient(span: &mut SpanRecord) {
44    if span.parent_span_id.is_some() {
45        return;
46    }
47    if let Some(parent) = current_span_context() {
48        span.trace_id = parent.trace_id;
49        span.run_id = parent.run_id;
50        span.parent_span_id = Some(parent.span_id);
51    }
52}
53
54impl SpanContext {
55    pub(crate) fn of(span: &SpanRecord) -> Self {
56        Self {
57            trace_id: span.trace_id,
58            run_id: span.run_id,
59            span_id: span.span_id,
60        }
61    }
62}
trace_weft/context.rs

trace_weft/
context.rs
