From f0cc5d1e2977309ac609aa5f919d0a6f4d324ca4 Mon Sep 17 00:00:00 2001 From: ilslv Date: Fri, 14 Apr 2023 09:03:26 +0200 Subject: [PATCH 01/13] Instrument `Future` inside `Drop` --- 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/src/executor/futures_01.rs | 23 ++++---- tracing-futures/src/lib.rs | 65 ++++++++++++++++------ tracing-futures/tests/std_future.rs | 4 ++ tracing/Cargo.toml | 1 + tracing/src/instrument.rs | 54 ++++++++++++++---- tracing/tests/instrument.rs | 58 +++++++++++++++++++ 10 files changed, 199 insertions(+), 40 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/src/executor/futures_01.rs b/tracing-futures/src/executor/futures_01.rs index 56ba6e3c42..4e9bd4d715 100644 --- a/tracing-futures/src/executor/futures_01.rs +++ b/tracing-futures/src/executor/futures_01.rs @@ -1,19 +1,10 @@ use crate::{Instrument, Instrumented, WithDispatch}; +use core::mem::ManuallyDrop; use futures_01::{ future::{ExecuteError, Executor}, 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 +12,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 = ManuallyDrop::into_inner(e.into_future().inner); + ExecuteError::new(kind, future) + }) } } @@ -32,7 +27,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..e2c863c403 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}; @@ -147,7 +151,10 @@ pub trait Instrument: Sized { /// /// [entered]: tracing::span::Span::enter() 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 @@ -243,7 +250,7 @@ pin_project! { #[derive(Debug, Clone)] pub struct Instrumented { #[pin] - inner: T, + inner: ManuallyDrop, span: Span, } } @@ -287,9 +294,11 @@ 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 mut this = self.project(); let _enter = this.span.enter(); - this.inner.poll(cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + inner.poll(cx) } } @@ -346,9 +355,11 @@ impl futures::Stream for Instrumented { self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_next(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_next(inner, cx) } } @@ -364,33 +375,41 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_ready(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::start_send(this.inner, item) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::start_send(inner, item) } fn poll_flush( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_flush(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_flush(inner, cx) } fn poll_close( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_close(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_close(inner, cx) } } @@ -419,21 +438,31 @@ 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 + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } } /// 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 + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { - self.inner + 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-futures/tests/std_future.rs b/tracing-futures/tests/std_future.rs index cd9656f153..0b00d123a3 100644 --- a/tracing-futures/tests/std_future.rs +++ b/tracing-futures/tests/std_future.rs @@ -9,6 +9,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 +28,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(); diff --git a/tracing/Cargo.toml b/tracing/Cargo.toml index 5afe42fc09..d2547f868e 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.28", 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..4deabd0cb8 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")] @@ -80,7 +84,10 @@ 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 @@ -289,9 +296,24 @@ pin_project! { #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Instrumented { #[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 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 Instrumented === @@ -300,9 +322,11 @@ impl Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - this.inner.poll(cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + inner.poll(cx) } } @@ -331,19 +355,29 @@ impl Instrumented { /// Get a pinned reference to the wrapped type. pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - self.project_ref().inner + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } } /// Get a pinned mutable reference to the wrapped type. pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { - self.project().inner + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { - self.inner + 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..727b17f363 --- /dev/null +++ b/tracing/tests/instrument.rs @@ -0,0 +1,58 @@ +// 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, Span}; +use tracing_mock::*; + +#[allow(clippy::let_underscore_future)] +#[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) { + assert!(!Span::current().is_none()); + } + } + + 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")) + .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")) + .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 + let _ = Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar")); + }); +} From 1cb523b87d3dc113ecc33fdf6cbee8ae63fbfec2 Mon Sep 17 00:00:00 2001 From: daxpedda Date: Fri, 14 Apr 2023 21:38:48 +0200 Subject: [PATCH 02/13] subscriber: add ability to disable ANSI without crate feature (#2532) ## 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: Eliza Weisman --- tracing-subscriber/src/fmt/fmt_subscriber.rs | 28 +++++++++++++++++--- tracing-subscriber/src/fmt/mod.rs | 18 ++++++++++--- 2 files changed, 40 insertions(+), 6 deletions(-) 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), From a0b5166b7b3d3564576e2ce1c0039a214f40b160 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 10:42:22 +0200 Subject: [PATCH 03/13] Review corrections --- tracing-futures/Cargo.toml | 1 + tracing-futures/src/executor/futures_01.rs | 3 +- tracing-futures/src/lib.rs | 109 +++++++++++++++------ tracing-futures/tests/std_future.rs | 59 +++++++++++ tracing/Cargo.toml | 2 +- tracing/src/instrument.rs | 41 ++++++-- tracing/tests/instrument.rs | 17 +++- 7 files changed, 186 insertions(+), 46 deletions(-) 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 4e9bd4d715..7d4b674af8 100644 --- a/tracing-futures/src/executor/futures_01.rs +++ b/tracing-futures/src/executor/futures_01.rs @@ -1,5 +1,4 @@ use crate::{Instrument, Instrumented, WithDispatch}; -use core::mem::ManuallyDrop; use futures_01::{ future::{ExecuteError, Executor}, Future, @@ -14,7 +13,7 @@ where let future = future.instrument(self.span.clone()); self.inner.execute(future).map_err(|e| { let kind = e.kind(); - let future = ManuallyDrop::into_inner(e.into_future().inner); + let future = e.into_future().into_inner(); ExecuteError::new(kind, future) }) } diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index e2c863c403..a048fefc7b 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -122,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 /// @@ -149,21 +149,22 @@ pub trait Instrument: Sized { /// # } /// ``` /// - /// [entered]: tracing::span::Span::enter() + /// [entered]: Span::enter() fn instrument(self, span: Span) -> Instrumented { - Instrumented { - inner: ManuallyDrop::new(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. /// @@ -187,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()) @@ -247,12 +248,53 @@ 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: 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 pinned mutable reference to the wrapped type. + fn inner_pin_mut(self) -> 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`. + unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + } +} + +#[cfg(feature = "std-future")] +impl<'a, T> InstrumentedProjRef<'a, T> { + /// Get a pinned reference to the wrapped type. + fn inner_pin_ref(self) -> 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`. + unsafe { self.inner.map_unchecked(|v| &**v) } + } } /// A future, stream, sink, or executor that has been instrumented with a `tracing` span. @@ -296,7 +338,7 @@ impl core::future::Future for Instrumented { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; inner.poll(cx) } @@ -357,9 +399,7 @@ impl futures::Stream for Instrumented { ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; - T::poll_next(inner, cx) + T::poll_next(this.inner_pin_mut(), cx) } } @@ -377,7 +417,7 @@ where ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_ready(inner, cx) } @@ -385,7 +425,7 @@ where fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::start_send(inner, item) } @@ -396,7 +436,7 @@ where ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_flush(inner, cx) } @@ -407,7 +447,7 @@ where ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_close(inner, cx) } @@ -438,22 +478,23 @@ impl Instrumented { #[cfg(feature = "std-future")] #[cfg_attr(docsrs, doc(cfg(feature = "std-future")))] pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } + self.project_ref().inner_pin_ref() } /// 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> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } + self.project().inner_pin_mut() } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { + // 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); @@ -599,6 +640,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(); @@ -618,6 +661,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(); @@ -642,6 +687,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 0b00d123a3..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::*; @@ -39,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 d2547f868e..2d7bf1120e 100644 --- a/tracing/Cargo.toml +++ b/tracing/Cargo.toml @@ -38,7 +38,7 @@ pin-project-lite = "0.2.9" [dev-dependencies] criterion = { version = "0.3.6", default_features = false } -futures = { version = "0.3.28", 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 4deabd0cb8..54d7e139cb 100644 --- a/tracing/src/instrument.rs +++ b/tracing/src/instrument.rs @@ -22,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 /// @@ -94,7 +94,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. /// /// This can be used to propagate the current span when spawning a new future. /// @@ -292,9 +292,13 @@ 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: ManuallyDrop, span: Span, @@ -304,7 +308,7 @@ pin_project! { fn drop(this: Pin<&mut Self>) { let this = this.project(); let _enter = this.span.enter(); - // SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because isn't + // 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. @@ -316,6 +320,26 @@ pin_project! { } } +impl<'a, T> InstrumentedProj<'a, T> { + /// Get a pinned mutable reference to the wrapped type. + fn inner_pin_mut(self) -> 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`. + unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + } +} + +impl<'a, T> InstrumentedProjRef<'a, T> { + /// Get a pinned reference to the wrapped type. + fn inner_pin_ref(self) -> 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`. + unsafe { self.inner.map_unchecked(|v| &**v) } + } +} + // === impl Instrumented === impl Future for Instrumented { @@ -324,7 +348,7 @@ impl Future for Instrumented { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; inner.poll(cx) } @@ -355,20 +379,21 @@ impl Instrumented { /// Get a pinned reference to the wrapped type. pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } + self.project_ref().inner_pin_ref() } /// Get a pinned mutable reference to the wrapped type. pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } + self.project().inner_pin_mut() } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { + // 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); diff --git a/tracing/tests/instrument.rs b/tracing/tests/instrument.rs index 727b17f363..9cb4dfbb3f 100644 --- a/tracing/tests/instrument.rs +++ b/tracing/tests/instrument.rs @@ -6,10 +6,9 @@ use std::{future::Future, pin::Pin, task}; use futures::FutureExt as _; -use tracing::{collect::with_default, Instrument as _, Level, Span}; +use tracing::{collect::with_default, Instrument as _, Level}; use tracing_mock::*; -#[allow(clippy::let_underscore_future)] #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)] #[test] fn span_on_drop() { @@ -18,7 +17,7 @@ fn span_on_drop() { impl Drop for AssertSpanOnDrop { fn drop(&mut self) { - assert!(!Span::current().is_none()); + tracing::info!("Drop"); } } @@ -35,11 +34,21 @@ fn span_on_drop() { 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() @@ -53,6 +62,6 @@ fn span_on_drop() { .unwrap(); // never polled - let _ = Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar")); + drop(Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar"))); }); } From 38620f67d9e732c24f018b30145437750e85b0eb Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 10:47:39 +0200 Subject: [PATCH 04/13] Fix `tracing_futures::Instrumented::into_inner` --- tracing-futures/src/lib.rs | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index a048fefc7b..4a5b1dcdaf 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -492,18 +492,23 @@ impl Instrumented { /// /// Note that this drops the span. pub fn into_inner(self) -> T { - // 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(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 } } From abcb488f284229ca813ad34fe7bbe838c32187ca Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 10:51:07 +0200 Subject: [PATCH 05/13] Corrections --- tracing-futures/src/lib.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index 4a5b1dcdaf..1a8d9c5f91 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -399,7 +399,9 @@ impl futures::Stream for Instrumented { ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - T::poll_next(this.inner_pin_mut(), cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_next(inner, cx) } } From 41c10bd2bbfea6bd3fa29c631d60900eb9cdd1e6 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 11:34:34 +0200 Subject: [PATCH 06/13] Refactor unsafe `Pin<&mut ManuallyDrop>` deref into a separate fns --- tracing-futures/src/lib.rs | 45 +++++++++++++++----------------------- tracing/src/instrument.rs | 25 +++++++++++---------- 2 files changed, 31 insertions(+), 39 deletions(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index 1a8d9c5f91..200c5a4885 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -277,23 +277,26 @@ pin_project! { #[cfg(feature = "std-future")] impl<'a, T> InstrumentedProj<'a, T> { - /// Get a pinned mutable reference to the wrapped type. - fn inner_pin_mut(self) -> Pin<&'a mut 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`. - unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + 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 pinned reference to the wrapped type. - fn inner_pin_ref(self) -> Pin<&'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`. - unsafe { self.inner.map_unchecked(|v| &**v) } + let inner = unsafe { self.inner.map_unchecked(|v| &**v) }; + (self.span, inner) } } @@ -336,10 +339,8 @@ impl core::future::Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll { - let mut this = self.project(); - let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); inner.poll(cx) } } @@ -397,10 +398,8 @@ impl futures::Stream for Instrumented { self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_next(inner, cx) } } @@ -417,18 +416,14 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::start_send(inner, item) } @@ -436,10 +431,8 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_flush(inner, cx) } @@ -447,10 +440,8 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_close(inner, cx) } } @@ -480,14 +471,14 @@ 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_pin_ref() + 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_pin_mut() + self.project().span_and_inner_pin_mut().1 } /// Consumes the `Instrumented`, returning the wrapped type. diff --git a/tracing/src/instrument.rs b/tracing/src/instrument.rs index 54d7e139cb..208c2bd0cf 100644 --- a/tracing/src/instrument.rs +++ b/tracing/src/instrument.rs @@ -321,22 +321,25 @@ pin_project! { } impl<'a, T> InstrumentedProj<'a, T> { - /// Get a pinned mutable reference to the wrapped type. - fn inner_pin_mut(self) -> Pin<&'a mut 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`. - unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + let inner = unsafe { self.inner.map_unchecked_mut(|v| &mut **v) }; + (self.span, inner) } } impl<'a, T> InstrumentedProjRef<'a, T> { - /// Get a pinned reference to the wrapped type. - fn inner_pin_ref(self) -> Pin<&'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`. - unsafe { self.inner.map_unchecked(|v| &**v) } + let inner = unsafe { self.inner.map_unchecked(|v| &**v) }; + (self.span, inner) } } @@ -346,10 +349,8 @@ impl Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let mut this = self.project(); - let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); inner.poll(cx) } } @@ -379,12 +380,12 @@ impl Instrumented { /// Get a pinned reference to the wrapped type. pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - self.project_ref().inner_pin_ref() + 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_pin_mut() + self.project().span_and_inner_pin_mut().1 } /// Consumes the `Instrumented`, returning the wrapped type. From e1d6d4eee768bb06f54cd33ff4b5a93371cf0160 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 11:38:14 +0200 Subject: [PATCH 07/13] Corrections --- tracing-futures/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index 200c5a4885..262d15fc87 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -399,7 +399,7 @@ impl futures::Stream for Instrumented { cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_next(inner, cx) } } @@ -417,13 +417,13 @@ where cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::start_send(inner, item) } @@ -432,7 +432,7 @@ where cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_flush(inner, cx) } @@ -441,7 +441,7 @@ where cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_close(inner, cx) } } From 813bd3ebd41360aa8b727d4880612066216b189d Mon Sep 17 00:00:00 2001 From: ilslv Date: Fri, 14 Apr 2023 09:03:26 +0200 Subject: [PATCH 08/13] Instrument `Future` inside `Drop` --- 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/src/executor/futures_01.rs | 23 ++++---- tracing-futures/src/lib.rs | 65 ++++++++++++++++------ tracing-futures/tests/std_future.rs | 4 ++ tracing/Cargo.toml | 1 + tracing/src/instrument.rs | 54 ++++++++++++++---- tracing/tests/instrument.rs | 58 +++++++++++++++++++ 10 files changed, 199 insertions(+), 40 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/src/executor/futures_01.rs b/tracing-futures/src/executor/futures_01.rs index 56ba6e3c42..4e9bd4d715 100644 --- a/tracing-futures/src/executor/futures_01.rs +++ b/tracing-futures/src/executor/futures_01.rs @@ -1,19 +1,10 @@ use crate::{Instrument, Instrumented, WithDispatch}; +use core::mem::ManuallyDrop; use futures_01::{ future::{ExecuteError, Executor}, 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 +12,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 = ManuallyDrop::into_inner(e.into_future().inner); + ExecuteError::new(kind, future) + }) } } @@ -32,7 +27,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..e2c863c403 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}; @@ -147,7 +151,10 @@ pub trait Instrument: Sized { /// /// [entered]: tracing::span::Span::enter() 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 @@ -243,7 +250,7 @@ pin_project! { #[derive(Debug, Clone)] pub struct Instrumented { #[pin] - inner: T, + inner: ManuallyDrop, span: Span, } } @@ -287,9 +294,11 @@ 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 mut this = self.project(); let _enter = this.span.enter(); - this.inner.poll(cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + inner.poll(cx) } } @@ -346,9 +355,11 @@ impl futures::Stream for Instrumented { self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_next(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_next(inner, cx) } } @@ -364,33 +375,41 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_ready(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::start_send(this.inner, item) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::start_send(inner, item) } fn poll_flush( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_flush(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_flush(inner, cx) } fn poll_close( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - T::poll_close(this.inner, cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_close(inner, cx) } } @@ -419,21 +438,31 @@ 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 + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } } /// 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 + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { - self.inner + 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-futures/tests/std_future.rs b/tracing-futures/tests/std_future.rs index cd9656f153..0b00d123a3 100644 --- a/tracing-futures/tests/std_future.rs +++ b/tracing-futures/tests/std_future.rs @@ -9,6 +9,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 +28,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(); diff --git a/tracing/Cargo.toml b/tracing/Cargo.toml index 5afe42fc09..d2547f868e 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.28", 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..4deabd0cb8 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")] @@ -80,7 +84,10 @@ 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 @@ -289,9 +296,24 @@ pin_project! { #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Instrumented { #[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 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 Instrumented === @@ -300,9 +322,11 @@ impl Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.project(); + let mut this = self.project(); let _enter = this.span.enter(); - this.inner.poll(cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + inner.poll(cx) } } @@ -331,19 +355,29 @@ impl Instrumented { /// Get a pinned reference to the wrapped type. pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - self.project_ref().inner + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } } /// Get a pinned mutable reference to the wrapped type. pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { - self.project().inner + // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { - self.inner + 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..727b17f363 --- /dev/null +++ b/tracing/tests/instrument.rs @@ -0,0 +1,58 @@ +// 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, Span}; +use tracing_mock::*; + +#[allow(clippy::let_underscore_future)] +#[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) { + assert!(!Span::current().is_none()); + } + } + + 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")) + .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")) + .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 + let _ = Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar")); + }); +} From b96105796505ec602b96d9f46897d8d01e0e3d44 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 10:42:22 +0200 Subject: [PATCH 09/13] Review corrections --- tracing-futures/Cargo.toml | 1 + tracing-futures/src/executor/futures_01.rs | 3 +- tracing-futures/src/lib.rs | 109 +++++++++++++++------ tracing-futures/tests/std_future.rs | 59 +++++++++++ tracing/Cargo.toml | 2 +- tracing/src/instrument.rs | 41 ++++++-- tracing/tests/instrument.rs | 17 +++- 7 files changed, 186 insertions(+), 46 deletions(-) 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 4e9bd4d715..7d4b674af8 100644 --- a/tracing-futures/src/executor/futures_01.rs +++ b/tracing-futures/src/executor/futures_01.rs @@ -1,5 +1,4 @@ use crate::{Instrument, Instrumented, WithDispatch}; -use core::mem::ManuallyDrop; use futures_01::{ future::{ExecuteError, Executor}, Future, @@ -14,7 +13,7 @@ where let future = future.instrument(self.span.clone()); self.inner.execute(future).map_err(|e| { let kind = e.kind(); - let future = ManuallyDrop::into_inner(e.into_future().inner); + let future = e.into_future().into_inner(); ExecuteError::new(kind, future) }) } diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index e2c863c403..a048fefc7b 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -122,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 /// @@ -149,21 +149,22 @@ pub trait Instrument: Sized { /// # } /// ``` /// - /// [entered]: tracing::span::Span::enter() + /// [entered]: Span::enter() fn instrument(self, span: Span) -> Instrumented { - Instrumented { - inner: ManuallyDrop::new(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. /// @@ -187,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()) @@ -247,12 +248,53 @@ 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: 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 pinned mutable reference to the wrapped type. + fn inner_pin_mut(self) -> 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`. + unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + } +} + +#[cfg(feature = "std-future")] +impl<'a, T> InstrumentedProjRef<'a, T> { + /// Get a pinned reference to the wrapped type. + fn inner_pin_ref(self) -> 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`. + unsafe { self.inner.map_unchecked(|v| &**v) } + } } /// A future, stream, sink, or executor that has been instrumented with a `tracing` span. @@ -296,7 +338,7 @@ impl core::future::Future for Instrumented { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; inner.poll(cx) } @@ -357,9 +399,7 @@ impl futures::Stream for Instrumented { ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; - T::poll_next(inner, cx) + T::poll_next(this.inner_pin_mut(), cx) } } @@ -377,7 +417,7 @@ where ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_ready(inner, cx) } @@ -385,7 +425,7 @@ where fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::start_send(inner, item) } @@ -396,7 +436,7 @@ where ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_flush(inner, cx) } @@ -407,7 +447,7 @@ where ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_close(inner, cx) } @@ -438,22 +478,23 @@ impl Instrumented { #[cfg(feature = "std-future")] #[cfg_attr(docsrs, doc(cfg(feature = "std-future")))] pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } + self.project_ref().inner_pin_ref() } /// 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> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } + self.project().inner_pin_mut() } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { + // 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); @@ -599,6 +640,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(); @@ -618,6 +661,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(); @@ -642,6 +687,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 0b00d123a3..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::*; @@ -39,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 d2547f868e..2d7bf1120e 100644 --- a/tracing/Cargo.toml +++ b/tracing/Cargo.toml @@ -38,7 +38,7 @@ pin-project-lite = "0.2.9" [dev-dependencies] criterion = { version = "0.3.6", default_features = false } -futures = { version = "0.3.28", 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 4deabd0cb8..54d7e139cb 100644 --- a/tracing/src/instrument.rs +++ b/tracing/src/instrument.rs @@ -22,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 /// @@ -94,7 +94,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. /// /// This can be used to propagate the current span when spawning a new future. /// @@ -292,9 +292,13 @@ 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: ManuallyDrop, span: Span, @@ -304,7 +308,7 @@ pin_project! { fn drop(this: Pin<&mut Self>) { let this = this.project(); let _enter = this.span.enter(); - // SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because isn't + // 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. @@ -316,6 +320,26 @@ pin_project! { } } +impl<'a, T> InstrumentedProj<'a, T> { + /// Get a pinned mutable reference to the wrapped type. + fn inner_pin_mut(self) -> 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`. + unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + } +} + +impl<'a, T> InstrumentedProjRef<'a, T> { + /// Get a pinned reference to the wrapped type. + fn inner_pin_ref(self) -> 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`. + unsafe { self.inner.map_unchecked(|v| &**v) } + } +} + // === impl Instrumented === impl Future for Instrumented { @@ -324,7 +348,7 @@ impl Future for Instrumented { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut this = self.project(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; inner.poll(cx) } @@ -355,20 +379,21 @@ impl Instrumented { /// Get a pinned reference to the wrapped type. pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project_ref().inner.map_unchecked(|v| &**v) } + self.project_ref().inner_pin_ref() } /// Get a pinned mutable reference to the wrapped type. pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { - // SAFETY: As long as `ManuallyDrop` does not move, `T` wont move. - unsafe { self.project().inner.map_unchecked_mut(|v| &mut **v) } + self.project().inner_pin_mut() } /// Consumes the `Instrumented`, returning the wrapped type. /// /// Note that this drops the span. pub fn into_inner(self) -> T { + // 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); diff --git a/tracing/tests/instrument.rs b/tracing/tests/instrument.rs index 727b17f363..9cb4dfbb3f 100644 --- a/tracing/tests/instrument.rs +++ b/tracing/tests/instrument.rs @@ -6,10 +6,9 @@ use std::{future::Future, pin::Pin, task}; use futures::FutureExt as _; -use tracing::{collect::with_default, Instrument as _, Level, Span}; +use tracing::{collect::with_default, Instrument as _, Level}; use tracing_mock::*; -#[allow(clippy::let_underscore_future)] #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)] #[test] fn span_on_drop() { @@ -18,7 +17,7 @@ fn span_on_drop() { impl Drop for AssertSpanOnDrop { fn drop(&mut self) { - assert!(!Span::current().is_none()); + tracing::info!("Drop"); } } @@ -35,11 +34,21 @@ fn span_on_drop() { 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() @@ -53,6 +62,6 @@ fn span_on_drop() { .unwrap(); // never polled - let _ = Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar")); + drop(Fut(Some(AssertSpanOnDrop)).instrument(tracing::span!(Level::TRACE, "bar"))); }); } From b24f38366612fccbee94fcbde799cabc99ae5f8b Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 10:47:39 +0200 Subject: [PATCH 10/13] Fix `tracing_futures::Instrumented::into_inner` --- tracing-futures/src/lib.rs | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index a048fefc7b..4a5b1dcdaf 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -492,18 +492,23 @@ impl Instrumented { /// /// Note that this drops the span. pub fn into_inner(self) -> T { - // 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(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 } } From daf27b50a40c75e4fad6452b13ce18bc3505d817 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 10:51:07 +0200 Subject: [PATCH 11/13] Corrections --- tracing-futures/src/lib.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index 4a5b1dcdaf..1a8d9c5f91 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -399,7 +399,9 @@ impl futures::Stream for Instrumented { ) -> futures::task::Poll> { let mut this = self.project(); let _enter = this.span.enter(); - T::poll_next(this.inner_pin_mut(), cx) + // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. + let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + T::poll_next(inner, cx) } } From 95f4ed7e0d2d31a88b78b618d40ec5ea95701677 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 11:34:34 +0200 Subject: [PATCH 12/13] Refactor unsafe `Pin<&mut ManuallyDrop>` deref into a separate fns --- tracing-futures/src/lib.rs | 45 +++++++++++++++----------------------- tracing/src/instrument.rs | 25 +++++++++++---------- 2 files changed, 31 insertions(+), 39 deletions(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index 1a8d9c5f91..200c5a4885 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -277,23 +277,26 @@ pin_project! { #[cfg(feature = "std-future")] impl<'a, T> InstrumentedProj<'a, T> { - /// Get a pinned mutable reference to the wrapped type. - fn inner_pin_mut(self) -> Pin<&'a mut 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`. - unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + 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 pinned reference to the wrapped type. - fn inner_pin_ref(self) -> Pin<&'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`. - unsafe { self.inner.map_unchecked(|v| &**v) } + let inner = unsafe { self.inner.map_unchecked(|v| &**v) }; + (self.span, inner) } } @@ -336,10 +339,8 @@ impl core::future::Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll { - let mut this = self.project(); - let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); inner.poll(cx) } } @@ -397,10 +398,8 @@ impl futures::Stream for Instrumented { self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_next(inner, cx) } } @@ -417,18 +416,14 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::start_send(inner, item) } @@ -436,10 +431,8 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_flush(inner, cx) } @@ -447,10 +440,8 @@ where self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> futures::task::Poll> { - let mut this = self.project(); + let (span, inner) = self.project().span_and_inner_pin_mut(); let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; T::poll_close(inner, cx) } } @@ -480,14 +471,14 @@ 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_pin_ref() + 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_pin_mut() + self.project().span_and_inner_pin_mut().1 } /// Consumes the `Instrumented`, returning the wrapped type. diff --git a/tracing/src/instrument.rs b/tracing/src/instrument.rs index 54d7e139cb..208c2bd0cf 100644 --- a/tracing/src/instrument.rs +++ b/tracing/src/instrument.rs @@ -321,22 +321,25 @@ pin_project! { } impl<'a, T> InstrumentedProj<'a, T> { - /// Get a pinned mutable reference to the wrapped type. - fn inner_pin_mut(self) -> Pin<&'a mut 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`. - unsafe { self.inner.map_unchecked_mut(|v| &mut **v) } + let inner = unsafe { self.inner.map_unchecked_mut(|v| &mut **v) }; + (self.span, inner) } } impl<'a, T> InstrumentedProjRef<'a, T> { - /// Get a pinned reference to the wrapped type. - fn inner_pin_ref(self) -> Pin<&'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`. - unsafe { self.inner.map_unchecked(|v| &**v) } + let inner = unsafe { self.inner.map_unchecked(|v| &**v) }; + (self.span, inner) } } @@ -346,10 +349,8 @@ impl Future for Instrumented { type Output = T::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let mut this = self.project(); - let _enter = this.span.enter(); - // SAFETY: As long as `ManuallyDrop` does not move, `T` won't move. - let inner = unsafe { this.inner.as_mut().map_unchecked_mut(|v| &mut **v) }; + let (span, inner) = self.project().span_and_inner_pin_mut(); + let _enter = span.enter(); inner.poll(cx) } } @@ -379,12 +380,12 @@ impl Instrumented { /// Get a pinned reference to the wrapped type. pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { - self.project_ref().inner_pin_ref() + 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_pin_mut() + self.project().span_and_inner_pin_mut().1 } /// Consumes the `Instrumented`, returning the wrapped type. From c8b56a8c58e6ef660b87599056610552230add21 Mon Sep 17 00:00:00 2001 From: ilslv Date: Mon, 17 Apr 2023 11:38:14 +0200 Subject: [PATCH 13/13] Corrections --- tracing-futures/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tracing-futures/src/lib.rs b/tracing-futures/src/lib.rs index 200c5a4885..262d15fc87 100644 --- a/tracing-futures/src/lib.rs +++ b/tracing-futures/src/lib.rs @@ -399,7 +399,7 @@ impl futures::Stream for Instrumented { cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_next(inner, cx) } } @@ -417,13 +417,13 @@ where cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_ready(inner, cx) } fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::start_send(inner, item) } @@ -432,7 +432,7 @@ where cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_flush(inner, cx) } @@ -441,7 +441,7 @@ where cx: &mut Context<'_>, ) -> futures::task::Poll> { let (span, inner) = self.project().span_and_inner_pin_mut(); - let _enter = this.span.enter(); + let _enter = span.enter(); T::poll_close(inner, cx) } }