Consider entering `Span` on `Drop` for `Instrumented`

[ilslv]

Feature Request

Crates

tracing, tracing-futures

Motivation

Sometimes it's desired to get Span inside Drop implementation. For example:

async {
    let _guard = deferred(|| {
        if std::thread::panicking() {
            tracing::warn!("thread panicked");
        }
    });

    // stuff..
}
.instrument(span)

The problem is there is difference between when Drop can be actually called: inside Future::poll() or on Drop of the entire Future:

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

use futures::FutureExt as _;
use tracing::{info_span, Instrument as _, Span};

fn main() {
    tracing_subscriber::fmt().init();

    let span = info_span!("span");

    // polled
    Fut(Some(LogSpanOnDrop))
        .instrument(span.clone())
        .now_or_never()
        .unwrap();

    // never polled
    let _ = Fut(Some(LogSpanOnDrop)).instrument(span);
}

#[derive(Clone, Debug)]
struct LogSpanOnDrop;

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

struct Fut(Option<LogSpanOnDrop>);

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(())
    }
}

Output:

2023-03-30T09:30:46.863224Z  INFO span: soc_common::future::instrumented::t: drop
2023-03-30T09:30:46.863279Z  INFO soc_common::future::instrumented::t: drop

Span is missing for a second log.

Proposal

Would you accept a PR adding Drop implementation for Instrumented, that will enter the Span and Drop inner Future?

Alternatives

Leave it as is.

[zohnannor]

Here's a contrived example of the same problem happening with #[tracing::instrument] showcasing two possible outcomes.

use std::future::{pending, ready};

use futures::FutureExt as _;

#[tracing::instrument]
async fn foo(switch: bool) -> i32 {
    let val = LogOnDrop;
    match switch {
        true => ready(1).await,
        false => pending::<i32>().await,
    }
}

#[tokio::test]
async fn main() {
    tracing_subscriber::fmt().init();

    // polled to completion
    assert!(foo(true).now_or_never().is_some());

    // polled once and thrown away mid-execution
    assert!(foo(false).now_or_never().is_none());
}

struct LogOnDrop;

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

Output:

2023-03-30T12:00:22.905817Z  INFO foo{switch=true}: soc_common::future::tests: drop
2023-03-30T12:00:22.905874Z  INFO soc_common::future::tests: drop

[hawkw]

Hmm, this isn't a bad idea. I agree that it's potentially desirable to have any events occurring inside a Drop implementation for an Instrumented future occur inside that span.

One potential challenge is that it would be necessary to actually ensure that the inner Future's drop implementation runs while the span is entered. This means we would need to actually drop the inner Future inside the drop method on Instrumented, while the guard is held, rather than after drop returns. We could solve this either by putting the Future inside an Option, which is taken in drop so that it can be dropped inside drop (haha), or using std::mem::ManuallyDrop. Using ManuallyDrop would require unsafe code, but using an Option would mean that we would need to unwrap the Option every time the future is polled, adding some small per-poll overhead. Using ManuallyDrop is probably a better solution here, since the behavior is simple enough and the use of unsafe code would be very low risk, as we can avoid the overhead of checking the Option on every poll.

[zohnannor]

We could possibly avoid Option check overhead by using .unwrap_or_else(|| unsafe { std::hint::unreachable_unchecked() }), or, alternatively, just .unwrap_or_else(|| unreachable!()), hinting at optimization to the compiler. ManuallyDrop solution would be just as good (in terms of performance), but I don't see the reason to it given the simpler option (pun intended).

[hawkw]

IMO using ManuallyDrop is probably simpler and more idiomatic than unreachable_unchecked — ManuallyDrop is intended for situations where it's necessary to manually control drop order, which is what we need to do here. But either solution would be fine.

[zohnannor]

Opened PR #2561. I propose to discuss the implementation details of the feature there.

[ilslv]

I actually myself already implemented version with ManuallyDrop to avoid runtime and memory overhead https://github.com/tokio-rs/tracing/pull/2562

[tyranron]

@ilslv as being released in 0.1.38, I guess we may close this issue now.

[ilslv]

@tyranron I would like to leave this open as a reminder to do a similar thing to the WithDispatch.

UPD: postponing the WithDispatch changes because it will be a breaking change: https://github.com/tokio-rs/tracing/issues/2578