From 8aae1c37b091963aafdd336b1168fe5a24c0b4f0 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 | 150 +++++++++++++++------ tracing-futures/tests/std_future.rs | 63 +++++++++ tracing/Cargo.toml | 1 + tracing/src/instrument.rs | 86 ++++++++++-- tracing/tests/instrument.rs | 67 +++++++++ 11 files changed, 361 insertions(+), 63 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 38675a451a..29481c4e7c 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(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::())), ) .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(); 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..acfec81f4b 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) .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(); diff --git a/tracing-futures/Cargo.toml b/tracing-futures/Cargo.toml index d4ef1f71b6..f3b4dc1d61 100644 --- a/tracing-futures/Cargo.toml +++ b/tracing-futures/Cargo.toml @@ -40,6 +40,7 @@ tokio-threadpool = "0.1.18" mio = "0.6.23" [dev-dependencies] +futures = "0.3.21" tokio-test = "0.4.2" tracing-core = { path = "../tracing-core", version = "0.2" } 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 5af13a0a5b..262d15fc87 100644 --- a/tracing-futures/src/lib.rs +++ b/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}; @@ -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::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::Span::current() - /// [entered]: tracing::span::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 WithCollector: 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. @@ -287,9 +339,9 @@ impl core::future::Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::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) } } @@ -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(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(); @@ -589,6 +659,8 @@ mod tests { .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(); @@ -613,6 +685,8 @@ mod tests { .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")) .run_with_handle(); with_default(collector, || { diff --git a/tracing-futures/tests/std_future.rs b/tracing-futures/tests/std_future.rs index cd9656f153..ebba8cf084 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::{collect::with_default, Level}; use tracing_mock::*; @@ -9,6 +12,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 +31,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(); @@ -35,3 +42,59 @@ 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 collector = collector::mock() + .enter(expect::span().named("foo")) + .event( + expect::event() + .with_contextual_parent(Some("foo")) + .at_level(Level::INFO), + ) + .exit(expect::span().named("foo")) + .enter(expect::span().named("foo")) + .exit(expect::span().named("foo")) + .drop_span(expect::span().named("foo")) + .enter(expect::span().named("bar")) + .event( + expect::event() + .with_contextual_parent(Some("bar")) + .at_level(Level::INFO), + ) + .exit(expect::span().named("bar")) + .drop_span(expect::span().named("bar")) + .only() + .run(); + + with_default(collector, || { + // 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 5afe42fc09..2d7bf1120e 100644 --- a/tracing/Cargo.toml +++ b/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.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 98f0da9bd1..208c2bd0cf 100644 --- a/tracing/src/instrument.rs +++ b/tracing/src/instrument.rs @@ -1,7 +1,11 @@ use crate::span::Span; -use core::pin::Pin; -use core::task::{Context, Poll}; -use core::{future::Future, marker::Sized}; +use core::{ + future::Future, + marker::Sized, + mem::{self, ManuallyDrop}, + pin::Pin, + task::{Context, Poll}, +}; use pin_project_lite::pin_project; #[cfg(feature = "std")] @@ -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 { - 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. /// @@ -285,13 +292,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 === @@ -300,9 +349,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) } } @@ -331,19 +380,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..9cb4dfbb3f --- /dev/null +++ b/tracing/tests/instrument.rs @@ -0,0 +1,67 @@ +// 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::{collect::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 collector = collector::mock() + .enter(expect::span().named("foo")) + .event( + expect::event() + .with_contextual_parent(Some("foo")) + .at_level(Level::INFO), + ) + .exit(expect::span().named("foo")) + .enter(expect::span().named("foo")) + .exit(expect::span().named("foo")) + .drop_span(expect::span().named("foo")) + .enter(expect::span().named("bar")) + .event( + expect::event() + .with_contextual_parent(Some("bar")) + .at_level(Level::INFO), + ) + .exit(expect::span().named("bar")) + .drop_span(expect::span().named("bar")) + .only() + .run(); + + with_default(collector, || { + // 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"))); + }); +}