Tasks spawned via spawn_local have useless location information

[kamirr]

Version
1.21.1

Platform

$ uname -a
Linux <privacy snip> 5.15.0-46-generic #49~20.04.1-Ubuntu SMP Thu Aug 4 19:15:44 UTC 2022 x86_64 x86_64 x86_64 GNU/Linux

Description
Tracing is not implemented properly for local tasks. Consider the following complete example:

async fn run() {
    for _ in 0..10 {
        tokio::task::spawn_local(async { println!("anon future"); }).await.unwrap();
    }
}

#[tokio::main(flavor="current_thread")]
async fn main() {
    console_subscriber::init();

    let local_set = tokio::task::LocalSet::new();
    local_set.run_until(run()).await;

    // program needs to remain live for tokio-console to work properly
    loop {
        tokio::time::sleep(std::time::Duration::from_secs_f32(60.0)).await;
    }
}

(Buildable project: tokio-spawn-local-repro.zip, run with RUSTFLAGS="--cfg tokio_unstable" cargo run)
tokio-console produces the following output:

From my investigation it appears that the problem stems from a lacking application of #[track_caller] in tokio/src/task/local.rs:

#[track_caller]
pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
where F: Future + 'static,
        F::Output: 'static
{
    CURRENT.with(|maybe_cx| { // line 317
        match maybe_cx.get() {
            None => panic!("`spawn_local` called from outside of a `task::LocalSet`"),
            Some(cx) => cx.spawn(future, name)
        }
    })
}

Notice that the closure has no #[track_caller] attribute, hence breaking the chain leading up to the call of Location::caller in Contex::spawn, meaning all locations will refer to the spawn_local_inner function body. #[track_caller] attribute cannot be currently applied to closures (it's an unstable feature), so another approach is necessary.

[hds]

I hadn't thought of this before, but for all the cases where the only closure that a call stack passes through is to access a LocalKey, I think we could do something like the following:

    #[track_caller]
    pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
    where F: Future + 'static,
          F::Output: 'static
    {
        CURRENT.with(|maybe_cx| {
            maybe_cx.get().map(|cx| cx.spawn(future, name))
        })
        .expect("`spawn_local` called from outside of a `task::LocalSet`")
    }

This brings the panic (in the form of an except, we could use a match statement to keep an explicit panic) outside the closure. It's certainly not easier to read, but it does have it's advantages.

[hawkw]

I hadn't thought of this before, but for all the cases where the only closure that a call stack passes through is to access a LocalKey, I think we could do something like the following:

    #[track_caller]
    pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
    where F: Future + 'static,
          F::Output: 'static
    {
        CURRENT.with(|maybe_cx| {
            maybe_cx.get().map(|cx| cx.spawn(future, name))
        })
        .expect("`spawn_local` called from outside of a `task::LocalSet`")
    }

This brings the panic (in the form of an except, we could use a match statement to keep an explicit panic) outside the closure. It's certainly not easier to read, but it does have it's advantages.

@hds this would fix the panic location, but it doesn't fix the incorrect spawn location for tokio-console. I think the only solution for that would be to change Context::spawn to also take a std::panic::Location argument, and get the location in spawn_local_inner prior to the closure, and then pass it to Context::spawn, like this:

    #[track_caller]
    pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
    where F: Future + 'static,
          F::Output: 'static
    {
        let location = std::panic::Location::caller();
        CURRENT.with(move |maybe_cx| {
            maybe_cx.get().map(|cx| cx.spawn(future, name, location))
        })
        .expect("`spawn_local` called from outside of a `task::LocalSet`")
    }

We probably need to do a check of other spawning functions that access a thread local similarly.

[Darksonn]

Another simpler solution is to just clone the Rc in CURRENT and call spawn outside the closure.

let maybe_cx = CURRENT.with(|maybe_cx| maybe_cx.get());

[hds]

And here I was thinking that we needed to keep that context inside the closure for some reason...

So this should solve both issues:

    #[track_caller]
    pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
    where F: Future + 'static,
          F::Output: 'static
    {
        match CURRENT.with(|maybe_cx| maybe_cx.get()) {
            None => panic!("`spawn_local` called from outside of a `task::LocalSet`"),
            Some(cx) => cx.spawn(future, name)
       }
    }

[hds]

@kamirr I've added what should be a fix for both these issues in #5034. If you'd like to have a look at it that would be awesome.