Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

tracing, tracing-futures: instrument Future inside Drop #2562

Merged
merged 14 commits into from
Apr 18, 2023
26 changes: 26 additions & 0 deletions tracing-attributes/tests/async_fn.rs
Original file line number Diff line number Diff line change
@@ -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::<TestImpl>())),
)
.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();
4 changes: 4 additions & 0 deletions tracing-attributes/tests/err.rs
Original file line number Diff line number Diff line change
@@ -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();
2 changes: 2 additions & 0 deletions tracing-attributes/tests/follows_from.rs
Original file line number Diff line number Diff line change
@@ -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();
2 changes: 2 additions & 0 deletions tracing-attributes/tests/ret.rs
Original file line number Diff line number Diff line change
@@ -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();
1 change: 1 addition & 0 deletions tracing-futures/Cargo.toml
Original file line number Diff line number Diff line change
@@ -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"] }
22 changes: 10 additions & 12 deletions tracing-futures/src/executor/futures_01.rs
Original file line number Diff line number Diff line change
@@ -4,24 +4,18 @@ use futures_01::{
Future,
};

macro_rules! deinstrument_err {
($e:expr) => {
$e.map_err(|e| {
let kind = e.kind();
let future = e.into_future().inner;
ExecuteError::new(kind, future)
})
};
}

impl<T, F> Executor<F> for Instrumented<T>
where
T: Executor<Instrumented<F>>,
F: Future<Item = (), Error = ()>,
{
fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
let future = future.instrument(self.span.clone());
deinstrument_err!(self.inner.execute(future))
self.inner.execute(future).map_err(|e| {
let kind = e.kind();
let future = e.into_future().into_inner();
ExecuteError::new(kind, future)
})
}
}

@@ -32,7 +26,11 @@ where
{
fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
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)
})
}
}

150 changes: 112 additions & 38 deletions tracing-futures/src/lib.rs
Original file line number Diff line number Diff line change
@@ -103,7 +103,11 @@
use pin_project_lite::pin_project;

#[cfg(feature = "std-future")]
use core::{pin::Pin, task::Context};
use core::{
mem::{self, ManuallyDrop},
pin::Pin,
task::Context,
};

#[cfg(feature = "std")]
use tracing::{dispatch, Dispatch};
@@ -118,13 +122,13 @@ pub mod executor;
///
/// [span]: mod@tracing::span
pub trait Instrument: Sized {
/// Instruments this type with the provided `Span`, returning an
/// `Instrumented` wrapper.
/// Instruments this type with the provided [`Span`], returning an
/// [`Instrumented`] wrapper.
///
/// If the instrumented type is a future, stream, or sink, the attached `Span`
/// will be [entered] every time it is polled. If the instrumented type
/// is a future executor, every future spawned on that executor will be
/// instrumented by the attached `Span`.
/// If the instrumented type is a future, stream, or sink, the attached
/// [`Span`] will be [entered] every time it is polled or [`Drop`]ped. If
/// the instrumented type is a future executor, every future spawned on that
/// executor will be instrumented by the attached [`Span`].
///
/// # Examples
///
@@ -145,18 +149,22 @@ pub trait Instrument: Sized {
/// # }
/// ```
///
/// [entered]: tracing::span::Span::enter()
/// [entered]: Span::enter()
fn instrument(self, span: Span) -> Instrumented<Self> {
Instrumented { inner: self, span }
#[cfg(feature = "std-future")]
let inner = ManuallyDrop::new(self);
#[cfg(not(feature = "std-future"))]
let inner = self;
Instrumented { inner, span }
}

/// Instruments this type with the [current] `Span`, returning an
/// `Instrumented` wrapper.
/// Instruments this type with the [current] [`Span`], returning an
/// [`Instrumented`] wrapper.
///
/// If the instrumented type is a future, stream, or sink, the attached `Span`
/// will be [entered] every time it is polled. If the instrumented type
/// is a future executor, every future spawned on that executor will be
/// instrumented by the attached `Span`.
/// If the instrumented type is a future, stream, or sink, the attached
/// [`Span`] will be [entered] every time it is polled or [`Drop`]ped. If
/// the instrumented type is a future executor, every future spawned on that
/// executor will be instrumented by the attached [`Span`].
///
/// This can be used to propagate the current span when spawning a new future.
///
@@ -180,8 +188,8 @@ pub trait Instrument: Sized {
/// # }
/// ```
///
/// [current]: tracing::span::Span::current()
/// [entered]: tracing::span::Span::enter()
/// [current]: Span::current()
/// [entered]: Span::enter()
#[inline]
fn in_current_span(self) -> Instrumented<Self> {
self.instrument(Span::current())
@@ -240,12 +248,56 @@ pub trait WithCollector: Sized {
#[cfg(feature = "std-future")]
pin_project! {
/// A future, stream, sink, or executor that has been instrumented with a `tracing` span.
#[project = InstrumentedProj]
#[project_ref = InstrumentedProjRef]
#[derive(Debug, Clone)]
pub struct Instrumented<T> {
// `ManuallyDrop` is used here to to enter instrument `Drop` by entering
// `Span` and executing `ManuallyDrop::drop`.
#[pin]
inner: T,
inner: ManuallyDrop<T>,
span: Span,
}

impl<T> PinnedDrop for Instrumented<T> {
fn drop(this: Pin<&mut Self>) {
let this = this.project();
let _enter = this.span.enter();
// SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because this isn't
// different from wrapping `T` in `Option` and calling
// `Pin::set(&mut this.inner, None)`, except avoiding
// additional memory overhead.
// 2. `ManuallyDrop::drop()` is safe, because
// `PinnedDrop::drop()` is guaranteed to be called only
// once.
unsafe { ManuallyDrop::drop(this.inner.get_unchecked_mut()) }
}
}
}

#[cfg(feature = "std-future")]
impl<'a, T> InstrumentedProj<'a, T> {
/// Get a mutable reference to the [`Span`] a pinned mutable reference to
/// the wrapped type.
fn span_and_inner_pin_mut(self) -> (&'a mut Span, Pin<&'a mut T>) {
// SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
// and `inner` is valid, because `ManuallyDrop::drop` is called
// only inside `Drop` of the `Instrumented`.
let inner = unsafe { self.inner.map_unchecked_mut(|v| &mut **v) };
(self.span, inner)
}
}

#[cfg(feature = "std-future")]
impl<'a, T> InstrumentedProjRef<'a, T> {
/// Get a reference to the [`Span`] a pinned reference to the wrapped type.
fn span_and_inner_pin_ref(self) -> (&'a Span, Pin<&'a T>) {
// SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
// and `inner` is valid, because `ManuallyDrop::drop` is called
// only inside `Drop` of the `Instrumented`.
let inner = unsafe { self.inner.map_unchecked(|v| &**v) };
(self.span, inner)
}
}

/// A future, stream, sink, or executor that has been instrumented with a `tracing` span.
@@ -287,9 +339,9 @@ impl<T: core::future::Future> core::future::Future for Instrumented<T> {
type Output = T::Output;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> core::task::Poll<Self::Output> {
let this = self.project();
let _enter = this.span.enter();
this.inner.poll(cx)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
inner.poll(cx)
}
}

@@ -346,9 +398,9 @@ impl<T: futures::Stream> futures::Stream for Instrumented<T> {
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> futures::task::Poll<Option<Self::Item>> {
let this = self.project();
let _enter = this.span.enter();
T::poll_next(this.inner, cx)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
T::poll_next(inner, cx)
}
}

@@ -364,33 +416,33 @@ where
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> futures::task::Poll<Result<(), Self::Error>> {
let this = self.project();
let _enter = this.span.enter();
T::poll_ready(this.inner, cx)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
T::poll_ready(inner, cx)
}

fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
let this = self.project();
let _enter = this.span.enter();
T::start_send(this.inner, item)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
T::start_send(inner, item)
}

fn poll_flush(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> futures::task::Poll<Result<(), Self::Error>> {
let this = self.project();
let _enter = this.span.enter();
T::poll_flush(this.inner, cx)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
T::poll_flush(inner, cx)
}

fn poll_close(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> futures::task::Poll<Result<(), Self::Error>> {
let this = self.project();
let _enter = this.span.enter();
T::poll_close(this.inner, cx)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
T::poll_close(inner, cx)
}
}

@@ -419,20 +471,36 @@ impl<T> Instrumented<T> {
#[cfg(feature = "std-future")]
#[cfg_attr(docsrs, doc(cfg(feature = "std-future")))]
pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
self.project_ref().inner
self.project_ref().span_and_inner_pin_ref().1
}

/// Get a pinned mutable reference to the wrapped type.
#[cfg(feature = "std-future")]
#[cfg_attr(docsrs, doc(cfg(feature = "std-future")))]
pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
self.project().inner
self.project().span_and_inner_pin_mut().1
}

/// Consumes the `Instrumented`, returning the wrapped type.
///
/// Note that this drops the span.
pub fn into_inner(self) -> T {
#[cfg(feature = "std-future")]
{
// To manually destructure `Instrumented` without `Drop`, we save
// pointers to the fields and use `mem::forget` to leave those pointers
// valid.
let span: *const Span = &self.span;
let inner: *const ManuallyDrop<T> = &self.inner;
mem::forget(self);
// SAFETY: Those pointers are valid for reads, because `Drop` didn't
// run, and properly aligned, because `Instrumented` isn't
// `#[repr(packed)]`.
let _span = unsafe { span.read() };
let inner = unsafe { inner.read() };
ManuallyDrop::into_inner(inner)
}
#[cfg(not(feature = "std-future"))]
self.inner
}
}
@@ -570,6 +638,8 @@ mod tests {
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.drop_span(expect::span().named("foo"))
.only()
.run_with_handle();
@@ -589,6 +659,8 @@ mod tests {
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.drop_span(expect::span().named("foo"))
.only()
.run_with_handle();
@@ -613,6 +685,8 @@ mod tests {
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.drop_span(expect::span().named("foo"))
.run_with_handle();
with_default(collector, || {
63 changes: 63 additions & 0 deletions tracing-futures/tests/std_future.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
use std::{future::Future, pin::Pin, task};

use futures::FutureExt as _;
use tracing::Instrument;
use tracing::{collect::with_default, Level};
use tracing_mock::*;
@@ -9,6 +12,8 @@ fn enter_exit_is_reasonable() {
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.drop_span(expect::span().named("foo"))
.only()
.run_with_handle();
@@ -26,6 +31,8 @@ fn error_ends_span() {
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.enter(expect::span().named("foo"))
.exit(expect::span().named("foo"))
.drop_span(expect::span().named("foo"))
.only()
.run_with_handle();
@@ -35,3 +42,59 @@ fn error_ends_span() {
});
handle.assert_finished();
}

#[test]
fn span_on_drop() {
#[derive(Clone, Debug)]
struct AssertSpanOnDrop;

impl Drop for AssertSpanOnDrop {
fn drop(&mut self) {
tracing::info!("Drop");
}
}

struct Fut(Option<AssertSpanOnDrop>);

impl Future for Fut {
type Output = ();

fn poll(mut self: Pin<&mut Self>, _: &mut task::Context<'_>) -> task::Poll<Self::Output> {
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")));
});
}
28 changes: 25 additions & 3 deletions tracing-subscriber/src/fmt/fmt_subscriber.rs
Original file line number Diff line number Diff line change
@@ -274,10 +274,32 @@ impl<C, N, E, W> Subscriber<C, N, E, W> {
}
}

/// 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
18 changes: 15 additions & 3 deletions tracing-subscriber/src/fmt/mod.rs
Original file line number Diff line number Diff line change
@@ -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<N, format::Format<L, T>, F, W> {
CollectorBuilder {
inner: self.inner.with_ansi(ansi),
1 change: 1 addition & 0 deletions tracing/Cargo.toml
Original file line number Diff line number Diff line change
@@ -38,6 +38,7 @@ pin-project-lite = "0.2.9"

[dev-dependencies]
criterion = { version = "0.3.6", default_features = false }
futures = { version = "0.3.21", default_features = false }
log = "0.4.17"
tracing-mock = { path = "../tracing-mock" }

86 changes: 73 additions & 13 deletions tracing/src/instrument.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
use crate::span::Span;
use core::pin::Pin;
use core::task::{Context, Poll};
use core::{future::Future, marker::Sized};
use core::{
future::Future,
marker::Sized,
mem::{self, ManuallyDrop},
pin::Pin,
task::{Context, Poll},
};
use pin_project_lite::pin_project;

#[cfg(feature = "std")]
@@ -18,7 +22,7 @@ pub trait Instrument: Sized {
/// `Instrumented` wrapper.
///
/// The attached [`Span`] will be [entered] every time the instrumented
/// [`Future`] is polled.
/// [`Future`] is polled or [`Drop`]ped.
///
/// # Examples
///
@@ -80,14 +84,17 @@ pub trait Instrument: Sized {
/// [disabled]: super::Span::is_disabled()
/// [`Future`]: std::future::Future
fn instrument(self, span: Span) -> Instrumented<Self> {
Instrumented { inner: self, span }
Instrumented {
inner: ManuallyDrop::new(self),
span,
}
}

/// Instruments this type with the [current] [`Span`], returning an
/// `Instrumented` wrapper.
///
/// The attached [`Span`] will be [entered] every time the instrumented
/// [`Future`] is polled.
/// [`Future`] is polled or [`Drop`]ped.
///
/// This can be used to propagate the current span when spawning a new future.
///
@@ -285,13 +292,55 @@ pin_project! {
///
/// [`Future`]: std::future::Future
/// [`Span`]: crate::Span
#[project = InstrumentedProj]
#[project_ref = InstrumentedProjRef]
#[derive(Debug, Clone)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Instrumented<T> {
// `ManuallyDrop` is used here to to enter instrument `Drop` by entering
// `Span` and executing `ManuallyDrop::drop`.
#[pin]
inner: T,
inner: ManuallyDrop<T>,
ilslv marked this conversation as resolved.
Show resolved Hide resolved
span: Span,
}

impl<T> PinnedDrop for Instrumented<T> {
fn drop(this: Pin<&mut Self>) {
let this = this.project();
let _enter = this.span.enter();
// SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because this isn't
// different from wrapping `T` in `Option` and calling
// `Pin::set(&mut this.inner, None)`, except avoiding
// additional memory overhead.
// 2. `ManuallyDrop::drop()` is safe, because
// `PinnedDrop::drop()` is guaranteed to be called only
// once.
unsafe { ManuallyDrop::drop(this.inner.get_unchecked_mut()) }
}
}
}

impl<'a, T> InstrumentedProj<'a, T> {
/// Get a mutable reference to the [`Span`] a pinned mutable reference to
/// the wrapped type.
fn span_and_inner_pin_mut(self) -> (&'a mut Span, Pin<&'a mut T>) {
// SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
// and `inner` is valid, because `ManuallyDrop::drop` is called
// only inside `Drop` of the `Instrumented`.
let inner = unsafe { self.inner.map_unchecked_mut(|v| &mut **v) };
(self.span, inner)
}
}

impl<'a, T> InstrumentedProjRef<'a, T> {
/// Get a reference to the [`Span`] a pinned reference to the wrapped type.
fn span_and_inner_pin_ref(self) -> (&'a Span, Pin<&'a T>) {
// SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
// and `inner` is valid, because `ManuallyDrop::drop` is called
// only inside `Drop` of the `Instrumented`.
let inner = unsafe { self.inner.map_unchecked(|v| &**v) };
(self.span, inner)
}
}

// === impl Instrumented ===
@@ -300,9 +349,9 @@ impl<T: Future> Future for Instrumented<T> {
type Output = T::Output;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
let _enter = this.span.enter();
this.inner.poll(cx)
let (span, inner) = self.project().span_and_inner_pin_mut();
let _enter = span.enter();
inner.poll(cx)
}
}

@@ -331,19 +380,30 @@ impl<T> Instrumented<T> {

/// Get a pinned reference to the wrapped type.
pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> {
self.project_ref().inner
self.project_ref().span_and_inner_pin_ref().1
}

/// Get a pinned mutable reference to the wrapped type.
pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
self.project().inner
self.project().span_and_inner_pin_mut().1
}

/// Consumes the `Instrumented`, returning the wrapped type.
///
/// Note that this drops the span.
pub fn into_inner(self) -> T {
self.inner
// To manually destructure `Instrumented` without `Drop`, we save
// pointers to the fields and use `mem::forget` to leave those pointers
// valid.
let span: *const Span = &self.span;
let inner: *const ManuallyDrop<T> = &self.inner;
mem::forget(self);
// SAFETY: Those pointers are valid for reads, because `Drop` didn't
// run, and properly aligned, because `Instrumented` isn't
// `#[repr(packed)]`.
let _span = unsafe { span.read() };
let inner = unsafe { inner.read() };
ManuallyDrop::into_inner(inner)
}
}

67 changes: 67 additions & 0 deletions tracing/tests/instrument.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
// These tests require the thread-local scoped dispatcher, which only works when
// we have a standard library. The behaviour being tested should be the same
// with the standard lib disabled.
#![cfg(feature = "std")]

use std::{future::Future, pin::Pin, task};

use futures::FutureExt as _;
use tracing::{collect::with_default, Instrument as _, Level};
use tracing_mock::*;

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn span_on_drop() {
#[derive(Clone, Debug)]
struct AssertSpanOnDrop;

impl Drop for AssertSpanOnDrop {
fn drop(&mut self) {
tracing::info!("Drop");
}
}

struct Fut(Option<AssertSpanOnDrop>);

impl Future for Fut {
type Output = ();

fn poll(mut self: Pin<&mut Self>, _: &mut task::Context<'_>) -> task::Poll<Self::Output> {
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")));
});
}