从头实现异步任务局部变量

# 从头实现异步任务局部变量 来源：https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/ 长话短说，假设我们有以下约束条件（先不管这些约束条件从何而来）： > - 我们想将某些数据与 Rust 中的异步任务（async task）关联起来，使得在轮询（poll）该任务时，可以在调用栈更深的地方检索到这些数据，而无需通过函数参数手动传递。 > - 我们不想使用 `tokio` 生态系统。也就是说，我们想做 `tokio::task_local!` (https://docs.rs/tokio/latest/tokio/macro.task_local.html) 所做的事，但不想用 tokio。怎么办？ 好吧……我找不到其他做这件事的方法，所以我只好自己做了 (:) ## 系好安全带，现在是做苹果派的时间 (https://www.goodreads.com/quotes/32952-if-you-wish-to-make-an-apple-pie-from-scratch) 要理解我们想要什么，首先必须理解为什么标准库提供的东西不够用。要理解标准库提供的东西，首先必须理解它所依赖的执行模型。我说过“从零开始”1 (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fn-other-reading) 嘛，毕竟 :P 大多数操作系统将代码执行组织成两个不同的单元：**进程（Processes）** 和 **线程（Threads）**。进程大致划分了“哪些内存可见”，而线程大致划分了“哪些指令正在运行”。一个进程可以包含一个或多个线程，当进程中有多个线程时，计算机科学家称之为“多线程架构 (https://en.wikipedia.org/wiki/Multithreading_(computer_architecture))”和“共享内存并行 (https://en.wikipedia.org/wiki/Parallel_programming_model#Shared_memory)”：字面上有多个线程，在同一个进程中共享内存。 这种进程/线程模型需要程序员做一些工作，以确保线程在运行时不会互相干扰。大多数普通代码假设自己是唯一能访问其所触碰内存的，如果其他线程趁它不注意在其工作上乱涂乱画，它会非常恼火。如果线程**不**触碰彼此的内存，或者有某种机制让它们协调彼此的触碰2 (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fn-gay)，这种性质就称为“线程安全（Thread-Safety）”。自然，线程安全是非常重要的性质（在编译时确保它是 Rust 的一大动机！），所以我们在实现总体目标时需要保持这个性质。 有多种策略可以确保线程安全，但最简单的是“直接在栈上使用变量啦”。每个线程的栈天然是唯一的3 (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fn-thread-stacks)，所以如果我们在栈上创建一个值，并通过参数向需要它的任何地方传递引用，boom，搞定，轻松加愉快。 但是记住，我们有一个特殊要求： > [……] 可以在调用栈更深的地方检索到这些数据，**而无需通过函数参数手动传递**。 所以我们真正想要的是某种“全局”变量，但又没有多个线程共享一个全局变量所带来的麻烦。幸运的是，这已经存在了——**线程局部存储（Thread-Local Storage）**，简称 TLS，创建的变量是“局部”的，意味着只能被单个线程访问，但又是“全局”的，因为任何函数都可以触碰它们。不同平台的 TLS 实现可能大不相同，我们只需要知道：每个线程有自己的位置来存放全局变量，再也没有脚趾相碰的问题了，大家皆大欢喜。C/C++ 使用 `thread_local` 关键字 (https://en.cppreference.com/w/c/thread/thread_local.html)，而 Rust 使用 `std::thread_local!` 宏 (https://doc.rust-lang.org/std/macro.thread_local.html)。 ## 线程局部变量就够用了吗？ 事实证明不够！在我们开始做苹果派之前，还需要引入更多概念，特别是操作系统线程和 Rust 任务之间的区别。 对于线程来说，操作系统会自动管理诸如“这个 CPU 核心上运行着哪个线程？”、“哪些线程可以运行但还没运行？”以及“好了，该切换到另一个线程了”之类的事情。然而，在操作系统层面切换线程相对于普通代码来说是非常慢的操作4 (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fn-os-slow)。因此，语言通常会在“用户空间”（即无需操作系统帮助）中创建自己的类线程结构，以避免这种减速，同时保持并发代码的其他好处。这些类线程结构的名字根据某些因素不同，要么叫绿色线程 (https://en.wikipedia.org/wiki/Green_thread) 要么叫协程 (https://en.wikipedia.org/wiki/Coroutine)。 ### 至于这些因素的旁注 主要因素是“它能否自行运行”（绿色线程，即抢占式多任务）还是“是否需要其他东西来驱动它”（协程，即协作式多任务）。绿色线程的常见例子是 Erlang (https://en.wikipedia.org/wiki/Erlang_(programming_language)) 的进程和 Go 的 goroutine (https://go.dev/tour/concurrency/1)（是的，他们的命名真的很混乱），而协程的例子包括 Python 的协程 (https://docs.python.org/3/library/asyncio-task.html)、C++ 的协程 (https://en.cppreference.com/w/cpp/language/coroutines.html) 以及 Rust 的 `Future` (https://doc.rust-lang.org/std/future/trait.Future.html) 和 `Iterator` (https://doc.rust-lang.org/std/iter/trait.Iterator.html) trait。Javascript 的 `Promise` (https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Promise) 对象很有趣，因为乍一看它们像是协程（有 `async`/`await` 语法），但实际上它们是绿色线程，因为它们会自行开始运行，而 `await` 语法只是将回调链在一起的 fancy 方式。 总之！！这些区别往往非常有争议，我肯定会收到一些有益的评论，指出我在这方面的错误 (: 暂时我将它们统称为“任务（Tasks）”，这样就能最大化惹恼别人 :P 总之！需要注意到的重要事情是：**任务不是线程**，所以我们不能只使用线程局部存储，因为： 1. 可能在同一个线程上运行多个任务，或者 2. 一个任务可能在其执行的不同时间点被分配到不同的线程上（出于吞吐量的考虑）。 具体来说，那种“任务开始运行时设置线程局部变量，任务结束时取消设置”的计划是行不通的，因为： 1. 其他在同一线程上开始运行的任务可能会共享同一个线程局部变量，或者 2. 任务可能在进行设置和尝试读取之间被转移到另一个线程。 这两种失败模式都会让我们伤心 :( 所以！讲完这些，我希望已经说服你：我们确实需要为任务局部变量构建一个特殊的结构。现在开始构建…… ## 船长，我们的策略是什么？ Rust 任务的工作方式是，**它们只会在其 `Future::poll()` (https://doc.rust-lang.org/std/future/trait.Future.html#tymethod.poll) 方法被调用时执行工作**5 (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fn-coroutine-green-thread)。这最终引导我们得出关键见解：为了拥有任务局部值，我们只需要在**调用 `poll()` 期间**设置某个全局值。使用线程局部值作为这个全局值是可行的，因为在单次 `poll()` 调用期间，**我们不能改变线程**。所以我们可以在那个单一线程上自由地保存/恢复： ### 关于即将出现的 `unsafe` 块的旁注 你应该忽略它们。它们是为了处理 `Pin` (https://doc.rust-lang.org/std/pin/index.html) 语义，在实际的代码库中，你应该使用 `pin_project` (https://docs.rs/pin-project/latest/pin_project/) crate 来处理，它提供了安全的方式。我展示不使用该 crate 的代码，只是为了证明背后没有魔法，只是标准库函数。 ```rust use std::cell::RefCell; use std::pin::Pin; use std::task::{Context, Poll}; use std::thread::LocalKey; /// 容器类型，通过包装现有 [`std::future::Future`] 实现它，并保存/恢复使其可用的任务局部变量。 pub struct Scoped<T, F: Future> { /// 保存当前任务局部值的线程局部存储。 tls: &'static LocalKey<RefCell<Option<T>>>, /// 当 future **没有**被轮询时：我们想存储的值在 `curr` 里面。 /// 当 future **正在**被轮询时：如果我们在另一个 [`Scoped`] 内部运行，则 `curr` 中保存之前的值。 curr: Option<T>, /// 我们正在包装的 future。 fut: F, } impl<T, F: Future> Future for Scoped<T, F> { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { // 安全性：`this` 永远不会移出其值，所以 `self` 保持固定。 let this = unsafe { &mut self.get_unchecked_mut() }; // 安全性：因为 `self` 仍然是固定的，所以 `fut` 也是固定的。 let fut = unsafe { Pin::new_unchecked(&mut this.fut) }; // 交换只移动 `self` 内部存储的值； // 它不会直接改变 `self` 的位置。 let curr = &mut this.curr; let mut swap = || this.tls.with_borrow_mut(|prev| std::mem::swap(curr, prev)); // 在轮询时，将值从“栈”（`this.curr`）交换到“全局内存”（`this.tls`）。 swap(); let out = fut.poll(cx); // 将值从“全局内存”换回“栈”，以便在稍后轮询时保存。 swap(); out } } ``` 我希望所有这些解释有助于理解我们为什么要这样做——这个实现并非凭空而来，以及为什么我们把保存/恢复逻辑放在 `poll()` 中而不是构造函数/`Drop` 实现中。 现在剩下的就是提供一些辅助函数，使这个包装器拥有更类似于 `tokio` 的 API。首先，我们定义一个辅助结构体，提供与 `tokio::task::LocalKey` (https://docs.rs/tokio/latest/tokio/task/struct.LocalKey.html#method.scope) 的 `scope()` 和 `with()` 函数非常相似的函数： ```rust pub struct ScopeBuilder<T> { tls: &'static LocalKey<RefCell<Option<T>>>, } impl<T> ScopeBuilder<T> { /// 为给定的线程局部存储创建一个新的任务包装器构建器 pub const fn new(tls: &'static LocalKey<RefCell<Option<T>>>) -> Self { Self { tls } } /// 给定构建器创建时提供的线程局部存储，构建一个 [`Scoped`] future，它将暴露给定的 `value`。 pub fn scope<F: Future>(&self, value: T, fut: F) -> Scoped<T, F> { Scoped { tls: self.tls, curr: Some(value), fut, } } /// 从线程局部存储中读取，如果我们在 [`Scoped`] future 内部，则返回 `Some`。 /// 接受一个回调函数以确保生命周期正确。 pub fn with<V>(&self, f: impl FnOnce(Option<&T>) -> V) -> V { self.tls.with_borrow(|value| f(value.as_ref())) } } ``` 然后，为了完整性，我们还定义一个类似于 `tokio::task_local!` (https://docs.rs/tokio/latest/tokio/macro.task_local.html) 的宏，它以符合人体工程学的方式同时创建线程局部存储和关联的 `ScopeBuilder`： ```rust #[macro_export] macro_rules! task_local { ($(static $ident:ident : $ty:ty ;)*) => { $( static $ident: $crate::ScopeBuilder<$ty> = { std::thread_local! { static LOCAL: std::cell::RefCell<Option<$ty>> = const { std::cell::RefCell::new(None) }; } $crate::ScopeBuilder::new(&LOCAL) }; )* } } ``` 组合在一起，用法看起来像这样： ```rust task_local! { static SOME_GLOBAL: usize; } /// 打印当前包含在 [`SOME_GLOBAL`] 中的值 fn print() { let v = SOME_GLOBAL.with(|v| v.map(|v| *v)); match v { Some(v) => println!("{v}"), None => println!("None"), } } #[tokio::main] async fn main() { // 定义一层层包装的任务，展示任务局部变量如何保存/恢复 print(); tokio::spawn(SOME_GLOBAL.scope(5138008, async { print(); tokio::spawn(SOME_GLOBAL.scope(69, async { print(); tokio::spawn(SOME_GLOBAL.scope(42, async { print(); })).await.unwrap(); print(); })).await.unwrap(); print(); })).await.unwrap(); print(); } ``` 注意每个 `.scope()` 调用外面包裹了 `tokio::spawn()`（以正确测试多个任务并发运行时的情况），但每个 `async { }` 调用外面没有（以确保 `print()` 调用与对应的 `.scope()` 在同一个任务中发生）。在 Rust 游乐场 (https://play.rust-lang.org/?version=stable&mode=debug&edition=2024&gist=02c3f51b299f07ac860e28dd16933c1d) 尝试一下，我们观察到以下输出： ``` None 5138008 69 42 69 5138008 None ``` 完全符合我们的设计 :) ## 结论 在这篇文章中，我展示了一种完整的（！！）在 Rust 中编写任务局部变量的方法，仅使用标准库，覆盖了“单个线程上的多个任务”的情况（通过从栈交换到线程局部存储），以及“单个任务在多个线程上”的情况；`std::thread::LocalKey` (https://doc.rust-lang.org/std/thread/struct.LocalKey.html) 代表“如何访问当前线程的线程局部存储”，因此生成的 `Scoped` future 实际上可以发送到其他线程，它将简单地使用它开始运行所在线程的 TLS。 这篇文章中的源代码直接来源于我正在制作的更大 crate 集合 (https://github.com/p0lyw0lf/driver/blob/725e822d8302ecf1af290aab9582c357c56f9e2a/packages/async-task-local/src/lib.rs) 的一部分，虽然我可能会考虑最终将其发布到 crates.io，但它足够小，如果你真的想要的话，可以直接 vendored 到任何项目中。 以上就是全部内容了，希望你喜欢并且可能还学到了一些东西！下次再见！！ 1. 我考虑过链接到 Kora 的《为 3DS 构建 AsyncIO 执行器》(https://blog.cat-girl.gay/3ds-async-part-one/) 来替代写整个引言部分，因为那篇文章涵盖了非常相似的内容，但是嘿，世界上再多一份解释也没什么坏处吧？↩ (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fnref-other-reading) 2. 例如：互斥锁、条件变量、信号量等。我不为我尽可能猥琐地措辞的习惯道歉，因为你仔细想想，**真正**想想，你也会意识到并发是 yur（就是那种感觉）。具体风味因语言而异：Rust => 慢炖，Go => 有毒，C++ => 注定完蛋。我陈述完毕。↩ (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fnref-gay) 3. 栈和“正在执行的代码”这两个概念紧密相连，至少在我们当前的语言/计算机架构下不可能有其他方式。↩ (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fnref-thread-stacks) 4. 好吧，我不打算深入探讨这个。原因**真的**很复杂，我承认我自己也没有完全理解。大概是什么计时器中断啦，流水线停顿啦，`%cr3` (https://en.wikipedia.org/wiki/Control_register#CR3) 啦？？只要记住“TL;DR 你应该知道这件事”就够了。↩ (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fnref-os-slow) 5. 注意这个定义实际上并不取决于任务是协程还是绿色线程！是的，Rust 的 Future 实际上是协程，而 `poll()` 方法往往表明了这一点，但基于绿色线程的系统也可以暴露具有类似保证的 API，很可能就是为了构造这种作用域包装器。↩ (https://wolfgirl.dev/blog/2026-06-16-async-task-locals-from-scratch/#fnref-coroutine-green-thread)