From 6b60a0410ab30e536d88d353fb6e1de6b53023fe 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 Co-authored-by: Eliza Weisman --- tracing-attributes/tests/async_fn.rs | 26 ++++ tracing-attributes/tests/err.rs | 4 + tracing-attributes/tests/follows_from.rs | 2 + tracing-attributes/tests/ret.rs | 2 + tracing-futures/Cargo.toml | 1 + tracing-futures/src/executor/futures_01.rs | 22 ++- tracing-futures/src/lib.rs | 152 +++++++++++++++------ tracing-futures/tests/std_future.rs | 61 +++++++++ tracing/Cargo.toml | 3 +- tracing/src/instrument.rs | 87 ++++++++++-- tracing/tests/instrument.rs | 65 +++++++++ 11 files changed, 360 insertions(+), 65 deletions(-) create mode 100644 tracing/tests/instrument.rs diff --git a/tracing-attributes/tests/async_fn.rs b/tracing-attributes/tests/async_fn.rs index c89963672c..d6d874ffd7 100644 --- a/tracing-attributes/tests/async_fn.rs +++ b/tracing-attributes/tests/async_fn.rs @@ -90,6 +90,8 @@ fn async_fn_only_enters_for_polls() { .exit(span::mock().named("test_async_fn")) .enter(span::mock().named("test_async_fn")) .exit(span::mock().named("test_async_fn")) + .enter(span::mock().named("test_async_fn")) + .exit(span::mock().named("test_async_fn")) .drop_span(span::mock().named("test_async_fn")) .done() .run_with_handle(); @@ -120,8 +122,12 @@ fn async_fn_nested() { .enter(span2.clone()) .event(event::mock().with_fields(field::mock("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) .done() .run_with_handle(); @@ -199,13 +205,19 @@ fn async_fn_with_async_trait() { .enter(span3.clone()) .event(event::mock().with_fields(field::mock("val").with_value(&2u64))) .exit(span3.clone()) + .enter(span3.clone()) + .exit(span3.clone()) .drop_span(span3) .new_span(span2.clone().with_field(field::mock("self"))) .enter(span2.clone()) .event(event::mock().with_fields(field::mock("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) .done() .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) .done() .run_with_handle(); @@ -331,8 +345,12 @@ fn async_fn_with_async_trait_and_fields_expressions_with_generic_parameter() { .with_field(field::mock("Self").with_value(&std::any::type_name::())), ) .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) .done() .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) .done() .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) .done() .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) .done() .run_with_handle(); diff --git a/tracing-attributes/tests/err.rs b/tracing-attributes/tests/err.rs index deb2f1d59d..485dd11961 100644 --- a/tracing-attributes/tests/err.rs +++ b/tracing-attributes/tests/err.rs @@ -57,6 +57,8 @@ fn test_async() { .enter(span.clone()) .event(event::mock().at_level(Level::ERROR)) .exit(span.clone()) + .enter(span.clone()) + .exit(span.clone()) .drop_span(span) .done() .run_with_handle(); @@ -111,6 +113,8 @@ fn test_mut_async() { .enter(span.clone()) .event(event::mock().at_level(Level::ERROR)) .exit(span.clone()) + .enter(span.clone()) + .exit(span.clone()) .drop_span(span) .done() .run_with_handle(); diff --git a/tracing-attributes/tests/follows_from.rs b/tracing-attributes/tests/follows_from.rs index da0eec6357..a589409ded 100644 --- a/tracing-attributes/tests/follows_from.rs +++ b/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) .done() .run_with_handle(); diff --git a/tracing-attributes/tests/ret.rs b/tracing-attributes/tests/ret.rs index f20b66797c..f56c80baaf 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) .done() .run_with_handle(); diff --git a/tracing-futures/Cargo.toml b/tracing-futures/Cargo.toml index a34b459b1e..efc87eb248 100644 --- a/tracing-futures/Cargo.toml +++ b/tracing-futures/Cargo.toml @@ -40,6 +40,7 @@ tokio-threadpool = { version = "0.1.18", optional = true } mio = { version = "0.6.23", optional = true } [dev-dependencies] +futures = "0.3.21" tokio-test = "0.4.2" tracing-core = { path = "../tracing-core", version = "0.1.28" } tracing-mock = { path = "../tracing-mock", features = ["tokio-test"] } diff --git a/tracing-futures/src/executor/futures_01.rs b/tracing-futures/src/executor/futures_01.rs index 56ba6e3c42..7d4b674af8 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 Executor for Instrumented where T: Executor>, @@ -21,7 +11,11 @@ where { fn execute(&self, future: F) -> Result<(), ExecuteError> { 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 = e.into_future().into_inner(); + ExecuteError::new(kind, future) + }) } } @@ -32,7 +26,11 @@ where { fn execute(&self, future: F) -> Result<(), ExecuteError> { 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) + }) } } diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index addf19f159..8065175d28 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -103,7 +103,11 @@ use pin_project_lite::pin_project; pub(crate) mod stdlib; #[cfg(feature = "std-future")] -use crate::stdlib::{pin::Pin, task::Context}; +use core::{ + mem::{self, ManuallyDrop}, + pin::Pin, + task::Context, +}; #[cfg(feature = "std")] use tracing::{dispatcher, Dispatch}; @@ -118,13 +122,13 @@ pub mod executor; /// /// [span]: mod@tracing::span pub trait Instrument: Sized { - /// Instruments this type with the provided `Span`, returning an - /// `Instrumented` wrapper. + /// 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`. + /// If the instrumented type is a future, stream, or sink, the attached + /// [`Span`] will be [entered] every time it is polled or [`Drop`]ped. If + /// the instrumented type is a future executor, every future spawned on that + /// executor will be instrumented by the attached [`Span`]. /// /// # Examples /// @@ -145,18 +149,22 @@ pub trait Instrument: Sized { /// # } /// ``` /// - /// [entered]: tracing::Span::enter + /// [entered]: Span::enter() fn instrument(self, span: Span) -> Instrumented { - Instrumented { inner: self, span } + #[cfg(feature = "std-future")] + let inner = ManuallyDrop::new(self); + #[cfg(not(feature = "std-future"))] + let inner = self; + Instrumented { inner, span } } - /// Instruments this type with the [current] `Span`, returning an - /// `Instrumented` wrapper. + /// 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`. + /// If the instrumented type is a future, stream, or sink, the attached + /// [`Span`] will be [entered] every time it is polled or [`Drop`]ped. 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. /// @@ -180,8 +188,8 @@ pub trait Instrument: Sized { /// # } /// ``` /// - /// [current]: tracing::Span::current - /// [entered]: tracing::Span::enter + /// [current]: Span::current() + /// [entered]: Span::enter() #[inline] fn in_current_span(self) -> Instrumented { self.instrument(Span::current()) @@ -240,12 +248,56 @@ pub trait WithSubscriber: Sized { #[cfg(feature = "std-future")] pin_project! { /// A future, stream, sink, or executor that has been instrumented with a `tracing` span. + #[project = InstrumentedProj] + #[project_ref = InstrumentedProjRef] #[derive(Debug, Clone)] pub struct Instrumented { + // `ManuallyDrop` is used here to to enter instrument `Drop` by entering + // `Span` and executing `ManuallyDrop::drop`. #[pin] - inner: T, + inner: ManuallyDrop, span: Span, } + + impl PinnedDrop for Instrumented { + 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()) } + } + } +} + +#[cfg(feature = "std-future")] +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` 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) + } +} + +#[cfg(feature = "std-future")] +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` 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) + } } /// A future, stream, sink, or executor that has been instrumented with a `tracing` span. @@ -286,10 +338,10 @@ impl Instrument for T {} impl crate::stdlib::future::Future for Instrumented { type Output = T::Output; - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> crate::stdlib::task::Poll { - let this = self.project(); - let _enter = this.span.enter(); - this.inner.poll(cx) + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll { + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); + inner.poll(cx) } } @@ -346,9 +398,9 @@ impl futures::Stream for Instrumented { self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); - let _enter = this.span.enter(); - T::poll_next(this.inner, cx) + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); + T::poll_next(inner, cx) } } @@ -364,33 +416,33 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); - let _enter = this.span.enter(); - T::poll_ready(this.inner, cx) + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); + T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { - let this = self.project(); - let _enter = this.span.enter(); - T::start_send(this.inner, item) + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); + T::start_send(inner, item) } fn poll_flush( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); - let _enter = this.span.enter(); - T::poll_flush(this.inner, cx) + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); + T::poll_flush(inner, cx) } fn poll_close( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); - let _enter = this.span.enter(); - T::poll_close(this.inner, cx) + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); + T::poll_close(inner, cx) } } @@ -419,20 +471,36 @@ impl Instrumented { #[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 + self.project_ref().span_and_inner_pin_ref().1 } /// 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 + 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 { + #[cfg(feature = "std-future")] + { + // 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 = &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) + } + #[cfg(not(feature = "std-future"))] self.inner } } @@ -570,6 +638,8 @@ mod tests { .exit(span::mock().named("foo")) .enter(span::mock().named("foo")) .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) .drop_span(span::mock().named("foo")) .done() .run_with_handle(); @@ -589,6 +659,8 @@ mod tests { .exit(span::mock().named("foo")) .enter(span::mock().named("foo")) .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) .drop_span(span::mock().named("foo")) .done() .run_with_handle(); @@ -613,6 +685,8 @@ mod tests { .exit(span::mock().named("foo")) .enter(span::mock().named("foo")) .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) .drop_span(span::mock().named("foo")) .run_with_handle(); with_default(subscriber, || { diff --git a/tracing-futures/tests/std_future.rs b/tracing-futures/tests/std_future.rs index a8d27bb6ba..da47569de4 100644 --- a/tracing-futures/tests/std_future.rs +++ b/tracing-futures/tests/std_future.rs @@ -1,3 +1,6 @@ +use std::{future::Future, pin::Pin, task}; + +use futures::FutureExt as _; use tracing::Instrument; use tracing::{subscriber::with_default, Level}; use tracing_mock::*; @@ -9,6 +12,8 @@ fn enter_exit_is_reasonable() { .exit(span::mock().named("foo")) .enter(span::mock().named("foo")) .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) .drop_span(span::mock().named("foo")) .done() .run_with_handle(); @@ -26,6 +31,8 @@ fn error_ends_span() { .exit(span::mock().named("foo")) .enter(span::mock().named("foo")) .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) .drop_span(span::mock().named("foo")) .done() .run_with_handle(); @@ -35,3 +42,57 @@ fn error_ends_span() { }); handle.assert_finished(); } + +#[test] +fn span_on_drop() { + #[derive(Clone, Debug)] + struct AssertSpanOnDrop; + + impl Drop for AssertSpanOnDrop { + fn drop(&mut self) { + tracing::info!("Drop"); + } + } + + struct Fut(Option); + + impl Future for Fut { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, _: &mut task::Context<'_>) -> task::Poll { + self.set(Fut(None)); + task::Poll::Ready(()) + } + } + + let subscriber = subscriber::mock() + .enter(span::mock().named("foo")) + .event( + event::mock() + .at_level(Level::INFO), + ) + .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) + .drop_span(span::mock().named("foo")) + .enter(span::mock().named("bar")) + .event( + event::mock() + .at_level(Level::INFO), + ) + .exit(span::mock().named("bar")) + .drop_span(span::mock().named("bar")) + .done() + .run(); + + with_default(subscriber, || { + // polled once + Fut(Some(AssertSpanOnDrop)) + .instrument(tracing::span!(Level::TRACE, "foo")) + .now_or_never() + .unwrap(); + + // never polled + drop(Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar"))); + }); +} diff --git a/tracing/Cargo.toml b/tracing/Cargo.toml index 7cf472048b..234c4ce0d5 100644 --- a/tracing/Cargo.toml +++ b/tracing/Cargo.toml @@ -34,7 +34,8 @@ tracing-attributes = { path = "../tracing-attributes", version = "0.1.23", optio pin-project-lite = "0.2.9" [dev-dependencies] -criterion = { version = "0.3.6", default-features = false } +criterion = { version = "0.3.6", default_features = false } +futures = { version = "0.3.21", default_features = false } log = "0.4.17" tracing-mock = { path = "../tracing-mock" } diff --git a/tracing/src/instrument.rs b/tracing/src/instrument.rs index 46e5f579cd..376a2c903d 100644 --- a/tracing/src/instrument.rs +++ b/tracing/src/instrument.rs @@ -1,10 +1,15 @@ -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::{self, ManuallyDrop}, + pin::Pin, + task::{Context, Poll}, +}; use pin_project_lite::pin_project; /// Attaches spans to a [`std::future::Future`]. @@ -18,7 +23,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 +85,17 @@ pub trait Instrument: Sized { /// [disabled]: super::Span::is_disabled() /// [`Future`]: std::future::Future fn instrument(self, span: Span) -> Instrumented { - 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 +260,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 { + // `ManuallyDrop` is used here to to enter instrument `Drop` by entering + // `Span` and executing `ManuallyDrop::drop`. #[pin] - inner: T, + inner: ManuallyDrop, span: Span, } + + impl PinnedDrop for Instrumented { + 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` 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` 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 +317,9 @@ impl Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - 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 +348,30 @@ impl Instrumented { /// 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 save + // pointers to the fields and use `mem::forget` to leave those pointers + // valid. + let span: *const Span = &self.span; + let inner: *const ManuallyDrop = &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) } } diff --git a/tracing/tests/instrument.rs b/tracing/tests/instrument.rs new file mode 100644 index 0000000000..1c4ebc9ffa --- /dev/null +++ b/tracing/tests/instrument.rs @@ -0,0 +1,65 @@ +// These tests require the thread-local scoped dispatcher, which only works when +// we have a standard library. The behaviour being tested should be the same +// with the standard lib disabled. +#![cfg(feature = "std")] + +use std::{future::Future, pin::Pin, task}; + +use futures::FutureExt as _; +use tracing::{subscriber::with_default, Instrument as _, Level}; +use tracing_mock::*; + +#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)] +#[test] +fn span_on_drop() { + #[derive(Clone, Debug)] + struct AssertSpanOnDrop; + + impl Drop for AssertSpanOnDrop { + fn drop(&mut self) { + tracing::info!("Drop"); + } + } + + struct Fut(Option); + + impl Future for Fut { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, _: &mut task::Context<'_>) -> task::Poll { + self.set(Fut(None)); + task::Poll::Ready(()) + } + } + + let subscriber = subscriber::mock() + .enter(span::mock().named("foo")) + .event( + event::mock() + .at_level(Level::INFO), + ) + .exit(span::mock().named("foo")) + .enter(span::mock().named("foo")) + .exit(span::mock().named("foo")) + .drop_span(span::mock().named("foo")) + .enter(span::mock().named("bar")) + .event( + event::mock() + .at_level(Level::INFO), + ) + .exit(span::mock().named("bar")) + .drop_span(span::mock().named("bar")) + .done() + .run(); + + with_default(subscriber, || { + // polled once + Fut(Some(AssertSpanOnDrop)) + .instrument(tracing::span!(Level::TRACE, "foo")) + .now_or_never() + .unwrap(); + + // never polled + drop(Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar"))); + }); +}