`Handle::dump()` panics with "RefCell already borrowed" when a traced task wakes a task (current_thread runtime)

[tillrohrmann]

Disclaimer: I used AI for analyzing this problem and creating the issue.

Version tokio 1.53.1, also reproduced on master (ea91b33)
Platform Linux x86_64 / aarch64

Description

Taking a task dump on a current_thread runtime panics if any traced task wakes a waker while it is being re-polled by the tracer:

thread '...' panicked at tokio/src/runtime/scheduler/current_thread/mod.rs:721:40:
RefCell already borrowed

Waking during a poll is legal, so any application that calls Handle::dump() on a current_thread runtime can hit this. We found it in production: a SIGUSR2-triggered task dump crashed a server (restatedev/restate#5235).

Reproducer

Drop this into tokio/tests/refcell_repro.rs and run RUSTFLAGS="--cfg tokio_unstable" cargo test --features full,taskdump --test refcell_repro:

#[test]
fn self_wake() {
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    rt.block_on(async {
        let task = tokio::spawn(std::future::poll_fn(|cx| {
            if tokio::runtime::Handle::is_tracing() {
                cx.waker().wake_by_ref();
            }
            std::task::Poll::<()>::Pending
        }));

        tokio::task::yield_now().await;
        let _dump = tokio::runtime::Handle::current().dump().await;

        task.abort();
    });
}

The dump().await panics, so the test fails there.

Two failure modes

1. The traced task wakes itself. The wake is deferred to the end of the poll, so the panic is raised in Harness::poll — outside the catch_unwind around the future.It escapes through Handle::dump() into whoever asked for the dump, and the traced task is left with a leaked ref-count and no queued notification.

Backtrace (trimmed):

5: RefCell::<Option<Box<Core>>>::borrow_mut
6: <Arc<current_thread::Handle> as Schedule>::schedule::{closure#0}
        at ./src/runtime/scheduler/current_thread/mod.rs:721:40
11: <Arc<current_thread::Handle> as Schedule>::schedule
        at ./src/runtime/scheduler/current_thread/mod.rs:719:9
12: <Arc<current_thread::Handle> as Schedule>::yield_now
        at ./src/runtime/task/mod.rs:322:14
13: Harness::<...>::poll
        at ./src/runtime/task/harness.rs:161:22
16: LocalNotified::<...>::run
        at ./src/runtime/task/mod.rs:523:13
17: trace::trace_owned::<...>
        at ./src/runtime/task/trace/mod.rs:475:64
...  Handle::dump
        at ./src/runtime/scheduler/current_thread/mod.rs:607:22

2. The traced task wakes a different, already-idle task. Now schedule() runs inside the future's own poll, so poll_future's catch_unwind swallows it. The dump reports success while the task is silently killed, and any lock held across the wake is poisoned:

dumper is_err=false
a.is_finished()=false
b.is_finished()=true          <-- silently terminated
lock poisoned=true
b: JoinError is_panic=true -> task 2 panicked with message "RefCell already borrowed"

This is what we saw in production. An h2 connection task woke a peer task while holding its streams mutex; the mutex was poisoned by the unwind, a later PoisonError unwrap() panicked, and a panic in a destructor escalated that to abort().

Root cause

current_thread::Handle::dump() holds a mutable borrow of the scheduler core across the whole tracing pass (current_thread/mod.rs:595-613):

let mut maybe_core = context.core.borrow_mut();   // borrow taken here
let core = ...;
let local = &mut core.tasks;
...
traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject)  // polls every task
    ...;
drop(maybe_core);                                // released only afterwards

trace_current_thread re-polls every live task, and Schedule::schedule for Arc<current_thread::Handle> takes the same borrow when a task is woken on its own runtime (current_thread/mod.rs:719-728]).

This only affects current_thread. The multi-thread scheduler moves the core out of the RefCell before running, so a wake during tracing finds None and falls back to the inject queue (multi_thread/worker.rs:1365). The normal current_thread polling path is fine too — Context::run_task puts the core into the RefCell and releases the borrow before task.run().

Suggested fix

Split draining from tracing so the core borrow is released before any task is polled. The borrow is only needed to drain the local queue, which must still happen before tracing to satisfy trace_owned's exclusive-access precondition:

let dequeued = {
    let mut maybe_core = context.core.borrow_mut();
    let core = /* ... */;
    drain_current_thread(&mut core.tasks, &self.shared.inject)
};  // borrow released before any poll

traces = trace_owned(&self.shared.owned, dequeued)
    .into_iter()
    .map(|(id, trace)| dump::Task::new(id, trace))
    .collect();

wake_deferred_tasks_and_free(context);

Wakes that arrive during tracing then take the normal path and land in the local queue, so the woken tasks are polled after the dump instead of panicking. I'm happy to open a PR.