aboutsummaryrefslogtreecommitdiff
path: root/src/instrument.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/instrument.rs')
-rw-r--r--src/instrument.rs370
1 files changed, 370 insertions, 0 deletions
diff --git a/src/instrument.rs b/src/instrument.rs
new file mode 100644
index 0000000..46e5f57
--- /dev/null
+++ b/src/instrument.rs
@@ -0,0 +1,370 @@
+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
+ }
+}
generated by cgit v1.2.3 (git 2.25.1) at 2024-05-24 03:18:28 +0000