1use crate::{2 dispatcher::{self, Dispatch},3 span::Span,4};5usecore::{6 future::Future,7 marker::Sized,8 mem::ManuallyDrop,9 pin::Pin,10 task::{Context, Poll},11};12usepin_project_lite::pin_project;1314/// Attaches spans to a [`std::future::Future`].15///16/// Extension trait allowing futures to be17/// instrumented with a `tracing` [span].18///19/// [span]: super::Span20pub traitInstrument: Sized {21/// Instruments this type with the provided [`Span`], returning an22 /// `Instrumented` wrapper.23 ///24 /// The attached [`Span`] will be [entered] every time the instrumented25 /// [`Future`] is polled or [`Drop`]ped.26 ///27 /// # Examples28 ///29 /// Instrumenting a future:30 ///31 /// ```rust32 /// use tracing::Instrument;33 ///34 /// # async fn doc() {35 /// let my_future = async {36 /// // ...37 /// };38 ///39 /// my_future40 /// .instrument(tracing::info_span!("my_future"))41 /// .await42 /// # }43 /// ```44 ///45 /// The [`Span::or_current`] combinator can be used in combination with46 /// `instrument` to ensure that the [current span] is attached to the47 /// 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 will62 /// // be within both "my_future" *and* "outer", since "outer" is63 /// // "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_future67 /// .instrument(tracing::debug_span!("my_future"))68 /// );69 ///70 /// // Using `Span::or_current` ensures the spawned task is instrumented71 /// // with the current span, if the new span passed to `instrument` is72 /// // 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_future77 /// .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::Future86fninstrument(self, span: Span) -> Instrumented<Self> {87 Instrumented {88 inner: ManuallyDrop::new(self),89 span,90 }91 }9293/// Instruments this type with the [current] [`Span`], returning an94 /// `Instrumented` wrapper.95 ///96 /// The attached [`Span`] will be [entered] every time the instrumented97 /// [`Future`] is polled or [`Drop`]ped.98 ///99 /// This can be used to propagate the current span when spawning a new future.100 ///101 /// # Examples102 ///103 /// ```rust104 /// 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::Span126 /// [`Future`]: std::future::Future127#[inline]128fnin_current_span(self) -> Instrumented<Self> {129self.instrument(Span::current())130 }131}132133/// Extension trait allowing futures to be instrumented with134/// a `tracing` [`Subscriber`](crate::Subscriber).135#[cfg_attr(docsrs, doc(cfg(feature ="std")))]136pub traitWithSubscriber: Sized {137/// Attaches the provided [`Subscriber`] to this type, returning a138 /// [`WithDispatch`] wrapper.139 ///140 /// The attached [`Subscriber`] will be set as the [default] when the returned141 /// [`Future`] is polled.142 ///143 /// # Examples144 ///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 /// future164 /// // Attach the other `Subscriber` to the future before awaiting it165 /// .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::Subscriber174 /// [default]: crate::dispatcher#setting-the-default-subscriber175 /// [`Future`]: std::future::Future176fnwith_subscriber<S>(self, subscriber: S) -> WithDispatch<Self>177where178S: Into<Dispatch>,179 {180 WithDispatch {181 inner:self,182 dispatcher: subscriber.into(),183 }184 }185186/// Attaches the current [default] [`Subscriber`] to this type, returning a187 /// [`WithDispatch`] wrapper.188 ///189 /// The attached `Subscriber` will be set as the [default] when the returned190 /// [`Future`] is polled.191 ///192 /// This can be used to propagate the current dispatcher context when193 /// spawning a new future that may run on a different thread.194 ///195 /// # Examples196 ///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 the206 /// // 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 may214 /// // 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 before218 /// // spawning it, ensures that the current thread's default `Subscriber` is219 /// // 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::Subscriber225 /// [default]: crate::dispatcher#setting-the-default-subscriber226 /// [`Future`]: std::future::Future227#[inline]228fnwith_current_subscriber(self) -> WithDispatch<Self> {229 WithDispatch {230 inner:self,231 dispatcher:crate::dispatcher::get_default(|default| default.clone()),232 }233 }234}235236pin_project! {237/// A [`Future`] that has been instrumented with a `tracing` [`Subscriber`].238 ///239 /// This type is returned by the [`WithSubscriber`] extension trait. See that240 /// trait's documentation for details.241 ///242 /// [`Future`]: std::future::Future243 /// [`Subscriber`]: crate::Subscriber244#[derive(Clone, Debug)]245 #[must_use ="futures do nothing unless you `.await` or poll them"]246 #[cfg_attr(docsrs, doc(cfg(feature ="std")))]247pub structWithDispatch<T> {248#[pin]249inner: T,250 dispatcher: Dispatch,251 }252}253254pin_project! {255/// A [`Future`] that has been instrumented with a `tracing` [`Span`].256 ///257 /// This type is returned by the [`Instrument`] extension trait. See that258 /// trait's documentation for details.259 ///260 /// [`Future`]: std::future::Future261 /// [`Span`]: crate::Span262#[project = InstrumentedProj]263 #[project_ref = InstrumentedProjRef]264 #[derive(Debug, Clone)]265 #[must_use ="futures do nothing unless you `.await` or poll them"]266pub structInstrumented<T> {267// `ManuallyDrop` is used here to to enter instrument `Drop` by entering268 // `Span` and executing `ManuallyDrop::drop`.269#[pin]270inner: ManuallyDrop<T>,271 span: Span,272 }273274impl<T> PinnedDropforInstrumented<T> {275fndrop(this: Pin<&mutSelf>) {276letthis = this.project();277let_enter = this.span.enter();278// SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because this isn't279 // different from wrapping `T` in `Option` and calling280 // `Pin::set(&mut this.inner, None)`, except avoiding281 // additional memory overhead.282 // 2. `ManuallyDrop::drop()` is safe, because283 // `PinnedDrop::drop()` is guaranteed to be called only284 // once.285unsafe{ ManuallyDrop::drop(this.inner.get_unchecked_mut()) }286 }287 }288}289290impl<'a, T> InstrumentedProj<'a, T> {291/// Get a mutable reference to the [`Span`] a pinned mutable reference to292 /// the wrapped type.293fnspan_and_inner_pin_mut(self) -> (&'amutSpan, Pin<&'amutT>) {294// SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move295 // and `inner` is valid, because `ManuallyDrop::drop` is called296 // only inside `Drop` of the `Instrumented`.297letinner =unsafe{self.inner.map_unchecked_mut(|v|&mut **v) };298 (self.span, inner)299 }300}301302impl<'a, T> InstrumentedProjRef<'a, T> {303/// Get a reference to the [`Span`] a pinned reference to the wrapped type.304fnspan_and_inner_pin_ref(self) -> (&'aSpan, Pin<&'aT>) {305// SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move306 // and `inner` is valid, because `ManuallyDrop::drop` is called307 // only inside `Drop` of the `Instrumented`.308letinner =unsafe{self.inner.map_unchecked(|v|&**v) };309 (self.span, inner)310 }311}312313// === impl Instrumented ===314315impl<T: Future> FutureforInstrumented<T> {316typeOutput = T::Output;317318fnpoll(self: Pin<&mutSelf>, cx:&mutContext<'_>) -> Poll<Self::Output> {319let(span, inner) =self.project().span_and_inner_pin_mut();320let_enter = span.enter();321 inner.poll(cx)322 }323}324325impl<T: Sized> InstrumentforT {}326327impl<T> Instrumented<T> {328/// Borrows the `Span` that this type is instrumented by.329pub fnspan(&self) ->&Span {330&self.span331 }332333/// Mutably borrows the `Span` that this type is instrumented by.334pub fnspan_mut(&mutself) ->&mutSpan {335&mutself.span336 }337338/// Borrows the wrapped type.339pub fninner(&self) ->&T {340&self.inner341 }342343/// Mutably borrows the wrapped type.344pub fninner_mut(&mutself) ->&mutT {345&mutself.inner346 }347348/// Get a pinned reference to the wrapped type.349pub fninner_pin_ref(self: Pin<&Self>) -> Pin<&T> {350self.project_ref().span_and_inner_pin_ref().1351}352353/// Get a pinned mutable reference to the wrapped type.354pub fninner_pin_mut(self: Pin<&mutSelf>) -> Pin<&mutT> {355self.project().span_and_inner_pin_mut().1356}357358/// Consumes the `Instrumented`, returning the wrapped type.359 ///360 /// Note that this drops the span.361pub fninto_inner(self) -> T {362// To manually destructure `Instrumented` without `Drop`, we363 // move it into a ManuallyDrop and use pointers to its fields364letthis = ManuallyDrop::new(self);365letspan:*constSpan =&this.span;366letinner:*constManuallyDrop<T> =&this.inner;367// SAFETY: Those pointers are valid for reads, because `Drop` didn't368 // run, and properly aligned, because `Instrumented` isn't369 // `#[repr(packed)]`.370let_span =unsafe{ span.read() };371letinner =unsafe{ inner.read() };372 ManuallyDrop::into_inner(inner)373 }374}375376// === impl WithDispatch ===377378#[cfg(feature ="std")]379#[cfg_attr(docsrs, doc(cfg(feature ="std")))]380impl<T: Future> FutureforWithDispatch<T> {381typeOutput = T::Output;382383fnpoll(self: Pin<&mutSelf>, cx:&mutContext<'_>) -> Poll<Self::Output> {384letthis =self.project();385letdispatcher = this.dispatcher;386letfuture = this.inner;387let_default = dispatcher::set_default(dispatcher);388 future.poll(cx)389 }390}391392#[cfg_attr(docsrs, doc(cfg(feature ="std")))]393impl<T: Sized> WithSubscriberforT {}394395#[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.399pub fndispatcher(&self) ->&Dispatch {400&self.dispatcher401 }402403/// Borrows the wrapped type.404pub fninner(&self) ->&T {405&self.inner406 }407408/// Mutably borrows the wrapped type.409pub fninner_mut(&mutself) ->&mutT {410&mutself.inner411 }412413/// Get a pinned reference to the wrapped type.414pub fninner_pin_ref(self: Pin<&Self>) -> Pin<&T> {415self.project_ref().inner416 }417418/// Get a pinned mutable reference to the wrapped type.419pub fninner_pin_mut(self: Pin<&mutSelf>) -> Pin<&mutT> {420self.project().inner421 }422423/// Consumes the `Instrumented`, returning the wrapped type.424 ///425 /// Note that this drops the span.426pub fninto_inner(self) -> T {427self.inner428 }429}