1 | use crate::stdlib::pin::Pin; |
2 | use crate::stdlib::task::{Context, Poll}; |
3 | use crate::stdlib::{future::Future, marker::Sized}; |
4 | use crate::{ |
5 | dispatcher::{self, Dispatch}, |
6 | span::Span, |
7 | }; |
8 | use pin_project_lite::pin_project; |
9 | |
10 | /// Attaches spans to a [`std::future::Future`]. |
11 | /// |
12 | /// Extension trait allowing futures to be |
13 | /// instrumented with a `tracing` [span]. |
14 | /// |
15 | /// [span]: super::Span |
16 | pub trait Instrument: Sized { |
17 | /// Instruments this type with the provided [`Span`], returning an |
18 | /// `Instrumented` wrapper. |
19 | /// |
20 | /// The attached [`Span`] will be [entered] every time the instrumented |
21 | /// [`Future`] is polled. |
22 | /// |
23 | /// # Examples |
24 | /// |
25 | /// Instrumenting a future: |
26 | /// |
27 | /// ```rust |
28 | /// use tracing::Instrument; |
29 | /// |
30 | /// # async fn doc() { |
31 | /// let my_future = async { |
32 | /// // ... |
33 | /// }; |
34 | /// |
35 | /// my_future |
36 | /// .instrument(tracing::info_span!("my_future" )) |
37 | /// .await |
38 | /// # } |
39 | /// ``` |
40 | /// |
41 | /// The [`Span::or_current`] combinator can be used in combination with |
42 | /// `instrument` to ensure that the [current span] is attached to the |
43 | /// future if the span passed to `instrument` is [disabled]: |
44 | /// |
45 | /// ``` |
46 | /// use tracing::Instrument; |
47 | /// # mod tokio { |
48 | /// # pub(super) fn spawn(_: impl std::future::Future) {} |
49 | /// # } |
50 | /// |
51 | /// let my_future = async { |
52 | /// // ... |
53 | /// }; |
54 | /// |
55 | /// let outer_span = tracing::info_span!("outer" ).entered(); |
56 | /// |
57 | /// // If the "my_future" span is enabled, then the spawned task will |
58 | /// // be within both "my_future" *and* "outer", since "outer" is |
59 | /// // "my_future"'s parent. However, if "my_future" is disabled, |
60 | /// // the spawned task will *not* be in any span. |
61 | /// tokio::spawn( |
62 | /// my_future |
63 | /// .instrument(tracing::debug_span!("my_future" )) |
64 | /// ); |
65 | /// |
66 | /// // Using `Span::or_current` ensures the spawned task is instrumented |
67 | /// // with the current span, if the new span passed to `instrument` is |
68 | /// // not enabled. This means that if the "my_future" span is disabled, |
69 | /// // the spawned task will still be instrumented with the "outer" span: |
70 | /// # let my_future = async {}; |
71 | /// tokio::spawn( |
72 | /// my_future |
73 | /// .instrument(tracing::debug_span!("my_future" ).or_current()) |
74 | /// ); |
75 | /// ``` |
76 | /// |
77 | /// [entered]: super::Span::enter() |
78 | /// [`Span::or_current`]: super::Span::or_current() |
79 | /// [current span]: super::Span::current() |
80 | /// [disabled]: super::Span::is_disabled() |
81 | /// [`Future`]: std::future::Future |
82 | fn instrument(self, span: Span) -> Instrumented<Self> { |
83 | Instrumented { inner: self, span } |
84 | } |
85 | |
86 | /// Instruments this type with the [current] [`Span`], returning an |
87 | /// `Instrumented` wrapper. |
88 | /// |
89 | /// The attached [`Span`] will be [entered] every time the instrumented |
90 | /// [`Future`] is polled. |
91 | /// |
92 | /// This can be used to propagate the current span when spawning a new future. |
93 | /// |
94 | /// # Examples |
95 | /// |
96 | /// ```rust |
97 | /// use tracing::Instrument; |
98 | /// |
99 | /// # mod tokio { |
100 | /// # pub(super) fn spawn(_: impl std::future::Future) {} |
101 | /// # } |
102 | /// # async fn doc() { |
103 | /// let span = tracing::info_span!("my_span" ); |
104 | /// let _enter = span.enter(); |
105 | /// |
106 | /// // ... |
107 | /// |
108 | /// let future = async { |
109 | /// tracing::debug!("this event will occur inside `my_span`" ); |
110 | /// // ... |
111 | /// }; |
112 | /// tokio::spawn(future.in_current_span()); |
113 | /// # } |
114 | /// ``` |
115 | /// |
116 | /// [current]: super::Span::current() |
117 | /// [entered]: super::Span::enter() |
118 | /// [`Span`]: crate::Span |
119 | /// [`Future`]: std::future::Future |
120 | #[inline ] |
121 | fn in_current_span(self) -> Instrumented<Self> { |
122 | self.instrument(Span::current()) |
123 | } |
124 | } |
125 | |
126 | /// Extension trait allowing futures to be instrumented with |
127 | /// a `tracing` [`Subscriber`](crate::Subscriber). |
128 | #[cfg_attr (docsrs, doc(cfg(feature = "std" )))] |
129 | pub trait WithSubscriber: Sized { |
130 | /// Attaches the provided [`Subscriber`] to this type, returning a |
131 | /// [`WithDispatch`] wrapper. |
132 | /// |
133 | /// The attached [`Subscriber`] will be set as the [default] when the returned |
134 | /// [`Future`] is polled. |
135 | /// |
136 | /// # Examples |
137 | /// |
138 | /// ``` |
139 | /// # use tracing::subscriber::NoSubscriber as MySubscriber; |
140 | /// # use tracing::subscriber::NoSubscriber as MyOtherSubscriber; |
141 | /// # async fn docs() { |
142 | /// use tracing::instrument::WithSubscriber; |
143 | /// |
144 | /// // Set the default `Subscriber` |
145 | /// let _default = tracing::subscriber::set_default(MySubscriber::default()); |
146 | /// |
147 | /// tracing::info!("this event will be recorded by the default `Subscriber`" ); |
148 | /// |
149 | /// // Create a different `Subscriber` and attach it to a future. |
150 | /// let other_subscriber = MyOtherSubscriber::default(); |
151 | /// let future = async { |
152 | /// tracing::info!("this event will be recorded by the other `Subscriber`" ); |
153 | /// // ... |
154 | /// }; |
155 | /// |
156 | /// future |
157 | /// // Attach the other `Subscriber` to the future before awaiting it |
158 | /// .with_subscriber(other_subscriber) |
159 | /// .await; |
160 | /// |
161 | /// // Once the future has completed, we return to the default `Subscriber`. |
162 | /// tracing::info!("this event will be recorded by the default `Subscriber`" ); |
163 | /// # } |
164 | /// ``` |
165 | /// |
166 | /// [`Subscriber`]: super::Subscriber |
167 | /// [default]: crate::dispatcher#setting-the-default-subscriber |
168 | /// [`Future`]: std::future::Future |
169 | fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self> |
170 | where |
171 | S: Into<Dispatch>, |
172 | { |
173 | WithDispatch { |
174 | inner: self, |
175 | dispatcher: subscriber.into(), |
176 | } |
177 | } |
178 | |
179 | /// Attaches the current [default] [`Subscriber`] to this type, returning a |
180 | /// [`WithDispatch`] wrapper. |
181 | /// |
182 | /// The attached `Subscriber` will be set as the [default] when the returned |
183 | /// [`Future`] is polled. |
184 | /// |
185 | /// This can be used to propagate the current dispatcher context when |
186 | /// spawning a new future that may run on a different thread. |
187 | /// |
188 | /// # Examples |
189 | /// |
190 | /// ``` |
191 | /// # mod tokio { |
192 | /// # pub(super) fn spawn(_: impl std::future::Future) {} |
193 | /// # } |
194 | /// # use tracing::subscriber::NoSubscriber as MySubscriber; |
195 | /// # async fn docs() { |
196 | /// use tracing::instrument::WithSubscriber; |
197 | /// |
198 | /// // Using `set_default` (rather than `set_global_default`) sets the |
199 | /// // default `Subscriber` for *this* thread only. |
200 | /// let _default = tracing::subscriber::set_default(MySubscriber::default()); |
201 | /// |
202 | /// let future = async { |
203 | /// // ... |
204 | /// }; |
205 | /// |
206 | /// // If a multi-threaded async runtime is in use, this spawned task may |
207 | /// // run on a different thread, in a different default `Subscriber`'s context. |
208 | /// tokio::spawn(future); |
209 | /// |
210 | /// // However, calling `with_current_subscriber` on the future before |
211 | /// // spawning it, ensures that the current thread's default `Subscriber` is |
212 | /// // propagated to the spawned task, regardless of where it executes: |
213 | /// # let future = async { }; |
214 | /// tokio::spawn(future.with_current_subscriber()); |
215 | /// # } |
216 | /// ``` |
217 | /// [`Subscriber`]: super::Subscriber |
218 | /// [default]: crate::dispatcher#setting-the-default-subscriber |
219 | /// [`Future`]: std::future::Future |
220 | #[inline ] |
221 | fn with_current_subscriber(self) -> WithDispatch<Self> { |
222 | WithDispatch { |
223 | inner: self, |
224 | dispatcher: crate::dispatcher::get_default(|default| default.clone()), |
225 | } |
226 | } |
227 | } |
228 | |
229 | pin_project! { |
230 | /// A [`Future`] that has been instrumented with a `tracing` [`Subscriber`]. |
231 | /// |
232 | /// This type is returned by the [`WithSubscriber`] extension trait. See that |
233 | /// trait's documentation for details. |
234 | /// |
235 | /// [`Future`]: std::future::Future |
236 | /// [`Subscriber`]: crate::Subscriber |
237 | #[derive(Clone, Debug)] |
238 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
239 | #[cfg_attr(docsrs, doc(cfg(feature = "std" )))] |
240 | pub struct WithDispatch<T> { |
241 | #[pin] |
242 | inner: T, |
243 | dispatcher: Dispatch, |
244 | } |
245 | } |
246 | |
247 | pin_project! { |
248 | /// A [`Future`] that has been instrumented with a `tracing` [`Span`]. |
249 | /// |
250 | /// This type is returned by the [`Instrument`] extension trait. See that |
251 | /// trait's documentation for details. |
252 | /// |
253 | /// [`Future`]: std::future::Future |
254 | /// [`Span`]: crate::Span |
255 | #[derive(Debug, Clone)] |
256 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
257 | pub struct Instrumented<T> { |
258 | #[pin] |
259 | inner: T, |
260 | span: Span, |
261 | } |
262 | } |
263 | |
264 | // === impl Instrumented === |
265 | |
266 | impl<T: Future> Future for Instrumented<T> { |
267 | type Output = T::Output; |
268 | |
269 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
270 | let this: Projection<'_, T> = self.project(); |
271 | let _enter: Entered<'_> = this.span.enter(); |
272 | this.inner.poll(cx) |
273 | } |
274 | } |
275 | |
276 | impl<T: Sized> Instrument for T {} |
277 | |
278 | impl<T> Instrumented<T> { |
279 | /// Borrows the `Span` that this type is instrumented by. |
280 | pub fn span(&self) -> &Span { |
281 | &self.span |
282 | } |
283 | |
284 | /// Mutably borrows the `Span` that this type is instrumented by. |
285 | pub fn span_mut(&mut self) -> &mut Span { |
286 | &mut self.span |
287 | } |
288 | |
289 | /// Borrows the wrapped type. |
290 | pub fn inner(&self) -> &T { |
291 | &self.inner |
292 | } |
293 | |
294 | /// Mutably borrows the wrapped type. |
295 | pub fn inner_mut(&mut self) -> &mut T { |
296 | &mut self.inner |
297 | } |
298 | |
299 | /// Get a pinned reference to the wrapped type. |
300 | pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { |
301 | self.project_ref().inner |
302 | } |
303 | |
304 | /// Get a pinned mutable reference to the wrapped type. |
305 | pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { |
306 | self.project().inner |
307 | } |
308 | |
309 | /// Consumes the `Instrumented`, returning the wrapped type. |
310 | /// |
311 | /// Note that this drops the span. |
312 | pub fn into_inner(self) -> T { |
313 | self.inner |
314 | } |
315 | } |
316 | |
317 | // === impl WithDispatch === |
318 | |
319 | #[cfg (feature = "std" )] |
320 | #[cfg_attr (docsrs, doc(cfg(feature = "std" )))] |
321 | impl<T: Future> Future for WithDispatch<T> { |
322 | type Output = T::Output; |
323 | |
324 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
325 | let this: Projection<'_, T> = self.project(); |
326 | let dispatcher: &mut Dispatch = this.dispatcher; |
327 | let future: Pin<&mut T> = this.inner; |
328 | let _default: DefaultGuard = dispatcher::set_default(dispatcher); |
329 | future.poll(cx) |
330 | } |
331 | } |
332 | |
333 | #[cfg_attr (docsrs, doc(cfg(feature = "std" )))] |
334 | impl<T: Sized> WithSubscriber for T {} |
335 | |
336 | #[cfg (feature = "std" )] |
337 | #[cfg_attr (docsrs, doc(cfg(feature = "std" )))] |
338 | impl<T> WithDispatch<T> { |
339 | /// Borrows the [`Dispatch`] that is entered when this type is polled. |
340 | pub fn dispatcher(&self) -> &Dispatch { |
341 | &self.dispatcher |
342 | } |
343 | |
344 | /// Borrows the wrapped type. |
345 | pub fn inner(&self) -> &T { |
346 | &self.inner |
347 | } |
348 | |
349 | /// Mutably borrows the wrapped type. |
350 | pub fn inner_mut(&mut self) -> &mut T { |
351 | &mut self.inner |
352 | } |
353 | |
354 | /// Get a pinned reference to the wrapped type. |
355 | pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { |
356 | self.project_ref().inner |
357 | } |
358 | |
359 | /// Get a pinned mutable reference to the wrapped type. |
360 | pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { |
361 | self.project().inner |
362 | } |
363 | |
364 | /// Consumes the `Instrumented`, returning the wrapped type. |
365 | /// |
366 | /// Note that this drops the span. |
367 | pub fn into_inner(self) -> T { |
368 | self.inner |
369 | } |
370 | } |
371 | |