JoinSet、LocalSet、task_local

主线位置：任务主线的外围。JoinSet 管一组任务的「完成 → 取结果」；LocalSet 是任务进队列、被 poll 的另一条路，不经过多线程调度器；task_local! 作用在单次 poll 内部。

flowchart TB
  wake[本地任务被唤醒] --> q1{"CURRENT 指向这个 LocalSet？"}
  q1 -->|"是，且正在 poll"| lq[进本地队列]
  q1 -->|否| q2{当前线程是 owner？}
  q2 -->|是| lq2[进本地队列，唤醒 LocalSet]
  q2 -->|否| rq[进远程队列，唤醒 LocalSet]
  lq --> tick[LocalSet 的 tick]
  lq2 --> tick
  rq --> tick
  tick -->|每次 poll 最多 61 个| run[在 owner 线程上 run]

JoinSet：只 poll 被唤醒的那几个

JoinSet 只有一个字段：

pub struct JoinSet<T> {
    inner: IdleNotifiedSet<JoinHandle<T>>,
}

最朴素的实现是一个 Vec<JoinHandle<T>>，join_next 时挨个 poll。一千个任务里只有一个完成了，也要 poll 一千次。IdleNotifiedSet 解决的就是这个问题：它记住哪些 JoinHandle 被唤醒过，join_next 只看这些。

两条链表

IdleNotifiedSet 的共享部分是一把锁里的两条侵入式链表和一个 waker（ListsInner）：

struct ListsInner<T> {
    notified: LinkedList<ListEntry<T>>,
    idle: LinkedList<ListEntry<T>>,
    /// Whenever an element in the `notified` list is woken, this waker will be
    /// notified and consumed, if it exists.
    waker: Option<Waker>,
}

每个元素是一个 Arc<ListEntry<T>>，里面放着 JoinHandle、链表指针、指回共享部分的 Arc，以及一个 my_list 字段记录自己在哪条链表上（枚举 List：Notified、Idle 或已取出的 Neither）。元素个数 length 不在锁里，因为只有持有 &mut JoinSet 的一方会增删元素。

stateDiagram-v2
  [*] --> Idle: spawn 后 insert_idle
  Idle --> Notified: 任务完成，entry 的 waker 被调用
  Notified --> Idle: join_next 取出准备 poll
  Idle --> Neither: poll 得到 Ready，remove
  Notified --> Neither: drop JoinSet，drain
  Idle --> Neither: drop JoinSet，drain
  Neither --> [*]

插入（insert）：JoinSet::spawn 就是 tokio::spawn 之后把 JoinHandle 放进 idle 链表，再用这个 entry 自己做的 waker 调 set_join_waker。任务完成时叫醒的不是用户的任务，而是这个 entry。如果设 waker 时任务已经完成了，set_join_waker 会立即唤醒一次。

唤醒（wake_by_ref）：加锁，如果 entry 在 idle 链表上，就把它挪到 notified 链表头部，然后取出 JoinSet 记下的那个 waker，在锁外调用它，这才叫醒等在 join_next 上的用户任务。已经在 notified 链表上的 entry 再被唤醒，什么都不做。

取出（pop_notified）：加锁，记下调用方的 waker（will_wake 相同就不替换，省一次克隆），从 notified 链表尾部弹出一个，放回 idle 链表。头部进、尾部出，所以 join_next 按被唤醒的先后顺序返回结果，先完成的先返回，不会有任务一直排不到。

poll_join_next 拿到 entry 后，用 entry 的 waker poll 那个 JoinHandle：

let res = entry.with_value_and_context(|jh, ctx| Pin::new(jh).poll(ctx));

if let Poll::Ready(res) = res {
    let _entry = entry.remove();
    Poll::Ready(Some(res))
} else {
    // A JoinHandle generally won't emit a wakeup without being ready unless
    // the coop limit has been reached. We yield to the executor in this
    // case.
    cx.waker().wake_by_ref();
    Poll::Pending
}

被唤醒的 JoinHandle 一般都已经就绪。例外是协作式预算用完：JoinHandle::poll 会扣预算（见协作式调度预算），扣不到就返回 Pending，并安排稍后唤醒 entry 的 waker，entry 会重新回到 notified 链表。JoinSet 这边叫醒自己再让出一次。

join_next 是 cancel safe 的：entry 只有在 JoinHandle 返回 Ready 的同一次 poll 里才被 remove，结果拿到手和从集合里删掉之间没有 .await。

drop 时全部 abort

impl<T> Drop for JoinSet<T> {
    fn drop(&mut self) {
        self.inner.drain(|join_handle| join_handle.abort());
    }
}

drain 在锁里把两条链表上的 entry 全部挪到一条临时链表，再逐个取出调用闭包。这是 JoinSet 和一堆 JoinHandle 最大的区别：drop JoinHandle 只是放手（任务继续跑），drop JoinSet 会取消所有还没完成的任务。不想取消就先调 detach_all。abort 只是置位并调度任务，future 要等任务下一次被运行时处理时才 drop（见 JoinHandle、取消与 JoinError）。

实验，三个任务按 30、10、20 毫秒完成；然后 drop 一个还有任务在跑的 JoinSet：

use std::time::Duration;
use tokio::task::JoinSet;
use tokio::time::sleep;

struct Guard(&'static str);
impl Drop for Guard {
  fn drop(&mut self) {
    println!("dropped: {}", self.0);
  }
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  let mut set = JoinSet::new();
  for (name, ms) in [("a", 30), ("b", 10), ("c", 20)] {
    set.spawn(async move {
      sleep(Duration::from_millis(ms)).await;
      name
    });
  }
  while let Some(r) = set.join_next().await {
    print!("{} ", r.unwrap());
  }
  println!();

  let mut set = JoinSet::new();
  set.spawn(async {
    let _g = Guard("task in JoinSet");
    sleep(Duration::from_secs(3600)).await;
  });
  tokio::task::yield_now().await; // let the task start
  println!("dropping JoinSet");
  drop(set);
  tokio::task::yield_now().await; // let the runtime process the abort
  println!("after drop");
}

b c a 
dropping JoinSet
dropped: task in JoinSet
after drop

结果按完成顺序返回，不是 spawn 顺序。drop JoinSet 之后，那个要睡一小时的任务在下一次调度时就被取消，它的局部变量随之析构。

LocalSet：跑在任务里的小调度器

tokio::spawn 要求 Send，因为多线程调度器会把任务交给任何一个 worker，还会被别的 worker 偷走。Rc、RefCell 这类 !Send 的值跨 .await 持有，future 就不是 Send 的。LocalSet 提供另一个执行器：任务只在创建 LocalSet 的那个线程上被 poll。

结构

pub struct LocalSet {
    /// Current scheduler tick.
    tick: Cell<u8>,

    /// State available from thread-local.
    context: Rc<Context>,

    /// This type should not be Send.
    _not_send: PhantomData<*const ()>,
}

（LocalSet，外层还包着一个 cfg_rt!。）跨线程共享的部分在 Shared：

两个队列初始容量都是 INITIAL_CAPACITY = 64。远程队列是 Option，LocalSet drop 时把它取走，之后其他线程的唤醒就直接丢弃。

LocalSet 不是运行时，它没有线程，也没有 IO 和定时器驱动。它本身是一个 Future（impl Future for LocalSet），或者通过 run_until 包住一个 future；block_on 就是 rt.block_on(self.run_until(future))。宿主运行时 poll 这个 future 时，LocalSet 才去跑自己队列里的任务。因为 LocalSet 是 !Send 的，包含它的 future 也是 !Send，不能交给 tokio::spawn，通常放在某个运行时的 block_on 里（current_thread 或多线程都可以）。

tick：一次最多 61 个

tick：

fn tick(&self) -> bool {
    for _ in 0..MAX_TASKS_PER_TICK {
        // Make sure we didn't hit an unhandled panic
        assert!(!self.context.unhandled_panic.get(), "a spawned task panicked and the LocalSet is configured to shutdown on unhandled panic");

        match self.next_task() {
            // ...
            Some(task) => crate::task::coop::budget(|| task.run()),
            // ...
            None => return false,
        }
    }

    true
}

每次被 poll 最多跑 MAX_TASKS_PER_TICK = 61 个任务，每个任务各有一份 128 的预算。61 个跑完队列还没空，就 wake_by_ref 自己、返回 Pending，把线程还给宿主运行时。这和单个任务的协作式预算是同一个道理：整个 LocalSet 在宿主看来只是一个任务，它不让出，宿主的 IO 驱动、定时器和其他任务就都动不了。

next_task 每 REMOTE_FIRST_INTERVAL = 31 个 tick 先看远程队列，平时先看本地队列，和 current_thread 用同一个数，理由也一样：本地任务互相唤醒能让本地队列一直不空，远程队列里的任务会饿死。

use std::cell::Cell;
use std::rc::Rc;
use tokio::task::{spawn_local, yield_now, LocalSet};

fn main() {
  let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
  let local = LocalSet::new();
  local.block_on(&rt, async {
    let done = Rc::new(Cell::new(0u32));
    for _ in 0..200 {
      let done = done.clone();
      spawn_local(async move { done.set(done.get() + 1) });
    }
    let mut seen = Vec::new();
    while done.get() < 200 {
      yield_now().await;
      seen.push(done.get());
    }
    println!("{seen:?}");
  });
}

!Send 的 future 为什么安全

任务的 future 不是 Send 的，但唤醒它的 Waker 一定是 Send + Sync，可能在任何线程上被调用。tokio 用下面几层保证 future 只在 owner 线程上被 poll：

1. LocalSet 是 !Send。 _not_send: PhantomData<*const ()> 让它不能被移到别的线程，而 poll 任务只发生在 LocalSet::poll / RunUntil::poll 里，也就只可能发生在创建它的线程上。

2. spawn_local 只能在 LocalSet 里面调。 线程本地变量 CURRENT 只在 LocalSet 被 poll 期间（或 LocalSet::enter 的守卫存活期间）指向它。spawn_local_inner 先看当前运行时是不是 LocalRuntime，不是就读 CURRENT，读不到就 panic：

   Copy

   match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) {
       None => panic!("`spawn_local` called from outside of a `task::LocalSet` or `runtime::LocalRuntime`"),
       Some(cx) => cx.spawn(future.take().unwrap(), meta)
   }
   

3. 调度按线程分流。 Shared::schedule 就是本页开头那张图：当前正在 poll 这个 LocalSet，直接进本地队列；在 owner 线程上但没在 poll，进本地队列并唤醒 LocalSet；在别的线程上，只能进加锁的远程队列并唤醒。别的线程只放入 Notified，从不 poll。

4. 只有 LocalNotified 能运行。 本地任务挂在 LocalOwnedTasks 上，它用 UnsafeCell 而不是锁，并且是 !Send + !Sync；它的 bind 不要求 future 是 Send。next_task 取出的 Notified 必须经过 assert_owner 换成 LocalNotified（带 PhantomData<*const ()>，不能离开当前线程）才能 run()。assert_owner 用 assert_eq! 检查任务的 owner id 是这个 LocalSet 的，release 构建里也生效。

5. 调试断言。 访问本地队列和 LocalOwnedTasks 前会调 assert_called_from_owner_thread 比较线程 ID。这是 debug_assert!，只在 debug 构建里检查，真正的保证来自前面几条。

Shared 里有 UnsafeCell，本来不是 Sync，源码用一行 unsafe impl Sync for Shared {} 绕过，上面的注释是：「This is safe because (and only because) we pinky pwomise to never touch the local run queue except from the thread that owns the LocalSet.」

LocalRuntime

1.54.1 里还有 LocalRuntime，通过 Builder::build_local 创建。在这个版本里它不需要 tokio_unstable。它内部就是一个 current_thread 调度器，额外记下创建线程的 ID（local_tid），自身用 PhantomData<*mut u8> 做成 !Send + !Sync。在它上面可以直接 spawn_local，不需要套一层 LocalSet；spawn_local_inner 里 handle.is_local() 那个分支就是给它的，在别的线程上调用会 panic「Local tasks can only be spawned on a LocalRuntime from the thread the runtime was created on」。用 new_multi_thread() 调 build_local 会 panic。文档说明它和 LocalSet 不兼容，不要在 LocalRuntime 里驱动 LocalSet。

实验：

fn main() {
  std::panic::set_hook(Box::new(|info| {
    println!("panic: {}", info.payload().downcast_ref::<&str>().copied()
      .or(info.payload().downcast_ref::<String>().map(|s| s.as_str())).unwrap_or("?"));
  }));
  let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();

  // 1. inside a runtime, but not inside a LocalSet
  let r = std::panic::catch_unwind(|| {
    rt.block_on(async {
      tokio::task::spawn_local(async {});
    })
  });
  println!("1. ok = {}", r.is_ok());

  // 2. inside a LocalSet: works, and !Send values are fine
  let local = tokio::task::LocalSet::new();
  local.block_on(&rt, async {
    let rc = std::rc::Rc::new(5);
    let v = tokio::task::spawn_local(async move { *rc + 1 }).await.unwrap();
    println!("2. spawn_local in LocalSet -> {v}");
  });

  // 3. LocalRuntime (no tokio_unstable needed in 1.54.1)
  let lrt = tokio::runtime::Builder::new_current_thread()
    .build_local(Default::default())
    .unwrap();
  let v = lrt.block_on(async {
    let rc = std::rc::Rc::new(7);
    tokio::task::spawn_local(async move { *rc * 2 }).await.unwrap()
  });
  println!("3. spawn_local on LocalRuntime -> {v}");
}

不加 --cfg tokio_unstable 构建运行：

panic: `spawn_local` called from outside of a `task::LocalSet` or `runtime::LocalRuntime`
1. ok = false
2. spawn_local in LocalSet -> 6
3. spawn_local on LocalRuntime -> 14

task_local!：每次 poll 换进换出

task_local! 展开成一个 static LocalKey，里面是一个普通的线程本地变量（__task_local_inner!）：

$vis static $name: $crate::task::LocalKey<$t> = {
    std::thread_local! {
        static __KEY: std::cell::RefCell<Option<$t>> = const { std::cell::RefCell::new(None) };
    }

    $crate::task::LocalKey { inner: __KEY }
};

任务会在不同线程上被 poll，值不能常驻在某个线程的线程本地变量里。scope 返回一个 TaskLocalFuture，值平时存在这个 future 自己的 slot 字段里。每次被 poll 时，通过 scope_inner 把值换进线程本地变量，poll 内层 future，再换回来：

fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr>
where
    F: FnOnce() -> R,
{
    struct Guard<'a, T: 'static> {
        local: &'static LocalKey<T>,
        slot: &'a mut Option<T>,
    }

    impl<'a, T: 'static> Drop for Guard<'a, T> {
        fn drop(&mut self) {
            // ...
            self.local.inner.with(|inner| {
                let mut ref_mut = inner.borrow_mut();
                mem::swap(self.slot, &mut *ref_mut);
            });
        }
    }

    self.inner.try_with(|inner| {
        inner
            .try_borrow_mut()
            .map(|mut ref_mut| mem::swap(slot, &mut *ref_mut))
    })??;

    let guard = Guard { local: self, slot };

    let res = f();

    drop(guard);

    Ok(res)
}

• 用 mem::swap 而不是「写入、之后清空」：换进去时，线程本地变量里原来的值（外层 scope 的值，或者 None）被换到 slot 里暂存，换回来时原样恢复。所以嵌套的 scope 能正确地遮蔽和还原外层的值。
• 换回放在 Guard 的 Drop 里：内层 future panic 时，栈展开也会把值换回 slot，线程本地变量不会残留这个任务的值，被同一线程上的下一个任务看到。
• 换进去用的是 try_borrow_mut。如果正在 KEY.with(|v| ...) 的闭包里又进入同一个 key 的 scope，RefCell 已经被借用，就 panic「cannot enter a task-local scope while the task-local storage is borrowed」。

TaskLocalFuture 的 PinnedDrop 也走一次 scope_inner：drop 内层 future 时值仍然可见，内层 future 的析构函数里可以读它。

这个设计的直接后果是：值只在 TaskLocalFuture::poll 的调用栈里可见。spawn 出去的任务由调度器单独 poll，不在这个调用栈里，自然看不到；.await 让出之后，值也会被换回 slot，同一线程上的其他任务看不到。

tokio::task_local! {
  static REQ_ID: u32;
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  REQ_ID.scope(42, async {
    println!("inside scope:        {:?}", REQ_ID.try_get());
    let h = tokio::spawn(async { REQ_ID.try_get().is_ok() });
    println!("in spawned task:     visible = {}", h.await.unwrap());
    let v = tokio::spawn(REQ_ID.scope(REQ_ID.get(), async { REQ_ID.get() }));
    println!("re-scoped in task:   {}", v.await.unwrap());
    tokio::task::yield_now().await;
    println!("after a yield:       {:?}", REQ_ID.try_get());
  }).await;
  println!("outside scope:       visible = {}", REQ_ID.try_get().is_ok());
}

inside scope:        Ok(42)
in spawned task:     visible = false
re-scoped in task:   42
after a yield:       Ok(42)
outside scope:       visible = false

要把值带进子任务，只能像第三行那样在 spawn 时显式再包一层 scope。yield_now 之后值还在，是因为下一次 poll 又换了进来。

动脑筋答案

M2 Max、rustc 1.99.0、--release，连跑 3 次，输出都一样：

[61, 122, 183, 200]

• block_on 第一次 poll RunUntil：先 poll 传进来的 future，它一口气 spawn_local 了 200 个任务（这时 CURRENT 指向这个 LocalSet，任务直接进本地队列），然后在 yield_now 处返回 Pending。接着 tick 跑了 61 个，队列没空，返回 true，RunUntil 叫醒自己。
• 第二次 poll：future 从 yield_now 醒来，看到 61，记下，再次让出；tick 又跑 61 个。第三次看到 122，第四次看到 183。
• 第四次 poll 的 tick 只剩 17 个，跑完队列空了。第五次 poll 时 future 看到 200，循环结束。

数组里的每一项就是 61 的整数倍，直到队列跑空。如果没有这个上限，第一次 tick 就会把 200 个全跑完，这期间宿主运行时的 IO 驱动和定时器一次也轮不到。

小结

• JoinSet 基于 IdleNotifiedSet：每个 JoinHandle 用自己 entry 的 waker，任务完成时 entry 从 idle 链表挪到 notified 链表；join_next 只从 notified 链表按 FIFO 取，不扫描全部任务。drop JoinSet 会 abort 所有任务。
• LocalSet 是一个 future，在宿主任务里每次 poll 最多跑 61 个任务，每 31 个 tick 先看一次远程队列。
• !Send 的安全性靠几层：LocalSet 本身 !Send；spawn_local 只认线程本地的 CURRENT；其他线程的唤醒只进加锁的远程队列；只有 owner 线程上换出的 LocalNotified 能 run。
• LocalRuntime 在 1.54.1 是稳定 API，本质是记住了线程 ID 的 current_thread 运行时，可以直接 spawn_local。
• task_local! 是线程本地的 RefCell<Option<T>>，TaskLocalFuture 每次 poll 用 mem::swap 把值换进去、用 drop 守卫换出来。值不会传给 spawn 的子任务。