micromegas_tracing/spans/
instrumented_future.rs

1//! Manual async span instrumentation using InstrumentedFuture wrapper
2
3use crate::dispatch::{on_begin_async_scope, on_end_async_scope};
4use crate::spans::{DynSpanName, SpanLocation, SpanMetadata, SpanName, StaticSpanName};
5use pin_project::pin_project;
6use std::cell::UnsafeCell;
7use std::future::Future;
8use std::pin::Pin;
9use std::task::{Context, Poll};
10
11thread_local! {
12    static ASYNC_CALL_STACK: UnsafeCell<Vec<u64>> = UnsafeCell ::new(vec![0]);
13}
14
15/// Returns the current span ID from the async call stack.
16/// Returns 0 (root) if no span is active.
17#[inline]
18pub fn current_span_id() -> u64 {
19    ASYNC_CALL_STACK.with(|stack_cell| {
20        let stack = unsafe { &*stack_cell.get() };
21        stack.last().copied().unwrap_or(0)
22    })
23}
24
25/// A future wrapper that establishes a span context on every poll.
26///
27/// Unlike an RAII guard, this pushes the parent span ID onto the
28/// thread-local async call stack before each poll and pops it after,
29/// which is correct across yield points and executor thread migration.
30///
31/// The stack is padded to match the parent's depth so that
32/// `InstrumentedFuture::new()` inside the spawned task computes
33/// `depth = stack.len() - 1 = parent_depth + 1`, preserving the
34/// logical nesting across spawn boundaries.
35#[pin_project]
36pub struct SpanContextFuture<F> {
37    #[pin]
38    future: F,
39    parent_span: u64,
40    parent_depth: u32,
41}
42
43impl<F: Future> Future for SpanContextFuture<F> {
44    type Output = F::Output;
45
46    #[inline]
47    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
48        let this = self.project();
49        ASYNC_CALL_STACK.with(|stack_cell| {
50            let stack = unsafe { &mut *stack_cell.get() };
51            let saved_len = stack.len();
52            // Pad stack so that children see depth = parent_depth.
53            // At the spawn point the stack had length parent_depth + 1;
54            // we recreate that, then push parent_span on top.
55            let target_len = *this.parent_depth as usize + 1;
56            while stack.len() < target_len.saturating_sub(1) {
57                stack.push(0);
58            }
59            stack.push(*this.parent_span);
60            let res = this.future.poll(cx);
61            stack.truncate(saved_len);
62            res
63        })
64    }
65}
66
67/// Spawns a future on the tokio runtime while preserving the current span context.
68///
69/// This is a wrapper around `tokio::spawn` that captures the current span ID
70/// before spawning and establishes it as the parent context in the spawned task.
71/// This ensures that instrumented async functions called within the spawned task
72/// will correctly report the spawning context as their parent.
73///
74/// The parent span ID is pushed onto the thread-local async call stack before
75/// each poll and popped after, so it works correctly across yield points and
76/// executor thread migration.
77///
78/// # Example
79/// ```ignore
80/// use micromegas_tracing::prelude::*;
81///
82/// #[span_fn]
83/// async fn parent_work() {
84///     // Spans created in child_work will show parent_work as their parent
85///     spawn_with_context(child_work()).await.unwrap();
86/// }
87///
88/// #[span_fn]
89/// async fn child_work() {
90///     // This span's parent will be parent_work, not root
91/// }
92/// ```
93#[cfg(feature = "tokio")]
94pub fn spawn_with_context<F>(future: F) -> tokio::task::JoinHandle<F::Output>
95where
96    F: Future + Send + 'static,
97    F::Output: Send + 'static,
98{
99    let (parent_span, parent_depth) = ASYNC_CALL_STACK.with(|stack_cell| {
100        let stack = unsafe { &*stack_cell.get() };
101        (
102            stack.last().copied().unwrap_or(0),
103            (stack.len().saturating_sub(1)) as u32,
104        )
105    });
106    tokio::spawn(SpanContextFuture {
107        future,
108        parent_span,
109        parent_depth,
110    })
111}
112
113/// Trait for adding instrumentation to futures
114pub trait InstrumentFuture: Future + Sized {
115    /// Instrument this future with the given span metadata
116    fn instrument(self, span_desc: &'static SpanMetadata) -> InstrumentedFuture<Self> {
117        InstrumentedFuture::new(self, span_desc)
118    }
119
120    /// Internal method for named instrumentation - do not use directly.
121    /// Use the `instrument_named!` macro for method-like syntax instead.
122    #[doc(hidden)]
123    fn __instrument_named_internal(
124        self,
125        span_location: &'static SpanLocation,
126        name: &'static str,
127    ) -> InstrumentedNamedFuture<Self> {
128        InstrumentedNamedFuture::new(self, span_location, name)
129    }
130
131    /// Internal method for runtime-named instrumentation - do not use directly.
132    /// Use the `instrument_named_dyn!` macro for method-like syntax instead.
133    #[doc(hidden)]
134    fn __instrument_named_dyn_internal<N: AsRef<str>>(
135        self,
136        span_location: &'static SpanLocation,
137        name: N,
138    ) -> InstrumentedNamedFuture<Self, DynSpanName<N>> {
139        InstrumentedNamedFuture::with_name(self, span_location, DynSpanName(name))
140    }
141}
142
143impl<F: Future> InstrumentFuture for F {}
144
145/// A wrapper that instruments a future with async span tracing
146#[pin_project]
147pub struct InstrumentedFuture<F> {
148    #[pin]
149    future: F,
150    desc: &'static SpanMetadata,
151    span_id: Option<u64>,
152    /// Parent span ID captured at future creation time
153    parent: u64,
154    /// Depth captured at future creation time
155    depth: u32,
156}
157
158impl<F> InstrumentedFuture<F> {
159    /// Create a new instrumented future
160    pub fn new(future: F, desc: &'static SpanMetadata) -> Self {
161        let (parent, depth) = ASYNC_CALL_STACK.with(|stack_cell| {
162            let stack = unsafe { &*stack_cell.get() };
163            assert!(!stack.is_empty());
164            (
165                stack[stack.len() - 1],
166                (stack.len().saturating_sub(1)) as u32,
167            )
168        });
169        Self {
170            future,
171            desc,
172            span_id: None,
173            parent,
174            depth,
175        }
176    }
177}
178
179impl<F> Future for InstrumentedFuture<F>
180where
181    F: Future,
182{
183    type Output = F::Output;
184
185    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
186        let this = self.project();
187        let parent = *this.parent;
188        let depth = *this.depth;
189        ASYNC_CALL_STACK.with(|stack_cell| {
190            let stack = unsafe { &mut *stack_cell.get() };
191            assert!(!stack.is_empty());
192            match this.span_id {
193                Some(span_id) => {
194                    stack.push(*span_id);
195                }
196                None => {
197                    // Begin the async span and store the span ID
198                    let span_id = on_begin_async_scope(this.desc, parent, depth);
199                    stack.push(span_id);
200                    *this.span_id = Some(span_id);
201                }
202            }
203            let res = match this.future.poll(cx) {
204                Poll::Ready(output) => {
205                    // End the async span when the future completes
206                    if let Some(span_id) = *this.span_id {
207                        on_end_async_scope(span_id, parent, this.desc, depth);
208                    }
209                    Poll::Ready(output)
210                }
211                Poll::Pending => Poll::Pending,
212            };
213            stack.pop();
214            res
215        })
216    }
217}
218
219/// A wrapper that instruments a future with named async span tracing
220#[pin_project]
221pub struct InstrumentedNamedFuture<F, N = StaticSpanName> {
222    #[pin]
223    future: F,
224    span_location: &'static SpanLocation,
225    name: N,
226    span_id: Option<u64>,
227    /// Parent span ID captured at future creation time
228    parent: u64,
229    /// Depth captured at future creation time
230    depth: u32,
231}
232
233impl<F> InstrumentedNamedFuture<F> {
234    /// Create a new instrumented named future
235    pub fn new(future: F, span_location: &'static SpanLocation, name: &'static str) -> Self {
236        Self::with_name(future, span_location, StaticSpanName(name))
237    }
238}
239
240impl<F, N: SpanName> InstrumentedNamedFuture<F, N> {
241    /// Create a new instrumented named future
242    pub fn with_name(future: F, span_location: &'static SpanLocation, name: N) -> Self {
243        let (parent, depth) = ASYNC_CALL_STACK.with(|stack_cell| {
244            let stack = unsafe { &*stack_cell.get() };
245            assert!(!stack.is_empty());
246            (
247                stack[stack.len() - 1],
248                (stack.len().saturating_sub(1)) as u32,
249            )
250        });
251        Self {
252            future,
253            span_location,
254            name,
255            span_id: None,
256            parent,
257            depth,
258        }
259    }
260}
261
262impl<F, N: SpanName> Future for InstrumentedNamedFuture<F, N>
263where
264    F: Future,
265{
266    type Output = F::Output;
267
268    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
269        let this = self.project();
270        let parent = *this.parent;
271        let depth = *this.depth;
272        ASYNC_CALL_STACK.with(|stack_cell| {
273            let stack = unsafe { &mut *stack_cell.get() };
274            assert!(!stack.is_empty());
275            match this.span_id {
276                Some(span_id) => {
277                    stack.push(*span_id);
278                }
279                None => {
280                    // Begin the async named span and store the span ID
281                    let span_id = this.name.begin_async(this.span_location, parent, depth);
282                    stack.push(span_id);
283                    *this.span_id = Some(span_id);
284                }
285            }
286            let res = match this.future.poll(cx) {
287                Poll::Ready(output) => {
288                    // End the async named span when the future completes
289                    if let Some(span_id) = *this.span_id {
290                        this.name
291                            .end_async(span_id, parent, this.span_location, depth);
292                    }
293                    Poll::Ready(output)
294                }
295                Poll::Pending => Poll::Pending,
296            };
297            stack.pop();
298            res
299        })
300    }
301}
micromegas_tracing/spans/instrumented_future.rs

micromegas_tracing/spans/
instrumented_future.rs
