tracing/src/instrument.rs

use crate::stdlib::pin::Pin;
use crate::stdlib::task::{Context, Poll};
use crate::stdlib::{future::Future, marker::Sized};
use crate::{
    dispatcher::{self, Dispatch},
    span::Span,
};
use pin_project_lite::pin_project;

/// Attaches spans to a [`std::future::Future`].
///
/// Extension trait allowing futures to be
/// instrumented with a `tracing` [span].
///
/// [span]: super::Span
pub trait Instrument: Sized {
    /// Instruments this type with the provided [`Span`], returning an
    /// `Instrumented` wrapper.
    ///
    /// The attached [`Span`] will be [entered] every time the instrumented
    /// [`Future`] is polled.
    ///
    /// # Examples
    ///
    /// Instrumenting a future:
    ///
    /// ```rust
    /// use tracing::Instrument;
    ///
    /// # async fn doc() {
    /// let my_future = async {
    ///     // ...
    /// };
    ///
    /// my_future
    ///     .instrument(tracing::info_span!("my_future"))
    ///     .await
    /// # }
    /// ```
    ///
    /// The [`Span::or_current`] combinator can be used in combination with
    /// `instrument` to ensure that the [current span] is attached to the
    /// future if the span passed to `instrument` is [disabled]:
    ///
    /// ```
    /// use tracing::Instrument;
    /// # mod tokio {
    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
    /// # }
    ///
    /// let my_future = async {
    ///     // ...
    /// };
    ///
    /// let outer_span = tracing::info_span!("outer").entered();
    ///
    /// // If the "my_future" span is enabled, then the spawned task will
    /// // be within both "my_future" *and* "outer", since "outer" is
    /// // "my_future"'s parent. However, if "my_future" is disabled,
    /// // the spawned task will *not* be in any span.
    /// tokio::spawn(
    ///     my_future
    ///         .instrument(tracing::debug_span!("my_future"))
    /// );
    ///
    /// // Using `Span::or_current` ensures the spawned task is instrumented
    /// // with the current span, if the new span passed to `instrument` is
    /// // not enabled. This means that if the "my_future"  span is disabled,
    /// // the spawned task will still be instrumented with the "outer" span:
    /// # let my_future = async {};
    /// tokio::spawn(
    ///    my_future
    ///         .instrument(tracing::debug_span!("my_future").or_current())
    /// );
    /// ```
    ///
    /// [entered]: super::Span::enter()
    /// [`Span::or_current`]: super::Span::or_current()
    /// [current span]: super::Span::current()
    /// [disabled]: super::Span::is_disabled()
    /// [`Future`]: std::future::Future
    fn instrument(self, span: Span) -> Instrumented<Self> {
        Instrumented { inner: self, span }
    }

    /// Instruments this type with the [current] [`Span`], returning an
    /// `Instrumented` wrapper.
    ///
    /// The attached [`Span`] will be [entered] every time the instrumented
    /// [`Future`] is polled.
    ///
    /// This can be used to propagate the current span when spawning a new future.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tracing::Instrument;
    ///
    /// # mod tokio {
    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
    /// # }
    /// # async fn doc() {
    /// let span = tracing::info_span!("my_span");
    /// let _enter = span.enter();
    ///
    /// // ...
    ///
    /// let future = async {
    ///     tracing::debug!("this event will occur inside `my_span`");
    ///     // ...
    /// };
    /// tokio::spawn(future.in_current_span());
    /// # }
    /// ```
    ///
    /// [current]: super::Span::current()
    /// [entered]: super::Span::enter()
    /// [`Span`]: crate::Span
    /// [`Future`]: std::future::Future
    #[inline]
    fn in_current_span(self) -> Instrumented<Self> {
        self.instrument(Span::current())
    }
}

/// Extension trait allowing futures to be instrumented with
/// a `tracing` [`Subscriber`](crate::Subscriber).
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
pub trait WithSubscriber: Sized {
    /// Attaches the provided [`Subscriber`] to this type, returning a
    /// [`WithDispatch`] wrapper.
    ///
    /// The attached [`Subscriber`] will be set as the [default] when the returned
    /// [`Future`] is polled.
    ///
    /// # Examples
    ///
    /// ```
    /// # use tracing::subscriber::NoSubscriber as MySubscriber;
    /// # use tracing::subscriber::NoSubscriber as MyOtherSubscriber;
    /// # async fn docs() {
    /// use tracing::instrument::WithSubscriber;
    ///
    /// // Set the default `Subscriber`
    /// let _default = tracing::subscriber::set_default(MySubscriber::default());
    ///
    /// tracing::info!("this event will be recorded by the default `Subscriber`");
    ///
    /// // Create a different `Subscriber` and attach it to a future.
    /// let other_subscriber = MyOtherSubscriber::default();
    /// let future = async {
    ///     tracing::info!("this event will be recorded by the other `Subscriber`");
    ///     // ...
    /// };
    ///
    /// future
    ///     // Attach the other `Subscriber` to the future before awaiting it
    ///     .with_subscriber(other_subscriber)
    ///     .await;
    ///
    /// // Once the future has completed, we return to the default `Subscriber`.
    /// tracing::info!("this event will be recorded by the default `Subscriber`");
    /// # }
    /// ```
    ///
    /// [`Subscriber`]: super::Subscriber
    /// [default]: crate::dispatcher#setting-the-default-subscriber
    /// [`Future`]: std::future::Future
    fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self>
    where
        S: Into<Dispatch>,
    {
        WithDispatch {
            inner: self,
            dispatcher: subscriber.into(),
        }
    }

    /// Attaches the current [default] [`Subscriber`] to this type, returning a
    /// [`WithDispatch`] wrapper.
    ///
    /// The attached `Subscriber` will be set as the [default] when the returned
    /// [`Future`] is polled.
    ///
    /// This can be used to propagate the current dispatcher context when
    /// spawning a new future that may run on a different thread.
    ///
    /// # Examples
    ///
    /// ```
    /// # mod tokio {
    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
    /// # }
    /// # use tracing::subscriber::NoSubscriber as MySubscriber;
    /// # async fn docs() {
    /// use tracing::instrument::WithSubscriber;
    ///
    /// // Using `set_default` (rather than `set_global_default`) sets the
    /// // default `Subscriber` for *this* thread only.
    /// let _default = tracing::subscriber::set_default(MySubscriber::default());
    ///
    /// let future = async {
    ///     // ...
    /// };
    ///
    /// // If a multi-threaded async runtime is in use, this spawned task may
    /// // run on a different thread, in a different default `Subscriber`'s context.
    /// tokio::spawn(future);
    ///
    /// // However, calling `with_current_subscriber` on the future before
    /// // spawning it, ensures that the current thread's default `Subscriber` is
    /// // propagated to the spawned task, regardless of where it executes:
    /// # let future = async { };
    /// tokio::spawn(future.with_current_subscriber());
    /// # }
    /// ```
    /// [`Subscriber`]: super::Subscriber
    /// [default]: crate::dispatcher#setting-the-default-subscriber
    /// [`Future`]: std::future::Future
    #[inline]
    fn with_current_subscriber(self) -> WithDispatch<Self> {
        WithDispatch {
            inner: self,
            dispatcher: crate::dispatcher::get_default(|default| default.clone()),
        }
    }
}

pin_project! {
    /// A [`Future`] that has been instrumented with a `tracing` [`Subscriber`].
    ///
    /// This type is returned by the [`WithSubscriber`] extension trait. See that
    /// trait's documentation for details.
    ///
    /// [`Future`]: std::future::Future
    /// [`Subscriber`]: crate::Subscriber
    #[derive(Clone, Debug)]
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
    pub struct WithDispatch<T> {
        #[pin]
        inner: T,
        dispatcher: Dispatch,
    }
}

pin_project! {
    /// A [`Future`] that has been instrumented with a `tracing` [`Span`].
    ///
    /// This type is returned by the [`Instrument`] extension trait. See that
    /// trait's documentation for details.
    ///
    /// [`Future`]: std::future::Future
    /// [`Span`]: crate::Span
    #[derive(Debug, Clone)]
    #[must_use = "futures do nothing unless you `.await` or poll them"]
    pub struct Instrumented<T> {
        #[pin]
        inner: T,
        span: Span,
    }
}

// === impl Instrumented ===

impl<T: Future> Future for Instrumented<T> {
    type Output = T::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let _enter = this.span.enter();
        this.inner.poll(cx)
    }
}

impl<T: Sized> Instrument for T {}

impl<T> Instrumented<T> {
    /// Borrows the `Span` that this type is instrumented by.
    pub fn span(&self) -> &Span {
        &self.span
    }

    /// Mutably borrows the `Span` that this type is instrumented by.
    pub fn span_mut(&mut self) -> &mut Span {
        &mut self.span
    }

    /// Borrows the wrapped type.
    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Mutably borrows the wrapped type.
    pub fn inner_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Get a pinned reference to the wrapped type.
    pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
        self.project_ref().inner
    }

    /// Get a pinned mutable reference to the wrapped type.
    pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
        self.project().inner
    }

    /// Consumes the `Instrumented`, returning the wrapped type.
    ///
    /// Note that this drops the span.
    pub fn into_inner(self) -> T {
        self.inner
    }
}

// === impl WithDispatch ===

#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl<T: Future> Future for WithDispatch<T> {
    type Output = T::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let dispatcher = this.dispatcher;
        let future = this.inner;
        let _default = dispatcher::set_default(dispatcher);
        future.poll(cx)
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl<T: Sized> WithSubscriber for T {}

#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl<T> WithDispatch<T> {
    /// Borrows the [`Dispatch`] that is entered when this type is polled.
    pub fn dispatcher(&self) -> &Dispatch {
        &self.dispatcher
    }

    /// Borrows the wrapped type.
    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Mutably borrows the wrapped type.
    pub fn inner_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Get a pinned reference to the wrapped type.
    pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
        self.project_ref().inner
    }

    /// Get a pinned mutable reference to the wrapped type.
    pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
        self.project().inner
    }

    /// Consumes the `Instrumented`, returning the wrapped type.
    ///
    /// Note that this drops the span.
    pub fn into_inner(self) -> T {
        self.inner
    }
}