Rust Async/Await能否像传统线程模型一样支持线程本地变量？

问题背景 #

Rust 的 async/await 模型通过线程复用提升并发效率，但异步任务在 .await 后可能被调度到不同线程恢复执行，导致依赖线程本地变量（如日志 MDC 上下文）的场景出现异常——线程切换后，原线程的本地变量无法被新线程继承，就像测试中出现的 MDC 值错乱问题：.await 前后线程 ID 变化，预期的 MDC 值与实际获取值不匹配。

测试依赖：

[dependencies]
log = { version = "0.4" }
log-mdc = { version = "0.1" }
tokio = { version = "1.3", features = ["full"] }
futures = { version = "0.3" }

可行解决方案 #

1. 用 tokio::task::LocalSet 绑定任务到单线程 #

如果业务场景允许任务始终在同一线程执行，可以使用 LocalSet 将异步任务限制在单个线程上，确保线程本地变量全程有效。

修改后的示例代码：

use tokio::task::LocalSet;

#[tokio::main]
async fn main() {
    fs_util::log::init_default().unwrap();

    let local = LocalSet::new();
    local.run_until(async {
        let mut join_handles = Vec::new();
        for _i in 1..10_000 {
            // 使用 spawn_local 替代 spawn，任务将在 LocalSet 绑定的线程执行
            join_handles.push(tokio::task::spawn_local(do_stuff()));
        }
        futures::future::join_all(join_handles).await;
        log::info!("Program DONE !!");
    }).await;
}

注意：这种方式会牺牲部分线程复用的优势，仅适合对线程绑定有强需求的场景。

2. 使用 async-local crate 实现异步任务上下文 #

async-local crate 提供了基于异步任务的上下文存储，替代线程本地变量，上下文会跟随异步任务的调度流程，而非绑定到特定线程。

首先添加依赖：

async-local = "0.1"

示例代码：

use async_local::AsyncLocal;
use log;
use std::thread;
use tokio;

// 定义全局异步上下文
static MDC_CONTEXT: AsyncLocal<String> = AsyncLocal::new();

#[tokio::main]
async fn main() {
    fs_util::log::init_default().unwrap();

    let mut join_handles = Vec::new();
    for _i in 1..10_000 {
        join_handles.push(tokio::spawn(do_stuff()));
    }
    futures::future::join_all(join_handles).await;
    log::info!("Program DONE !!");
}

async fn do_stuff() {
    let thread_id = thread::current().id();
    let mdc_val = format!("mdcval={thread_id:?}");

    // 将上下文绑定到当前异步任务
    MDC_CONTEXT.set(mdc_val.clone(), async {
        compute_the_answer().await;

        // 从异步上下文中获取值，不受线程切换影响
        let retrieved_val = MDC_CONTEXT.get().unwrap();
        if retrieved_val == mdc_val {
            log::info!("Context preserved across await");
        } else {
            log::error!("Context lost!");
        }
    }).await;
}

async fn compute_the_answer() -> i32 {
    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
    42
}

3. 自定义 Future 包装器自动恢复上下文 #

编写自定义 Future 包装器，在任务被调度执行时自动将上下文注入当前线程的本地存储，实现透明的上下文传递，无需修改业务逻辑。

示例包装器代码：

use futures::Future;
use log_mdc;
use std::pin::Pin;
use std::task::{Context, Poll};

// 包装器：持有原 Future 和需要传递的 MDC 上下文
struct MdcdFuture<F> {
    inner: F,
    mdc: Vec<(String, String)>,
}

impl<F: Future> Future for MdcdFuture<F> {
    type Output = F::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // 将 MDC 上下文设置到当前线程
        for (key, val) in &self.mdc {
            log_mdc::insert(key, val);
        }

        // 执行原 Future
        let result = Pin::new(&mut self.inner).poll(cx);

        // 清理当前线程的 MDC，避免污染其他任务
        for (key, _) in &self.mdc {
            log_mdc::remove(key);
        }

        result
    }
}

// 便捷包装方法
fn with_mdc<F: Future>(future: F) -> MdcdFuture<F> {
    // 捕获当前线程的 MDC 上下文
    let mdc = log_mdc::iter().map(|(k, v)| (k.to_string(), v.to_string())).collect();
    MdcdFuture { inner: future, mdc }
}

// 使用方式：调度任务前包装
async fn main() {
    // ... 初始化代码 ...
    join_handles.push(tokio::spawn(with_mdc(do_stuff())));
    // ...
}

内容的提问来源于stack exchange，提问作者Andrei Matei