From 9e2dd0f5c838e277a55a894f89aa12e6f313d89e Mon Sep 17 00:00:00 2001
From: Ilya Salauyeu <47687266+ilslv@users.noreply.github.com>
Date: Wed, 19 Apr 2023 00:55:46 +0200
Subject: [PATCH] tracing, tracing-futures: instrument `Future` inside `Drop`
 (#2562)

## Motivation

Currently it is not possible to disable ANSI in `fmt::Subscriber`
without enabling the "ansi" crate feature. This makes it difficult for
users to implement interoperable settings that are controllable with
crate features without having to pull in the dependencies "ansi" does.

I hit this while writing an application with multiple logging options
set during compile-time and I wanted to cut down on dependencies if
possible.

## Solution

This changes `fmt::Subscriber::with_ansi()` to not require the "ansi"
feature flag. This way, `with_ansi(false)` can be called even when the
"ansi" feature is disabled. Calling `with_ansi(true)` when the "ansi"
feature is not enabled will panic in debug mode, or print a warning if
debug assertions are disabled.

Co-authored-by: daxpedda <daxpedda@gmail.com>
Co-authored-by: Eliza Weisman <eliza@buoyant.io>
---
 tracing-attributes/tests/async_fn.rs |  6 ++++++
 tracing-attributes/tests/err.rs      |  4 ++++
 tracing/src/instrument.rs            | 22 ++++++++++++++++++++--
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/tracing-attributes/tests/async_fn.rs b/tracing-attributes/tests/async_fn.rs
index 60d772ebd3..7c23b963a4 100644
--- a/tracing-attributes/tests/async_fn.rs
+++ b/tracing-attributes/tests/async_fn.rs
@@ -124,6 +124,8 @@ fn async_fn_nested() {
         .exit(span2.clone())
         .enter(span2.clone())
         .exit(span2.clone())
+        .enter(span2.clone())
+        .exit(span2.clone())
         .drop_span(span2)
         .exit(span.clone())
         .enter(span.clone())
@@ -207,6 +209,8 @@ fn async_fn_with_async_trait() {
         .exit(span3.clone())
         .enter(span3.clone())
         .exit(span3.clone())
+        .enter(span3.clone())
+        .exit(span3.clone())
         .drop_span(span3)
         .new_span(span2.clone().with_fields(expect::field("self")))
         .enter(span2.clone())
@@ -214,6 +218,8 @@ fn async_fn_with_async_trait() {
         .exit(span2.clone())
         .enter(span2.clone())
         .exit(span2.clone())
+        .enter(span2.clone())
+        .exit(span2.clone())
         .drop_span(span2)
         .exit(span.clone())
         .enter(span.clone())
diff --git a/tracing-attributes/tests/err.rs b/tracing-attributes/tests/err.rs
index c270141ed1..d5d508f545 100644
--- a/tracing-attributes/tests/err.rs
+++ b/tracing-attributes/tests/err.rs
@@ -81,6 +81,8 @@ fn test_async() {
         .exit(span.clone())
         .enter(span.clone())
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
@@ -137,6 +139,8 @@ fn test_mut_async() {
         .exit(span.clone())
         .enter(span.clone())
         .exit(span.clone())
+        .enter(span.clone())
+        .exit(span.clone())
         .drop_span(span)
         .only()
         .run_with_handle();
diff --git a/tracing/src/instrument.rs b/tracing/src/instrument.rs
index 31bd90b1e3..3e3807a6a4 100644
--- a/tracing/src/instrument.rs
+++ b/tracing/src/instrument.rs
@@ -6,6 +6,13 @@ use core::{
     pin::Pin,
     task::{Context, Poll},
 };
+use core::{
+    future::Future,
+    marker::Sized,
+    mem::{self, ManuallyDrop},
+    pin::Pin,
+    task::{Context, Poll},
+};
 use pin_project_lite::pin_project;
 
 #[cfg(feature = "std")]
@@ -451,14 +458,25 @@ impl<T> WithDispatch<T> {
 
     /// 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 save
+        // pointers to the fields and use `mem::forget` to leave those pointers
+        // valid.
+        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)
     }
 }