1use crate::stdlib::pin::Pin;
2use crate::stdlib::task::{Context, Poll};
3use crate::stdlib::{future::Future, marker::Sized};
4use crate::{
5 dispatcher::{self, Dispatch},
6 span::Span,
7};
8use 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
16pub 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")))]
129pub 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
229pin_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
247pin_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
266impl<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
276impl<T: Sized> Instrument for T {}
277
278impl<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")))]
321impl<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")))]
334impl<T: Sized> WithSubscriber for T {}
335
336#[cfg(feature = "std")]
337#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
338impl<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