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