tokio/util/
trace.rs

1cfg_rt! {
2    use std::marker::PhantomData;
3
4    #[derive(Copy, Clone)]
5    pub(crate) struct SpawnMeta<'a> {
6        /// The name of the task
7        #[cfg(all(tokio_unstable, feature = "tracing"))]
8        pub(crate) name: Option<&'a str>,
9        /// The original size of the future or function being spawned
10        #[cfg(all(tokio_unstable, feature = "tracing"))]
11        pub(crate) original_size: usize,
12        /// The source code location where the task was spawned.
13        ///
14        /// This is wrapped in a type that may be empty when `tokio_unstable` is
15        /// not enabled.
16        pub(crate) spawned_at: crate::runtime::task::SpawnLocation,
17        _pd: PhantomData<&'a ()>,
18    }
19
20    impl<'a> SpawnMeta<'a> {
21        /// Create new spawn meta with a name and original size (before possible auto-boxing)
22        #[cfg(all(tokio_unstable, feature = "tracing"))]
23        #[track_caller]
24        pub(crate) fn new(name: Option<&'a str>, original_size: usize) -> Self {
25            Self {
26                name,
27                original_size,
28                spawned_at: crate::runtime::task::SpawnLocation::capture(),
29                _pd: PhantomData,
30            }
31        }
32
33        /// Create a new unnamed spawn meta with the original size (before possible auto-boxing)
34        #[track_caller]
35        pub(crate) fn new_unnamed(original_size: usize) -> Self {
36            #[cfg(not(all(tokio_unstable, feature = "tracing")))]
37            let _original_size = original_size;
38
39            Self {
40                #[cfg(all(tokio_unstable, feature = "tracing"))]
41                name: None,
42                #[cfg(all(tokio_unstable, feature = "tracing"))]
43                original_size,
44                spawned_at: crate::runtime::task::SpawnLocation::capture(),
45                _pd: PhantomData,
46            }
47        }
48    }
49
50    cfg_trace! {
51        use core::{
52            pin::Pin,
53            task::{Context, Poll},
54        };
55        use pin_project_lite::pin_project;
56        use std::mem;
57        use std::future::Future;
58        use tracing::instrument::Instrument;
59        pub(crate) use tracing::instrument::Instrumented;
60
61        #[inline]
62        pub(crate) fn task<F>(task: F, kind: &'static str, meta: SpawnMeta<'_>, id: u64) -> Instrumented<F> {
63            fn get_span(kind: &'static str, spawn_meta: SpawnMeta<'_>, id: u64, task_size: usize) -> tracing::Span {
64                let original_size = if spawn_meta.original_size != task_size {
65                    Some(spawn_meta.original_size)
66                } else {
67                    None
68                };
69                tracing::trace_span!(
70                    target: "tokio::task",
71                    parent: None,
72                    "runtime.spawn",
73                    %kind,
74                    task.name = %spawn_meta.name.unwrap_or_default(),
75                    task.id = id,
76                    original_size.bytes = original_size,
77                    size.bytes = task_size,
78                    loc.file = spawn_meta.spawned_at.0.file(),
79                    loc.line = spawn_meta.spawned_at.0.line(),
80                    loc.col = spawn_meta.spawned_at.0.column(),
81                )
82            }
83            use tracing::instrument::Instrument;
84            let span = get_span(kind, meta, id, mem::size_of::<F>());
85            task.instrument(span)
86        }
87
88        #[inline]
89        pub(crate) fn blocking_task<Fn, Fut>(task: Fut, spawn_meta: SpawnMeta<'_>, id: u64) -> Instrumented<Fut> {
90            let fn_size = mem::size_of::<Fn>();
91            let original_size = if spawn_meta.original_size != fn_size {
92                Some(spawn_meta.original_size)
93            } else {
94                None
95            };
96
97            let span = tracing::trace_span!(
98                target: "tokio::task::blocking",
99                "runtime.spawn",
100                kind = %"blocking",
101                task.name = %spawn_meta.name.unwrap_or_default(),
102                task.id = id,
103                "fn" = %std::any::type_name::<Fn>(),
104                original_size.bytes = original_size,
105                size.bytes = fn_size,
106                loc.file = spawn_meta.spawned_at.0.file(),
107                loc.line = spawn_meta.spawned_at.0.line(),
108                loc.col = spawn_meta.spawned_at.0.column(),
109            );
110            task.instrument(span)
111
112        }
113
114        pub(crate) fn async_op<P,F>(inner: P, resource_span: tracing::Span, source: &str, poll_op_name: &'static str, inherits_child_attrs: bool) -> InstrumentedAsyncOp<F>
115        where P: FnOnce() -> F {
116            resource_span.in_scope(|| {
117                let async_op_span = tracing::trace_span!("runtime.resource.async_op", source = source, inherits_child_attrs = inherits_child_attrs);
118                let enter = async_op_span.enter();
119                let async_op_poll_span = tracing::trace_span!("runtime.resource.async_op.poll");
120                let inner = inner();
121                drop(enter);
122                let tracing_ctx = AsyncOpTracingCtx {
123                    async_op_span,
124                    async_op_poll_span,
125                    resource_span: resource_span.clone(),
126                };
127                InstrumentedAsyncOp {
128                    inner,
129                    tracing_ctx,
130                    poll_op_name,
131                }
132            })
133        }
134
135        #[derive(Debug, Clone)]
136        pub(crate) struct AsyncOpTracingCtx {
137            pub(crate) async_op_span: tracing::Span,
138            pub(crate) async_op_poll_span: tracing::Span,
139            pub(crate) resource_span: tracing::Span,
140        }
141
142
143        pin_project! {
144            #[derive(Debug, Clone)]
145            pub(crate) struct InstrumentedAsyncOp<F> {
146                #[pin]
147                pub(crate) inner: F,
148                pub(crate) tracing_ctx: AsyncOpTracingCtx,
149                pub(crate) poll_op_name: &'static str
150            }
151        }
152
153        impl<F: Future> Future for InstrumentedAsyncOp<F> {
154            type Output = F::Output;
155
156            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
157                let this = self.project();
158                let poll_op_name = &*this.poll_op_name;
159                let _res_enter = this.tracing_ctx.resource_span.enter();
160                let _async_op_enter = this.tracing_ctx.async_op_span.enter();
161                let _async_op_poll_enter = this.tracing_ctx.async_op_poll_span.enter();
162                trace_poll_op!(poll_op_name, this.inner.poll(cx))
163            }
164        }
165    }
166
167    cfg_not_trace! {
168        #[inline]
169        pub(crate) fn task<F>(task: F, _kind: &'static str, _meta: SpawnMeta<'_>, _id: u64) -> F {
170            // nop
171            task
172        }
173
174        #[inline]
175        pub(crate) fn blocking_task<Fn, Fut>(task: Fut, _spawn_meta: SpawnMeta<'_>, _id: u64) -> Fut {
176            let _ = PhantomData::<&Fn>;
177            // nop
178            task
179        }
180    }
181}
182
183cfg_time! {
184    #[track_caller]
185    pub(crate) fn caller_location() -> Option<&'static std::panic::Location<'static>> {
186        #[cfg(all(tokio_unstable, feature = "tracing"))]
187        return Some(std::panic::Location::caller());
188        #[cfg(not(all(tokio_unstable, feature = "tracing")))]
189        None
190    }
191}