tracing, tracing-futures: instrument Future inside Drop

tracing-attributes/tests/async_fn.rs

@@ -90,6 +90,8 @@ fn async_fn_only_enters_for_polls() {
         .exit(expect::span().named("test_async_fn"))
         .enter(expect::span().named("test_async_fn"))
         .exit(expect::span().named("test_async_fn"))
+        .enter(expect::span().named("test_async_fn"))
+        .exit(expect::span().named("test_async_fn"))
         .drop_span(expect::span().named("test_async_fn"))
         .only()
         .run_with_handle();
@@ -120,8 +122,12 @@ fn async_fn_nested() {
         .enter(span2.clone())
         .event(expect::event().with_fields(expect::field("nested").with_value(&true)))
         .exit(span2.clone())
+        .enter(span2.clone())
+        .exit(span2.clone())
         .drop_span(span2)
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -199,13 +205,19 @@ fn async_fn_with_async_trait() {
         .enter(span3.clone())
         .event(expect::event().with_fields(expect::field("val").with_value(&2u64)))
         .exit(span3.clone())
+        .enter(span3.clone())
+        .exit(span3.clone())
         .drop_span(span3)
         .new_span(span2.clone().with_field(expect::field("self")))
         .enter(span2.clone())
         .event(expect::event().with_fields(expect::field("val").with_value(&5u64)))
         .exit(span2.clone())
+        .enter(span2.clone())
+        .exit(span2.clone())
         .drop_span(span2)
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -256,6 +268,8 @@ fn async_fn_with_async_trait_and_fields_expressions() {
         )
         .enter(span.clone())
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -331,8 +345,12 @@ fn async_fn_with_async_trait_and_fields_expressions_with_generic_parameter() {
                 .with_field(expect::field("Self").with_value(&std::any::type_name::<TestImpl>())),
         )
         .enter(span4.clone())
+        .exit(span4.clone())
+        .enter(span4.clone())
         .exit(span4)
         .exit(span2.clone())
+        .enter(span2.clone())
+        .exit(span2.clone())
         .drop_span(span2)
         .new_span(
             span3
@@ -341,6 +359,8 @@ fn async_fn_with_async_trait_and_fields_expressions_with_generic_parameter() {
         )
         .enter(span3.clone())
         .exit(span3.clone())
+        .enter(span3.clone())
+        .exit(span3.clone())
         .drop_span(span3)
         .only()
         .run_with_handle();
@@ -382,6 +402,8 @@ fn out_of_scope_fields() {
         .new_span(span.clone())
         .enter(span.clone())
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -417,6 +439,8 @@ fn manual_impl_future() {
         .enter(span.clone())
         .event(poll_event())
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -448,6 +472,8 @@ fn manual_box_pin() {
         .enter(span.clone())
         .event(poll_event())
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();

tracing-attributes/tests/err.rs

@@ -78,6 +78,8 @@ fn test_async() {
         .enter(span.clone())
         .event(expect::event().at_level(Level::ERROR))
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -132,6 +134,8 @@ fn test_mut_async() {
         .enter(span.clone())
         .event(expect::event().at_level(Level::ERROR))
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();

tracing-attributes/tests/follows_from.rs

@@ -58,6 +58,8 @@ fn follows_from_async_test() {
         .follows_from(consequence.clone(), cause_b)
         .follows_from(consequence.clone(), cause_c)
         .enter(consequence.clone())
+        .exit(consequence.clone())
+        .enter(consequence.clone())
         .exit(consequence)
         .only()
         .run_with_handle();

tracing-attributes/tests/ret.rs

@@ -138,6 +138,8 @@ fn test_async() {
                 .at_level(Level::INFO),
         )
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();

tracing-futures/src/executor/futures_01.rs

@@ -1,27 +1,22 @@
 use crate::{Instrument, Instrumented, WithDispatch};
+use core::mem::ManuallyDrop;
 use futures_01::{
     future::{ExecuteError, Executor},
     Future,
 };
 
-macro_rules! deinstrument_err {
-    ($e:expr) => {
-        $e.map_err(|e| {
-            let kind = e.kind();
-            let future = e.into_future().inner;
-            ExecuteError::new(kind, future)
-        })
-    };
-}
-
 impl<T, F> Executor<F> for Instrumented<T>
 where
     T: Executor<Instrumented<F>>,
     F: Future<Item = (), Error = ()>,
 {
     fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
         let future = future.instrument(self.span.clone());
-        deinstrument_err!(self.inner.execute(future))
+        self.inner.execute(future).map_err(|e| {
+            let kind = e.kind();
+            let future = ManuallyDrop::into_inner(e.into_future().inner);
+            ExecuteError::new(kind, future)
+        })
     }
 }
 
@@ -32,7 +27,11 @@ where
 {
     fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
         let future = self.with_dispatch(future);
-        deinstrument_err!(self.inner.execute(future))
+        self.inner.execute(future).map_err(|e| {
+            let kind = e.kind();
+            let future = e.into_future().inner;
+            ExecuteError::new(kind, future)
+        })
     }
 }
 

tracing-futures/src/lib.rs

@@ -103,7 +103,11 @@
 use pin_project_lite::pin_project;
 
 #[cfg(feature = "std-future")]
-use core::{pin::Pin, task::Context};
+use core::{
+    mem::{self, ManuallyDrop},
+    pin::Pin,
+    task::Context,
+};
 
 #[cfg(feature = "std")]
 use tracing::{dispatch, Dispatch};
@@ -147,7 +151,10 @@ pub trait Instrument: Sized {
     ///
     /// [entered]: tracing::span::Span::enter()
     fn instrument(self, span: Span) -> Instrumented<Self> {
-        Instrumented { inner: self, span }
+        Instrumented {
+            inner: ManuallyDrop::new(self),
+            span,
+        }
     }
 
     /// Instruments this type with the [current] `Span`, returning an
@@ -243,7 +250,7 @@ pin_project! {
     #[derive(Debug, Clone)]
     pub struct Instrumented<T> {
         #[pin]
-        inner: T,
+        inner: ManuallyDrop<T>,
         span: Span,
     }
 }
@@ -287,9 +294,11 @@ impl<T: core::future::Future> core::future::Future for Instrumented<T> {
     type Output = T::Output;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll<Self::Output> {
-        let this = self.project();
+        let mut this = self.project();
         let _enter = this.span.enter();
-        this.inner.poll(cx)
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) };
+        inner.poll(cx)
     }
 }
 
@@ -346,9 +355,11 @@ impl<T: futures::Stream> futures::Stream for Instrumented<T> {
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> futures::task::Poll<Option<Self::Item>> {
-        let this = self.project();
+        let mut this = self.project();
         let _enter = this.span.enter();
-        T::poll_next(this.inner, cx)
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) };
+        T::poll_next(inner, cx)
     }
 }
 
@@ -364,33 +375,41 @@ where
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> futures::task::Poll<Result<(), Self::Error>> {
-        let this = self.project();
+        let mut this = self.project();
         let _enter = this.span.enter();
-        T::poll_ready(this.inner, cx)
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) };
+        T::poll_ready(inner, cx)
     }
 
     fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
-        let this = self.project();
+        let mut this = self.project();
         let _enter = this.span.enter();
-        T::start_send(this.inner, item)
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) };
+        T::start_send(inner, item)
     }
 
     fn poll_flush(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> futures::task::Poll<Result<(), Self::Error>> {
-        let this = self.project();
+        let mut this = self.project();
         let _enter = this.span.enter();
-        T::poll_flush(this.inner, cx)
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) };
+        T::poll_flush(inner, cx)
     }
 
     fn poll_close(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> futures::task::Poll<Result<(), Self::Error>> {
-        let this = self.project();
+        let mut this = self.project();
         let _enter = this.span.enter();
-        T::poll_close(this.inner, cx)
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) };
+        T::poll_close(inner, cx)
     }
 }
 
@@ -419,21 +438,31 @@ impl<T> Instrumented<T> {
     #[cfg(feature = "std-future")]
     #[cfg_attr(docsrs, doc(cfg(feature = "std-future")))]
     pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
-        self.project_ref().inner
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        unsafe { self.project_ref().inner.map_unchecked(|v| &**v) }
     }
 
     /// Get a pinned mutable reference to the wrapped type.
     #[cfg(feature = "std-future")]
     #[cfg_attr(docsrs, doc(cfg(feature = "std-future")))]
     pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
-        self.project().inner
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` wont move.
+        unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) }
     }
 
     /// Consumes the `Instrumented`, returning the wrapped type.
     ///
     /// Note that this drops the span.
     pub fn into_inner(self) -> T {
-        self.inner
+        let span: *const Span = &self.span;
+        let inner: *const ManuallyDrop<T> = &self.inner;
+        mem::forget(self);
+        // SAFETY: Those pointers are valid for reads, because `Drop` didn't
+        //         run, and properly aligned, because `Instrumented` isn't
+        //         `#[repr(packed)]`.
+        let _span = unsafe { span.read() };
+        let inner = unsafe { inner.read() };
+        ManuallyDrop::into_inner(inner)
     }
 }
 

tracing-futures/tests/std_future.rs

@@ -9,6 +9,8 @@ fn enter_exit_is_reasonable() {
         .exit(expect::span().named("foo"))
         .enter(expect::span().named("foo"))
         .exit(expect::span().named("foo"))
+        .enter(expect::span().named("foo"))
+        .exit(expect::span().named("foo"))
         .drop_span(expect::span().named("foo"))
         .only()
         .run_with_handle();
@@ -26,6 +28,8 @@ fn error_ends_span() {
         .exit(expect::span().named("foo"))
         .enter(expect::span().named("foo"))
         .exit(expect::span().named("foo"))
+        .enter(expect::span().named("foo"))
+        .exit(expect::span().named("foo"))
         .drop_span(expect::span().named("foo"))
         .only()
         .run_with_handle();

tracing/Cargo.toml

@@ -38,6 +38,7 @@ pin-project-lite = "0.2.9"
 
 [dev-dependencies]
 criterion = { version = "0.3.6", default_features = false }
+futures = { version = "0.3.28", default_features = false }
 log = "0.4.17"
 tracing-mock = { path = "../tracing-mock" }