From f1385275ac4ec7706a4c493a99fb5f2f25c74afa Mon Sep 17 00:00:00 2001
From: zohnannor <zohnannor@gmail.com>
Date: Thu, 30 Mar 2023 17:44:50 +0300
Subject: [PATCH 1/4] tracing: instrument futures with `Span`s in `Drop`

---
 tracing/src/instrument.rs | 33 ++++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/tracing/src/instrument.rs b/tracing/src/instrument.rs
index 98f0da9bd1..8cbf693563 100644
--- a/tracing/src/instrument.rs
+++ b/tracing/src/instrument.rs
@@ -1,4 +1,5 @@
 use crate::span::Span;
+use core::hint::unreachable_unchecked;
 use core::pin::Pin;
 use core::task::{Context, Poll};
 use core::{future::Future, marker::Sized};
@@ -80,7 +81,10 @@ pub trait Instrument: Sized {
     /// [disabled]: super::Span::is_disabled()
     /// [`Future`]: std::future::Future
     fn instrument(self, span: Span) -> Instrumented<Self> {
-        Instrumented { inner: self, span }
+        Instrumented {
+            inner: Some(self),
+            span,
+        }
     }
 
     /// Instruments this type with the [current] [`Span`], returning an
@@ -289,9 +293,17 @@ pin_project! {
     #[must_use = "futures do nothing unless you `.await` or poll them"]
     pub struct Instrumented<T> {
         #[pin]
-        inner: T,
+        inner: Option<T>,
         span: Span,
     }
+
+    impl<T> PinnedDrop for Instrumented<T> {
+        fn drop(this: Pin<&mut Self>) {
+            let mut this = this.project();
+            let _enter = this.span.enter();
+            this.inner.set(None);
+        }
+    }
 }
 
 // === impl Instrumented ===
@@ -302,7 +314,8 @@ impl<T: Future> Future for Instrumented<T> {
     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)
+        // SAFETY:
+        unsafe { this.inner.as_pin_mut().unwrap_unchecked().poll(cx) }
     }
 }
 
@@ -321,32 +334,34 @@ impl<T> Instrumented<T> {
 
     /// Borrows the wrapped type.
     pub fn inner(&self) -> &T {
-        &self.inner
+        unsafe { self.inner.as_ref().unwrap_unchecked() }
     }
 
     /// Mutably borrows the wrapped type.
     pub fn inner_mut(&mut self) -> &mut T {
-        &mut self.inner
+        unsafe { self.inner.as_mut().unwrap_unchecked() }
     }
 
     /// Get a pinned reference to the wrapped type.
     pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
-        self.project_ref().inner
+        unsafe { self.project_ref().inner.as_pin_ref().unwrap_unchecked() }
     }
 
     /// Get a pinned mutable reference to the wrapped type.
     pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
-        self.project().inner
+        unsafe { self.project().inner.as_pin_mut().unwrap_unchecked() }
     }
 
     /// Consumes the `Instrumented`, returning the wrapped type.
     ///
     /// Note that this drops the span.
-    pub fn into_inner(self) -> T {
-        self.inner
+    pub fn into_inner(mut self) -> T {
+        unsafe { self.inner.take().unwrap_unchecked() }
     }
 }
 
+unsafe fn unwrap_() {}
+
 // === impl WithDispatch ===
 
 #[cfg(feature = "std")]

From 12ad199c6b9ae0d22749a4328c18e664c85a0a47 Mon Sep 17 00:00:00 2001
From: zohnannor <zohnannor@gmail.com>
Date: Thu, 30 Mar 2023 17:46:38 +0300
Subject: [PATCH 2/4] tracing-futures: instrument futures with `Span`s in
 `Drop`

---
 tracing-futures/src/executor/futures_01.rs |  49 ++++---
 tracing-futures/src/executor/futures_03.rs |  11 +-
 tracing-futures/src/lib.rs                 | 145 +++++++++++++--------
 3 files changed, 129 insertions(+), 76 deletions(-)

diff --git a/tracing-futures/src/executor/futures_01.rs b/tracing-futures/src/executor/futures_01.rs
index 56ba6e3c42..73225e65ff 100644
--- a/tracing-futures/src/executor/futures_01.rs
+++ b/tracing-futures/src/executor/futures_01.rs
@@ -4,16 +4,6 @@ use futures_01::{
     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>>,
@@ -21,7 +11,11 @@ where
 {
     fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
         let future = future.instrument(self.span.clone());
-        deinstrument_err!(self.inner.execute(future))
+        self.inner
+            .as_ref()
+            .unwrap()
+            .execute(future)
+            .map_err(|e| ExecuteError::new(e.kind(), e.into_future().inner.take().unwrap()))
     }
 }
 
@@ -32,7 +26,9 @@ 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| ExecuteError::new(e.kind(), e.into_future().inner))
     }
 }
 
@@ -59,7 +55,7 @@ mod tokio {
         ) -> Result<(), SpawnError> {
             // TODO: get rid of double box somehow?
             let future = Box::new(future.instrument(self.span.clone()));
-            self.inner.spawn(future)
+            self.inner.as_mut().unwrap().spawn(future)
         }
     }
 
@@ -68,11 +64,14 @@ mod tokio {
         T: TypedExecutor<Instrumented<F>>,
     {
         fn spawn(&mut self, future: F) -> Result<(), SpawnError> {
-            self.inner.spawn(future.instrument(self.span.clone()))
+            self.inner
+                .as_mut()
+                .unwrap()
+                .spawn(future.instrument(self.span.clone()))
         }
 
         fn status(&self) -> Result<(), SpawnError> {
-            self.inner.status()
+            self.inner.as_ref().unwrap().status()
         }
     }
 
@@ -90,7 +89,7 @@ mod tokio {
             F: Future<Item = (), Error = ()> + Send + 'static,
         {
             let future = future.instrument(self.span.clone());
-            self.inner.spawn(future);
+            self.inner.as_mut().unwrap().spawn(future);
             self
         }
 
@@ -116,7 +115,7 @@ mod tokio {
             E: Send + 'static,
         {
             let future = future.instrument(self.span.clone());
-            self.inner.block_on(future)
+            self.inner.as_mut().unwrap().block_on(future)
         }
 
         /// Return an instrumented handle to the runtime's executor.
@@ -127,7 +126,11 @@ mod tokio {
         /// `tokio::runtime::TaskExecutor`, but instruments the spawned
         /// futures prior to spawning them.
         pub fn executor(&self) -> Instrumented<TaskExecutor> {
-            self.inner.executor().instrument(self.span.clone())
+            self.inner
+                .as_ref()
+                .unwrap()
+                .executor()
+                .instrument(self.span.clone())
         }
     }
 
@@ -141,7 +144,7 @@ mod tokio {
             F: Future<Item = (), Error = ()> + 'static,
         {
             let future = future.instrument(self.span.clone());
-            self.inner.spawn(future);
+            self.inner.as_mut().unwrap().spawn(future);
             self
         }
 
@@ -176,7 +179,7 @@ mod tokio {
             E: 'static,
         {
             let future = future.instrument(self.span.clone());
-            self.inner.block_on(future)
+            self.inner.as_mut().unwrap().block_on(future)
         }
 
         /// Get a new instrumented handle to spawn futures on the single-threaded
@@ -189,7 +192,11 @@ mod tokio {
         /// `tokio::runtime::current_thread::Handle`, but instruments the spawned
         /// futures prior to spawning them.
         pub fn handle(&self) -> Instrumented<current_thread::Handle> {
-            self.inner.handle().instrument(self.span.clone())
+            self.inner
+                .as_ref()
+                .unwrap()
+                .handle()
+                .instrument(self.span.clone())
         }
     }
 
diff --git a/tracing-futures/src/executor/futures_03.rs b/tracing-futures/src/executor/futures_03.rs
index 94a75c3cc6..6e7ba0fb8e 100644
--- a/tracing-futures/src/executor/futures_03.rs
+++ b/tracing-futures/src/executor/futures_03.rs
@@ -15,7 +15,10 @@ where
     /// tasks.
     fn spawn_obj(&self, future: FutureObj<'static, ()>) -> Result<(), SpawnError> {
         let future = future.instrument(self.span.clone());
-        self.inner.spawn_obj(FutureObj::new(Box::new(future)))
+        self.inner
+            .as_ref()
+            .unwrap()
+            .spawn_obj(FutureObj::new(Box::new(future)))
     }
 
     /// Determines whether the executor is able to spawn new tasks.
@@ -26,7 +29,7 @@ where
     /// not guaranteed, to yield an error.
     #[inline]
     fn status(&self) -> Result<(), SpawnError> {
-        self.inner.status()
+        self.inner.as_ref().unwrap().status()
     }
 }
 
@@ -74,6 +77,8 @@ where
     fn spawn_local_obj(&self, future: LocalFutureObj<'static, ()>) -> Result<(), SpawnError> {
         let future = future.instrument(self.span.clone());
         self.inner
+            .as_ref()
+            .unwrap()
             .spawn_local_obj(LocalFutureObj::new(Box::new(future)))
     }
 
@@ -85,7 +90,7 @@ where
     /// not guaranteed, to yield an error.
     #[inline]
     fn status_local(&self) -> Result<(), SpawnError> {
-        self.inner.status_local()
+        self.inner.as_ref().unwrap().status_local()
     }
 }
 
diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs
index 5af13a0a5b..c35d533973 100644
--- a/tracing-futures/src/lib.rs
+++ b/tracing-futures/src/lib.rs
@@ -102,6 +102,7 @@
 #[cfg(feature = "std-future")]
 use pin_project_lite::pin_project;
 
+use core::hint::unreachable_unchecked;
 #[cfg(feature = "std-future")]
 use core::{pin::Pin, task::Context};
 
@@ -147,7 +148,10 @@ pub trait Instrument: Sized {
     ///
     /// [entered]: tracing::span::Span::enter()
     fn instrument(self, span: Span) -> Instrumented<Self> {
-        Instrumented { inner: self, span }
+        Instrumented {
+            inner: Some(self),
+            span,
+        }
     }
 
     /// Instruments this type with the [current] `Span`, returning an
@@ -243,16 +247,24 @@ pin_project! {
     #[derive(Debug, Clone)]
     pub struct Instrumented<T> {
         #[pin]
-        inner: T,
+        inner: Option<T>,
         span: Span,
     }
+
+    impl<T> PinnedDrop for Instrumented<T> {
+        fn drop(this: Pin<&mut Self>) {
+            let mut this = this.project();
+            let _enter = this.span.enter();
+            this.inner.set(None);
+        }
+    }
 }
 
 /// A future, stream, sink, or executor that has been instrumented with a `tracing` span.
 #[cfg(not(feature = "std-future"))]
 #[derive(Debug, Clone)]
 pub struct Instrumented<T> {
-    inner: T,
+    inner: Option<T>,
     span: Span,
 }
 
@@ -289,7 +301,7 @@ impl<T: core::future::Future> core::future::Future for Instrumented<T> {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll<Self::Output> {
         let this = self.project();
         let _enter = this.span.enter();
-        this.inner.poll(cx)
+        this.inner.as_pin_mut().unwrap().poll(cx)
     }
 }
 
@@ -301,7 +313,7 @@ impl<T: futures_01::Future> futures_01::Future for Instrumented<T> {
 
     fn poll(&mut self) -> futures_01::Poll<Self::Item, Self::Error> {
         let _enter = self.span.enter();
-        self.inner.poll()
+        unsafe { self.inner.as_mut().unwrap_unchecked().poll() }
     }
 }
 
@@ -313,7 +325,7 @@ impl<T: futures_01::Stream> futures_01::Stream for Instrumented<T> {
 
     fn poll(&mut self) -> futures_01::Poll<Option<Self::Item>, Self::Error> {
         let _enter = self.span.enter();
-        self.inner.poll()
+        unsafe { self.inner.as_mut().unwrap_unchecked().poll() }
     }
 }
 
@@ -328,12 +340,12 @@ impl<T: futures_01::Sink> futures_01::Sink for Instrumented<T> {
         item: Self::SinkItem,
     ) -> futures_01::StartSend<Self::SinkItem, Self::SinkError> {
         let _enter = self.span.enter();
-        self.inner.start_send(item)
+        unsafe { self.inner.as_mut().unwrap_unchecked().start_send(item) }
     }
 
     fn poll_complete(&mut self) -> futures_01::Poll<(), Self::SinkError> {
         let _enter = self.span.enter();
-        self.inner.poll_complete()
+        unsafe { self.inner.as_mut().unwrap_unchecked().poll_complete() }
     }
 }
 
@@ -346,9 +358,12 @@ 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)
+        T::poll_next(
+            unsafe { this.inner.as_mut().as_pin_mut().unwrap_unchecked() },
+            cx,
+        )
     }
 }
 
@@ -364,33 +379,45 @@ 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)
+        T::poll_ready(
+            unsafe { this.inner.as_mut().as_pin_mut().unwrap_unchecked() },
+            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)
+        T::start_send(
+            unsafe { this.inner.as_mut().as_pin_mut().unwrap_unchecked() },
+            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)
+        T::poll_flush(
+            unsafe { this.inner.as_mut().as_pin_mut().unwrap_unchecked() },
+            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)
+        T::poll_close(
+            unsafe { this.inner.as_mut().as_pin_mut().unwrap_unchecked() },
+            cx,
+        )
     }
 }
 
@@ -407,33 +434,33 @@ impl<T> Instrumented<T> {
 
     /// Borrows the wrapped type.
     pub fn inner(&self) -> &T {
-        &self.inner
+        unsafe { self.inner.as_ref().unwrap_unchecked() }
     }
 
     /// Mutably borrows the wrapped type.
     pub fn inner_mut(&mut self) -> &mut T {
-        &mut self.inner
+        unsafe { self.inner.as_mut().unwrap_unchecked() }
     }
 
     /// Get a pinned reference to the wrapped type.
     #[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
+        unsafe { self.project_ref().inner.as_pin_ref().unwrap_unchecked() }
     }
 
     /// 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
+        unsafe { self.project().inner.as_pin_mut().unwrap_unchecked() }
     }
 
     /// Consumes the `Instrumented`, returning the wrapped type.
     ///
     /// Note that this drops the span.
-    pub fn into_inner(self) -> T {
-        self.inner
+    pub fn into_inner(mut self) -> T {
+        unsafe { self.inner.take().unwrap_unchecked() }
     }
 }
 
@@ -565,12 +592,15 @@ mod tests {
 
         #[test]
         fn future_enter_exit_is_reasonable() {
+            let span = expect::span().named("foo");
             let (collector, handle) = collector::mock()
-                .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"))
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .drop_span(span)
                 .only()
                 .run_with_handle();
             with_default(collector, || {
@@ -584,12 +614,15 @@ mod tests {
 
         #[test]
         fn future_error_ends_span() {
+            let span = expect::span().named("foo");
             let (collector, handle) = collector::mock()
-                .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"))
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .drop_span(span)
                 .only()
                 .run_with_handle();
             with_default(collector, || {
@@ -604,16 +637,19 @@ mod tests {
 
         #[test]
         fn stream_enter_exit_is_reasonable() {
+            let span = expect::span().named("foo");
             let (collector, handle) = collector::mock()
-                .enter(expect::span().named("foo"))
-                .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"))
-                .enter(expect::span().named("foo"))
-                .exit(expect::span().named("foo"))
-                .drop_span(expect::span().named("foo"))
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .drop_span(span)
                 .run_with_handle();
             with_default(collector, || {
                 stream::iter_ok::<_, ()>(&[1, 2, 3])
@@ -665,16 +701,21 @@ mod tests {
 
         #[test]
         fn stream_enter_exit_is_reasonable() {
+            let span = expect::span().named("foo");
             let (collector, handle) = collector::mock()
-                .enter(expect::span().named("foo"))
-                .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"))
-                .enter(expect::span().named("foo"))
-                .exit(expect::span().named("foo"))
-                .drop_span(expect::span().named("foo"))
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .enter(span.clone())
+                .exit(span.clone())
+                .drop_span(span)
                 .run_with_handle();
             with_default(collector, || {
                 Instrument::instrument(stream::iter(&[1, 2, 3]), tracing::trace_span!("foo"))

From 12d8405e8c99bf6c9b1e2f5760c3a2e6fdbdbad8 Mon Sep 17 00:00:00 2001
From: zohnannor <zohnannor@gmail.com>
Date: Thu, 30 Mar 2023 17:49:29 +0300
Subject: [PATCH 3/4] tracing-attributes: fix tests to account for new
 `Instrumented` behavior

---
 tracing-attributes/tests/async_fn.rs     | 39 ++++++++++++++++++++----
 tracing-attributes/tests/err.rs          |  4 +++
 tracing-attributes/tests/follows_from.rs |  5 ++-
 tracing-attributes/tests/ret.rs          |  2 ++
 4 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/tracing-attributes/tests/async_fn.rs b/tracing-attributes/tests/async_fn.rs
index 38675a451a..d2dac9dabb 100644
--- a/tracing-attributes/tests/async_fn.rs
+++ b/tracing-attributes/tests/async_fn.rs
@@ -83,14 +83,17 @@ fn repro_1831_2() -> impl Future<Output = Result<(), Infallible>> {
 
 #[test]
 fn async_fn_only_enters_for_polls() {
+    let span = expect::span().named("test_async_fn");
     let (collector, handle) = collector::mock()
-        .new_span(expect::span().named("test_async_fn"))
-        .enter(expect::span().named("test_async_fn"))
+        .new_span(span.clone())
+        .enter(span.clone())
         .event(expect::event().with_fields(expect::field("awaiting").with_value(&true)))
-        .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"))
+        .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
+        .drop_span(span)
         .only()
         .run_with_handle();
     with_default(collector, || {
@@ -120,8 +123,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 +206,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 +269,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 +346,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 +360,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 +403,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 +440,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 +473,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();
diff --git a/tracing-attributes/tests/err.rs b/tracing-attributes/tests/err.rs
index ffd30b3742..8777502b26 100644
--- a/tracing-attributes/tests/err.rs
+++ b/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();
diff --git a/tracing-attributes/tests/follows_from.rs b/tracing-attributes/tests/follows_from.rs
index 266f7b59a3..ce9b6d32a7 100644
--- a/tracing-attributes/tests/follows_from.rs
+++ b/tracing-attributes/tests/follows_from.rs
@@ -58,7 +58,10 @@ fn follows_from_async_test() {
         .follows_from(consequence.clone(), cause_b)
         .follows_from(consequence.clone(), cause_c)
         .enter(consequence.clone())
-        .exit(consequence)
+        .exit(consequence.clone())
+        .enter(consequence.clone())
+        .exit(consequence.clone())
+        .drop_span(consequence)
         .only()
         .run_with_handle();
 
diff --git a/tracing-attributes/tests/ret.rs b/tracing-attributes/tests/ret.rs
index 0a2d2ed191..b5c4b4e07b 100644
--- a/tracing-attributes/tests/ret.rs
+++ b/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();

From 8af97e1941fe0ec22cef32c25c46234a07157dbd Mon Sep 17 00:00:00 2001
From: zohnannor <zohnannor@gmail.com>
Date: Thu, 30 Mar 2023 17:51:10 +0300
Subject: [PATCH 4/4] tracing-futures: fix tests to account for new
 `Instrumented` behavior

---
 tracing-futures/tests/std_future.rs | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/tracing-futures/tests/std_future.rs b/tracing-futures/tests/std_future.rs
index cd9656f153..0645bb29fc 100644
--- a/tracing-futures/tests/std_future.rs
+++ b/tracing-futures/tests/std_future.rs
@@ -4,12 +4,15 @@ use tracing_mock::*;
 
 #[test]
 fn enter_exit_is_reasonable() {
+    let span = expect::span().named("foo");
     let (collector, handle) = collector::mock()
-        .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"))
+        .enter(span.clone())
+        .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
+        .drop_span(span)
         .only()
         .run_with_handle();
     with_default(collector, || {
@@ -21,12 +24,15 @@ fn enter_exit_is_reasonable() {
 
 #[test]
 fn error_ends_span() {
+    let span = expect::span().named("foo");
     let (collector, handle) = collector::mock()
-        .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"))
+        .enter(span.clone())
+        .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
+        .drop_span(span)
         .only()
         .run_with_handle();
     with_default(collector, || {