tokio-rs · davidbarsky · Sep 25, 2020 · Jul 12, 2020 · Jul 12, 2020 · Jul 12, 2020
@@ -31,8 +31,7 @@ use std::env;
 use std::error::Error;
 use std::net::SocketAddr;
 
-use tracing::{debug, info, info_span, trace_span, warn};
-use tracing_futures::Instrument;
+use tracing::{debug, info, info_span, trace_span, warn, Instrument as _};
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {

@@ -32,9 +32,7 @@ use tokio::{
     self, io,
     net::{TcpListener, TcpStream},
 };
-use tracing::{debug, debug_span, info, warn};
-use tracing_attributes::instrument;
-use tracing_futures::Instrument;
+use tracing::{debug, debug_span, info, instrument, warn, Instrument as _};
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 

@@ -6,8 +6,7 @@ use hyper::{
     Body, Server,
 };
 use std::str;
-use tracing::{debug, info, span, Level};
-use tracing_futures::Instrument;
+use tracing::{debug, info, span, Instrument as _, Level};
 
 async fn echo(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
     let span = span!(

@@ -8,8 +8,7 @@
 /// cargo run --example tokio-spawny-thing
 /// ```
 use futures::future::try_join_all;
-use tracing::{debug, info, instrument, span, Level};
-use tracing_futures::Instrument;
+use tracing::{debug, info, instrument, span, Instrument as _, Level};
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 

@@ -42,8 +42,7 @@ use std::{
 };
 use tokio::{time, try_join};
 use tower::{Service, ServiceBuilder, ServiceExt};
-use tracing::{self, debug, error, info, span, trace, warn, Level, Span};
-use tracing_futures::Instrument;
+use tracing::{self, debug, error, info, span, trace, warn, Instrument as _, Level, Span};
 use tracing_subscriber::{filter::EnvFilter, reload::Handle};
 use tracing_tower::{request_span, request_span::make};
 

@@ -45,7 +45,6 @@ quote = "1"
 
 [dev-dependencies]
 tracing = { path = "../tracing", version = "0.1" }
-tracing-futures = { path = "../tracing-futures", version = "0.2" }
 tokio-test = { version = "0.2.0" }
 tracing-core = { path = "../tracing-core", version = "0.1"}
 async-trait = "0.1"

@@ -386,7 +386,7 @@ fn gen_body(
         if err {
             quote_spanned! {block.span()=>
                 let __tracing_attr_span = #span;
-                tracing_futures::Instrument::instrument(async move {
+                tracing::Instrument::instrument(async move {
                     match async move { #block }.await {
                         Ok(x) => Ok(x),
                         Err(e) => {
@@ -399,7 +399,7 @@ fn gen_body(
         } else {
             quote_spanned!(block.span()=>
                 let __tracing_attr_span = #span;
-                    tracing_futures::Instrument::instrument(
+                    tracing::Instrument::instrument(
                         async move { #block },
                         __tracing_attr_span
                     )

@@ -87,8 +87,10 @@ pub(crate) mod stdlib;
 #[cfg(feature = "std-future")]
 use crate::stdlib::{pin::Pin, task::Context};
 
-use tracing::dispatcher;
-use tracing::{Dispatch, Span};
+#[cfg(feature = "std")]
+use tracing::{dispatcher, Dispatch};
+
+use tracing::Span;
 
 /// Implementations for `Instrument`ed future executors.
 pub mod executor;
@@ -650,8 +652,7 @@ mod tests {
                 .drop_span(span::mock().named("foo"))
                 .run_with_handle();
             with_default(subscriber, || {
-                stream::iter(&[1, 2, 3])
-                    .instrument(tracing::trace_span!("foo"))
+                Instrument::instrument(stream::iter(&[1, 2, 3]), tracing::trace_span!("foo"))
                     .for_each(|_| future::ready(()))
                     .now_or_never()
                     .unwrap();
@@ -671,8 +672,7 @@ mod tests {
                 .drop_span(span::mock().named("foo"))
                 .run_with_handle();
             with_default(subscriber, || {
-                sink::drain()
-                    .instrument(tracing::trace_span!("foo"))
+                Instrument::instrument(sink::drain(), tracing::trace_span!("foo"))
                     .send(1u8)
                     .now_or_never()
                     .unwrap()

@@ -1,8 +1,8 @@
 mod support;
 use support::*;
 
+use tracing::Instrument;
 use tracing::{subscriber::with_default, Level};
-use tracing_futures::Instrument;
 
 #[test]
 fn enter_exit_is_reasonable() {

@@ -4,7 +4,7 @@ use futures::future::Future;
 use std::marker::PhantomData;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use tracing_futures::Instrument;
+use tracing::Instrument;
 
 #[derive(Debug)]
 pub struct Service<S, R, G = fn(&R) -> tracing::Span>
@@ -236,7 +236,7 @@ where
 {
     type Response = S::Response;
     type Error = S::Error;
-    type Future = tracing_futures::Instrumented<S::Future>;
+    type Future = tracing::instrument::Instrumented<S::Future>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.inner.poll_ready(cx)

@@ -31,6 +31,7 @@ tracing-core = { path = "../tracing-core", version = "0.1.11", default-features
 log = { version = "0.4", optional = true }
 tracing-attributes = { path = "../tracing-attributes", version = "0.1.9", optional = true }
 cfg-if = "0.1.10"
+pin-project-lite = "0.1"
 
 [dev-dependencies]
 futures = "0.1"

@@ -0,0 +1,201 @@
+use crate::stdlib::pin::Pin;
+use crate::stdlib::task::{Context, Poll};
+use crate::stdlib::{future::Future, marker::Sized};
+use crate::{dispatcher, span::Span, Dispatch};
+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]: https://docs.rs/tracing/latest/tracing/span/index.html
+pub trait Instrument: Sized {
+    /// Instruments this type with the provided `Span`, returning an
+    /// `Instrumented` wrapper.
+    ///
+    /// If the instrumented type is a future, stream, or sink, the attached `Span`
+    /// will be [entered] every time it is polled. If the instrumented type
+    /// is a future executor, every future spawned on that executor will be
+    /// instrumented by the attached `Span`.
+    ///
+    /// # Examples
+    ///
+    /// Instrumenting a future:
+    ///
+    // TODO: ignored until async-await is stable...
+    /// ```rust,ignore
+    /// use tracing::Instrument;
+    ///
+    /// # async fn doc() {
+    /// let my_future = async {
+    ///     // ...
+    /// };
+    ///
+    /// my_future
+    ///     .instrument(tracing::info_span!("my_future"))
+    ///     .await
+    /// # }
+    /// ```
+    ///
+    /// [entered]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.enter
+    fn instrument(self, span: Span) -> Instrumented<Self> {
+        Instrumented { inner: self, span }
+    }
+
+    /// Instruments this type with the [current] `Span`, returning an
+    /// `Instrumented` wrapper.
+    ///
+    /// If the instrumented type is a future, stream, or sink, the attached `Span`
+    /// will be [entered] every time it is polled. If the instrumented type
+    /// is a future executor, every future spawned on that executor will be
+    /// instrumented by the attached `Span`.
+    ///
+    /// This can be used to propagate the current span when spawning a new future.
+    ///
+    /// # Examples
+    ///
+    // TODO: ignored until async-await is stable...
+    /// ```rust,ignore
+    /// use tracing::Instrument;
+    ///
+    /// # 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]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.current
+    /// [entered]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.enter
+    #[inline]
+    fn in_current_span(self) -> Instrumented<Self> {
+        self.instrument(Span::current())
+    }
+}
+
+/// Extension trait allowing futures to be instrumented with
+/// a `tracing` [`Subscriber`].
+///
+/// [`Subscriber`]: https://docs.rs/tracing/latest/tracing/subscriber/trait.Subscriber.html
+pub trait WithSubscriber: Sized {
+    /// Attaches the provided [`Subscriber`] to this type, returning a
+    /// `WithDispatch` wrapper.
+    ///
+    /// When the wrapped type is a future, stream, or sink, the attached
+    /// subscriber will be set as the [default] while it is being polled.
+    /// When the wrapped type is an executor, the subscriber will be set as the
+    /// default for any futures spawned on that executor.
+    ///
+    /// [`Subscriber`]: https://docs.rs/tracing/latest/tracing/subscriber/trait.Subscriber.html
+    /// [default]: https://docs.rs/tracing/latest/tracing/dispatcher/index.html#setting-the-default-subscriber
+    fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self>
+    where
+        S: Into<Dispatch>,
+    {
+        WithDispatch {
+            inner: self,
+            dispatch: subscriber.into(),
+        }
+    }
+
+    /// Attaches the current [default] [`Subscriber`] to this type, returning a
+    /// `WithDispatch` wrapper.
+    ///
+    /// When the wrapped type is a future, stream, or sink, the attached
+    /// subscriber will be set as the [default] while it is being polled.
+    /// When the wrapped type is an executor, the subscriber will be set as the
+    /// default for any futures spawned on that executor.
+    ///
+    /// This can be used to propagate the current dispatcher context when
+    /// spawning a new future.
+    ///
+    /// [`Subscriber`]: https://docs.rs/tracing/latest/tracing/subscriber/trait.Subscriber.html
+    /// [default]: https://docs.rs/tracing/latest/tracing/dispatcher/index.html#setting-the-default-subscriber
+    #[inline]
+    fn with_current_subscriber(self) -> WithDispatch<Self> {
+        WithDispatch {
+            inner: self,
+            dispatch: dispatcher::get_default(|default| default.clone()),
+        }
+    }
+}
+
+pin_project! {
+    /// A future, stream, sink, or executor that has been instrumented with a
+    /// `tracing` subscriber.
+    #[derive(Clone, Debug)]
+    pub struct WithDispatch<T> {
+        #[pin]
+        inner: T,
+        dispatch: Dispatch,
+    }
+}
+
+pin_project! {
+    /// A future that has been instrumented with a `tracing` span.
+    #[derive(Debug, Clone)]
+    pub struct Instrumented<T> {
+        #[pin]
+        inner: T,
+        span: Span,
+    }
+}
+
+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
+    }
+}
@@ -834,7 +834,10 @@ pub use log;
 #[doc(hidden)]
 use tracing_core::*;
 
-pub use self::{dispatcher::Dispatch, event::Event, field::Value, subscriber::Subscriber};
+pub use self::{
+    dispatcher::Dispatch, event::Event, field::Value, instrument::Instrument,
+    subscriber::Subscriber,
+};
 
 #[doc(hidden)]
 pub use self::span::Id;
@@ -858,6 +861,8 @@ mod macros;
 
 pub mod dispatcher;
 pub mod field;
+/// Attach a span to a `std::future::Future`.
+pub mod instrument;
 pub mod level_filters;
 pub mod span;
 pub(crate) mod stdlib;