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-subscriber/src/fmt/fmt_subscriber.rs b/tracing-subscriber/src/fmt/fmt_subscriber.rs index 9887b711ff..1b89039261 100644 --- a/tracing-subscriber/src/fmt/fmt_subscriber.rs +++ b/tracing-subscriber/src/fmt/fmt_subscriber.rs @@ -274,10 +274,32 @@ impl Subscriber { } } - /// Enable ANSI terminal colors for formatted output. - #[cfg(feature = "ansi")] - #[cfg_attr(docsrs, doc(cfg(feature = "ansi")))] + /// Sets whether or not the formatter emits ANSI terminal escape codes + /// for colors and other text formatting. + /// + /// Enabling ANSI escapes (calling `with_ansi(true)`) requires the "ansi" + /// crate feature flag. Calling `with_ansi(true)` without the "ansi" + /// feature flag enabled will panic if debug assertions are enabled, or + /// print a warning otherwise. + /// + /// This method itself is still available without the feature flag. This + /// is to allow ANSI escape codes to be explicitly *disabled* without + /// having to opt-in to the dependencies required to emit ANSI formatting. + /// This way, code which constructs a formatter that should never emit + /// ANSI escape codes can ensure that they are not used, regardless of + /// whether or not other crates in the dependency graph enable the "ansi" + /// feature flag. pub fn with_ansi(self, ansi: bool) -> Self { + #[cfg(not(feature = "ansi"))] + if ansi { + const ERROR: &str = + "tracing-subscriber: the `ansi` crate feature is required to enable ANSI terminal colors"; + #[cfg(debug_assertions)] + panic!("{}", ERROR); + #[cfg(not(debug_assertions))] + eprintln!("{}", ERROR); + } + Subscriber { is_ansi: ansi, ..self diff --git a/tracing-subscriber/src/fmt/mod.rs b/tracing-subscriber/src/fmt/mod.rs index ba5066d2ec..38acf4d092 100644 --- a/tracing-subscriber/src/fmt/mod.rs +++ b/tracing-subscriber/src/fmt/mod.rs @@ -610,9 +610,21 @@ where } } - /// Enable ANSI terminal colors for formatted output. - #[cfg(feature = "ansi")] - #[cfg_attr(docsrs, doc(cfg(feature = "ansi")))] + /// Sets whether or not the formatter emits ANSI terminal escape codes + /// for colors and other text formatting. + /// + /// Enabling ANSI escapes (calling `with_ansi(true)`) requires the "ansi" + /// crate feature flag. Calling `with_ansi(true)` without the "ansi" + /// feature flag enabled will panic if debug assertions are enabled, or + /// print a warning otherwise. + /// + /// This method itself is still available without the feature flag. This + /// is to allow ANSI escape codes to be explicitly *disabled* without + /// having to opt-in to the dependencies required to emit ANSI formatting. + /// This way, code which constructs a formatter that should never emit + /// ANSI escape codes can ensure that they are not used, regardless of + /// whether or not other crates in the dependency graph enable the "ansi" + /// feature flag. pub fn with_ansi(self, ansi: bool) -> CollectorBuilder, F, W> { CollectorBuilder { inner: self.inner.with_ansi(ansi), 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"))); + }); +}