JoinHandle、取消与 JoinError

主线位置：任务主线的最后一段，「完成，JoinHandle 拿到结果」。也包括主线的岔路：任务被 abort。

sequenceDiagram
  participant J as JoinHandle（等结果的任务）
  participant S as 任务状态
  participant W as worker（跑任务）
  J->>S: poll：未完成，存 waker，置 JOIN_WAKER
  W->>S: transition_to_complete（置 COMPLETE）
  W->>J: JOIN_WAKER 为 1，wake_join
  W->>S: 清 JOIN_WAKER，waker 字段交还 JoinHandle
  J->>S: 再次 poll：已完成，take_output

JoinHandle 是什么

task/join.rs 第 163 行：

pub struct JoinHandle<T> {
    raw: RawTask,
    _p: PhantomData<T>,
}

一个 RawTask 加上输出类型。它占任务的一个引用计数，状态里的 JOIN_INTEREST 位表示它还活着。T: Send 时它才是 Send 的：输出要能跟着它跨线程。

poll：输出从 vtable 里「倒」出来

第 327 行：

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    ready!(crate::trace::trace_leaf(cx));
    let mut ret = Poll::Pending;

    // Keep track of task budget
    let coop = ready!(crate::task::coop::poll_proceed(cx));

    // Try to read the task output. If the task is not yet complete, the
    // waker is stored and is notified once the task does complete.
    //
    // The function must go via the vtable, which requires erasing generic
    // types. To do this, the function "return" is placed on the stack
    // **before** calling the function and is passed into the function using
    // `*mut ()`.
    //
    // Safety:
    //
    // The type of `T` must match the task's output type.
    unsafe {
        self.raw.try_read_output(&mut ret, cx.waker());
    }

    if ret.is_ready() {
        coop.made_progress();
    }

    ret
}

两个细节：

• .await 一个 JoinHandle 也消耗协作式预算。预算用完时，即使任务已经完成，也先返回 Pending，让出 worker。
• 输出怎么穿过类型擦除：vtable 里的函数签名不能带 T，所以调用方先在自己栈上放一个 Poll<Result<T>>，把它的地址当 *mut () 传进去。vtable 那头的 try_read_output::<T, S> 知道真实类型，把指针转回来，往里写。两边类型对得上，靠的是 JoinHandle<T> 只能由 new_task 和对应的 future 一起造出来。

真正干活的是 can_read_output：任务已完成就返回 true，调用方从 Stage::Finished 里取走输出（换成 Consumed）；没完成就把 waker 存进 Trailer，返回 false。

完成后再 poll 一次会 panic（JoinHandle polled after completion），因为输出已经被取走了。

JOIN_WAKER：两个线程，一个 waker 字段

Trailer 里的 waker 字段（UnsafeCell<Option<Waker>>，没有锁）有两个访问者：

• JoinHandle 所在的线程：每次 poll 可能要写一个新 waker。
• 跑任务的 worker：任务完成时要读它，调 wake。

这两件事可能同时发生。tokio 用状态里的 JOIN_WAKER 位做访问控制，task/mod.rs 的模块文档里写了 7 条规则，归纳起来是：

表里假设 JoinHandle 还在。它已经被 drop（JOIN_INTEREST 为 0）而任务完成了，waker 字段归运行时独占，由运行时负责清掉。

另外两条：只有任务未完成时 JoinHandle 才能改 JOIN_WAKER，只有任务完成后运行时才能改它。 完成这一刻（COMPLETE 从 0 变 1）就是所有权交接点。

JoinHandle 要换 waker 时，按三步走：

1. 如果 JOIN_WAKER 是 1，先把它清成 0，拿回独占权（unset_waker）。
2. 写入新 waker。
3. 把 JOIN_WAKER 置回 1（set_join_waker）。

第 1、3 步都是 CAS，任务未完成时才成功。任一步失败，说明任务就在这期间完成了：第 3 步失败时把刚写的 waker 清掉，然后按「已完成」处理，直接去读输出。can_read_output 里还有一个快路径：已存的 waker 和这次的 will_wake 相同，就什么都不做。同一个任务反复 poll 同一个 JoinHandle 时，大多走这条路。

运行时这边，complete 里看到 JOIN_WAKER 为 1 就调 wake_join，然后用 unset_waker_after_complete（一条 fetch_and）清掉它，把字段交还给 JoinHandle。如果清的时候发现 JOIN_INTEREST 已经是 0（JoinHandle 刚好在这期间被 drop 了），就由运行时自己把 waker drop 掉。

drop JoinHandle：快路径

tokio::spawn(fut); 不接返回值，是最常见的写法之一。这时 JoinHandle 在 spawn 返回后立刻 drop，Drop 先试快路径：

impl<T> Drop for JoinHandle<T> {
    fn drop(&mut self) {
        if self.raw.state().drop_join_handle_fast().is_ok() {
            return;
        }

        self.raw.drop_join_handle_slow();
    }
}

drop_join_handle_fast：

pub(super) fn drop_join_handle_fast(&self) -> Result<(), ()> {
    use std::sync::atomic::Ordering::Relaxed;

    // Relaxed is acceptable as if this function is called and succeeds,
    // then nothing has been done w/ the join handle.
    //
    // The moment the join handle is used (polled), the `JOIN_WAKER` flag is
    // set, at which point the CAS will fail.
    //
    // Given this, there is no risk if this operation is reordered.
    self.val
        .compare_exchange_weak(
            INITIAL_STATE,
            (INITIAL_STATE - REF_ONE) & !JOIN_INTEREST,
            Release,
            Relaxed,
        )
        .map(|_| ())
        .map_err(|_| ())
}

它赌的是：状态还是一模一样的初始值，即任务还没被 poll 过、没人 clone 过 waker、JoinHandle 也没被 poll 过。这时不可能有输出要丢、也不可能有 waker 要清，只要扣一个引用、清 JOIN_INTEREST，一次 CAS 完事。用 compare_exchange_weak 是因为失败也无所谓：伪失败只是多走一次慢路径。

慢路径 drop_join_handle_slow 先做 transition_to_join_handle_dropped：清 JOIN_INTEREST，任务未完成时顺便清 JOIN_WAKER。返回两个布尔值：

• drop_output：任务已完成，输出还在，由 JoinHandle 负责 drop。drop 包在 catch_unwind 里，panic 被吞掉：用户都不要这个 JoinHandle 了，显然也不关心输出析构时的 panic。
• drop_waker：waker 字段归自己独占了，清掉它。

为什么输出要由 JoinHandle 当场 drop，而不是留在任务里等释放时一起 drop？源码注释给的理由：输出可能不是 Send 的（LocalSet 上的任务，见 JoinSet、LocalSet、task_local），它只能在 JoinHandle 所在的线程上 drop；留在任务里，最后一个引用可能是某个在任意线程上被 drop 的 waker。

abort：一次置位，交给调度器

JoinHandle::abort 和 AbortHandle::abort 都是一行 self.raw.remote_abort()（harness.rs 第 118 行）：

pub(super) fn remote_abort(&self) {
    if self.state().transition_to_notified_and_cancel() {
        // The transition has created a new ref-count, which we turn into
        // a Notified and pass to the task.
        //
        // Since the caller holds a ref-count, the task cannot be destroyed
        // before the call to `schedule` returns even if the call drops the
        // `Notified` internally.
        self.schedule();
    }
}

abort 自己从不 drop future，只是置 CANCELLED，必要时让任务进队列，由调度器在任务所属的线程上完成取消。按调用时任务的状态，结果是：

最后两行就是文档那句「已经完成的任务被取消，.await 仍可能正常返回」。所以 abort 只是「请求」，想确认任务确实没有产生副作用，要看 .await 的结果是不是 is_cancelled()。

let a = tokio::spawn(async {
  std::thread::sleep(Duration::from_millis(100));
  42
});
let b = tokio::spawn(async {
  std::thread::sleep(Duration::from_millis(100));
  tokio::task::yield_now().await;
  42
});
std::thread::sleep(Duration::from_millis(10));
a.abort();
b.abort();
println!("A: {:?}", a.await);
println!("B: {:?}", b.await);

spawn_blocking 的任务没有 .await 点，取消只能在它开始运行之前生效；一旦开始，abort 就没有作用。

AbortHandle 也只包了一个 RawTask，JoinHandle::abort_handle() 和 AbortHandle::clone 都是 ref_inc，drop 时扣掉。它不碰 JOIN_INTEREST，所以只拿着 AbortHandle 不影响任务完成时丢弃输出。

JoinError

task/error.rs：

pub struct JoinError {
    repr: Repr,
    id: Id,
}

enum Repr {
    Cancelled,
    Panic(SyncWrapper<Box<dyn Any + Send + 'static>>),
}

只有两种：取消，或者 panic。panic 的 payload 是 Box<dyn Any + Send>，不是 Sync，但 JoinError 需要是 Sync（才能放进 anyhow::Error 这类要求 Send + Sync 的错误类型里）。SyncWrapper 的办法是不提供 &T 访问：拿着 &SyncWrapper 什么也做不了，跨线程共享引用也就是安全的，只有 into_inner 按值取出。

把 panic 原样抛到等待方，标准写法是：

match handle.await {
  Ok(v) => v,
  Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()),
  Err(e) => return Err(e.into()), // 取消：转成调用方自己的错误类型
}

Display 的输出带任务 ID，排查日志时能和 tokio::task::id() 对上：

task 13 panicked with message "boom"
task 14 was cancelled

动脑筋答案

完整程序：

use std::time::Duration;

#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() {
  // A：一次 poll 里同步忙 100ms，然后直接返回
  let a = tokio::spawn(async {
    std::thread::sleep(Duration::from_millis(100));
    42
  });
  // B：同样忙 100ms，但返回之前还有一个 .await
  let b = tokio::spawn(async {
    std::thread::sleep(Duration::from_millis(100));
    tokio::task::yield_now().await;
    42
  });
  // main 不是 worker 线程，这里同步睡 10ms 不影响两个任务
  std::thread::sleep(Duration::from_millis(10));
  a.abort();
  b.abort();
  println!("A: {:?}", a.await);
  println!("B: {:?}", b.await);
}

M2 Max、rustc 1.99.0、--release，连跑 5 次，每次都是：

A: Ok(42)
B: Err(JoinError::Cancelled(Id(4)))

• A：abort 时 A 正在 worker 上那次 poll 里睡觉。abort 只置 CANCELLED，打断不了它。100ms 后这次 poll 返回 Ready(42)，poll_inner 拿到 Ready 直接走完成流程，不看 CANCELLED，结果保留下来。这就是上表「这次 poll 返回了 Ready」那一行。
• B：睡完以后碰到 yield_now().await，这次 poll 返回 Pending，transition_to_idle 看到 CANCELLED，返回 Cancelled，poll 的线程当场 drop future。B 的 42 永远算不出来。
• Id(4) 是 B 的任务 ID，全局自增，数字因程序而异。

这个结果依赖时序：A 得在 abort 之前就开始被 poll。第一版程序在 main 里用的是 tokio::time::sleep 等 10ms，结果不稳定，5 次里 B 有 4 次是 Ok(42)。在 abort 前打印耗时，5 次都在 100ms 到 105ms 之间：两个 worker 都在同步睡觉，10ms 的定时器要等其中一个睡完空下来才触发，这时 B 可能已经跑完了。换成 std::thread::sleep 以后，main 线程自己计时，5 次都和上面一样。

小结

• JoinHandle 是一个 RawTask 加输出类型。poll 消耗协作式预算；输出靠「调用方在栈上放好槽位，按 *mut () 传进 vtable」穿过类型擦除。
• waker 字段由 JOIN_WAKER 位做访问控制：任务未完成时归 JoinHandle，完成后归运行时，用完交还。will_wake 相同时跳过重存。
• drop JoinHandle 先试快路径：状态还是初始值就一次 CAS 完事，覆盖了 tokio::spawn(fut); 这种不要返回值的写法。慢路径负责丢弃输出和 waker，非 Send 输出因此不会跑到别的线程上析构。
• abort 只置 CANCELLED，由调度器在正确的线程上 drop future。poll 中的任务要等这次 poll 返回；这次 poll 若产出了结果，结果优先。
• JoinError 只有取消和 panic 两种，SyncWrapper 让装着 panic payload 的它也是 Sync。