1 files changed, 72 insertions, 13 deletions

diff --git a/vendor/tracing/src/instrument.rs b/vendor/tracing/src/instrument.rs
index 46e5f579c..87c74016b 100644
--- a/vendor/tracing/src/instrument.rs
+++ b/vendor/tracing/src/instrument.rs
@@ -1,10 +1,14 @@
-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 core::{
+    future::Future,
+    marker::Sized,
+    mem::ManuallyDrop,
+    pin::Pin,
+    task::{Context, Poll},
+};
 use pin_project_lite::pin_project;
 
 /// Attaches spans to a [`std::future::Future`].
@@ -18,7 +22,7 @@ pub trait Instrument: Sized {
     /// `Instrumented` wrapper.
     ///
     /// The attached [`Span`] will be [entered] every time the instrumented
-    /// [`Future`] is polled.
+    /// [`Future`] is polled or [`Drop`]ped.
     ///
     /// # Examples
     ///
@@ -80,14 +84,17 @@ 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: ManuallyDrop::new(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.
+    /// [`Future`] is polled or [`Drop`]ped.
     ///
     /// This can be used to propagate the current span when spawning a new future.
     ///
@@ -252,13 +259,55 @@ pin_project! {
     ///
     /// [`Future`]: std::future::Future
     /// [`Span`]: crate::Span
+    #[project = InstrumentedProj]
+    #[project_ref = InstrumentedProjRef]
     #[derive(Debug, Clone)]
     #[must_use = "futures do nothing unless you `.await` or poll them"]
     pub struct Instrumented<T> {
+        // `ManuallyDrop` is used here to to enter instrument `Drop` by entering
+        // `Span` and executing `ManuallyDrop::drop`.
         #[pin]
-        inner: T,
+        inner: ManuallyDrop<T>,
         span: Span,
     }
+
+    impl<T> PinnedDrop for Instrumented<T> {
+        fn drop(this: Pin<&mut Self>) {
+            let this = this.project();
+            let _enter = this.span.enter();
+            // SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because this isn't
+            //             different from wrapping `T` in `Option` and calling
+            //             `Pin::set(&mut this.inner, None)`, except avoiding
+            //             additional memory overhead.
+            //         2. `ManuallyDrop::drop()` is safe, because
+            //            `PinnedDrop::drop()` is guaranteed to be called only
+            //            once.
+            unsafe { ManuallyDrop::drop(this.inner.get_unchecked_mut()) }
+        }
+    }
+}
+
+impl<'a, T> InstrumentedProj<'a, T> {
+    /// Get a mutable reference to the [`Span`] a pinned mutable reference to
+    /// the wrapped type.
+    fn span_and_inner_pin_mut(self) -> (&'a mut Span, Pin<&'a mut T>) {
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
+        //         and `inner` is valid, because `ManuallyDrop::drop` is called
+        //         only inside `Drop` of the `Instrumented`.
+        let inner = unsafe { self.inner.map_unchecked_mut(|v| &mut **v) };
+        (self.span, inner)
+    }
+}
+
+impl<'a, T> InstrumentedProjRef<'a, T> {
+    /// Get a reference to the [`Span`] a pinned reference to the wrapped type.
+    fn span_and_inner_pin_ref(self) -> (&'a Span, Pin<&'a T>) {
+        // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
+        //         and `inner` is valid, because `ManuallyDrop::drop` is called
+        //         only inside `Drop` of the `Instrumented`.
+        let inner = unsafe { self.inner.map_unchecked(|v| &**v) };
+        (self.span, inner)
+    }
 }
 
 // === impl Instrumented ===
@@ -267,9 +316,9 @@ 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)
+        let (span, inner) = self.project().span_and_inner_pin_mut();
+        let _enter = span.enter();
+        inner.poll(cx)
     }
 }
 
@@ -298,19 +347,29 @@ impl<T> Instrumented<T> {
 
     /// Get a pinned reference to the wrapped type.
     pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
-        self.project_ref().inner
+        self.project_ref().span_and_inner_pin_ref().1
     }
 
     /// Get a pinned mutable reference to the wrapped type.
     pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
-        self.project().inner
+        self.project().span_and_inner_pin_mut().1
     }
 
     /// Consumes the `Instrumented`, returning the wrapped type.
     ///
     /// Note that this drops the span.
     pub fn into_inner(self) -> T {
-        self.inner
+        // To manually destructure `Instrumented` without `Drop`, we
+        // move it into a ManuallyDrop and use pointers to its fields
+        let this = ManuallyDrop::new(self);
+        let span: *const Span = &this.span;
+        let inner: *const ManuallyDrop<T> = &this.inner;
+        // 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)
     }
 }