1 | cfg_trace! { |
2 | cfg_rt! { |
3 | use core::{ |
4 | pin::Pin, |
5 | task::{Context, Poll}, |
6 | }; |
7 | use pin_project_lite::pin_project; |
8 | use std::future::Future; |
9 | pub(crate) use tracing::instrument::Instrumented; |
10 | |
11 | #[inline ] |
12 | #[track_caller ] |
13 | pub(crate) fn task<F>(task: F, kind: &'static str, name: Option<&str>, id: u64) -> Instrumented<F> { |
14 | #[track_caller ] |
15 | fn get_span(kind: &'static str, name: Option<&str>, id: u64) -> tracing::Span { |
16 | let location = std::panic::Location::caller(); |
17 | tracing::trace_span!( |
18 | target: "tokio::task" , |
19 | parent: None, |
20 | "runtime.spawn" , |
21 | %kind, |
22 | task.name = %name.unwrap_or_default(), |
23 | task.id = id, |
24 | loc.file = location.file(), |
25 | loc.line = location.line(), |
26 | loc.col = location.column(), |
27 | ) |
28 | } |
29 | use tracing::instrument::Instrument; |
30 | let span = get_span(kind, name, id); |
31 | task.instrument(span) |
32 | } |
33 | |
34 | 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> |
35 | where P: FnOnce() -> F { |
36 | resource_span.in_scope(|| { |
37 | let async_op_span = tracing::trace_span!("runtime.resource.async_op" , source = source, inherits_child_attrs = inherits_child_attrs); |
38 | let enter = async_op_span.enter(); |
39 | let async_op_poll_span = tracing::trace_span!("runtime.resource.async_op.poll" ); |
40 | let inner = inner(); |
41 | drop(enter); |
42 | let tracing_ctx = AsyncOpTracingCtx { |
43 | async_op_span, |
44 | async_op_poll_span, |
45 | resource_span: resource_span.clone(), |
46 | }; |
47 | InstrumentedAsyncOp { |
48 | inner, |
49 | tracing_ctx, |
50 | poll_op_name, |
51 | } |
52 | }) |
53 | } |
54 | |
55 | #[derive (Debug, Clone)] |
56 | pub(crate) struct AsyncOpTracingCtx { |
57 | pub(crate) async_op_span: tracing::Span, |
58 | pub(crate) async_op_poll_span: tracing::Span, |
59 | pub(crate) resource_span: tracing::Span, |
60 | } |
61 | |
62 | |
63 | pin_project! { |
64 | #[derive(Debug, Clone)] |
65 | pub(crate) struct InstrumentedAsyncOp<F> { |
66 | #[pin] |
67 | pub(crate) inner: F, |
68 | pub(crate) tracing_ctx: AsyncOpTracingCtx, |
69 | pub(crate) poll_op_name: &'static str |
70 | } |
71 | } |
72 | |
73 | impl<F: Future> Future for InstrumentedAsyncOp<F> { |
74 | type Output = F::Output; |
75 | |
76 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
77 | let this = self.project(); |
78 | let poll_op_name = &*this.poll_op_name; |
79 | let _res_enter = this.tracing_ctx.resource_span.enter(); |
80 | let _async_op_enter = this.tracing_ctx.async_op_span.enter(); |
81 | let _async_op_poll_enter = this.tracing_ctx.async_op_poll_span.enter(); |
82 | trace_poll_op!(poll_op_name, this.inner.poll(cx)) |
83 | } |
84 | } |
85 | } |
86 | } |
87 | cfg_time! { |
88 | #[track_caller ] |
89 | pub(crate) fn caller_location() -> Option<&'static std::panic::Location<'static>> { |
90 | #[cfg (all(tokio_unstable, feature = "tracing" ))] |
91 | return Some(std::panic::Location::caller()); |
92 | #[cfg (not(all(tokio_unstable, feature = "tracing" )))] |
93 | None |
94 | } |
95 | } |
96 | |
97 | cfg_not_trace! { |
98 | cfg_rt! { |
99 | #[inline ] |
100 | pub(crate) fn task<F>(task: F, _: &'static str, _name: Option<&str>, _: u64) -> F { |
101 | // nop |
102 | task |
103 | } |
104 | } |
105 | } |
106 | |