如何通过 contextvars 模块在异步任务间安全地传递上下文状态

在现代 Python 编程中，尤其是使用 asyncio 进行高并发开发时，管理状态和上下文是一个常见挑战。如果你习惯使用全局变量来存储请求相关的信息（例如，请求 ID、用户会话数据），在异步环境中会立即遇到问题：当事件循环在不同任务之间切换时，全局变量会被污染，导致一个任务读取到另一个任务设置的值。

Python 3.7 引入了 contextvars 模块，它提供了“上下文变量”（Context Variables），完美解决了这一问题。contextvars 允许数据在同一个异步任务（或线程）的执行栈中保持独立，从而实现了任务本地存储（Task-Local Storage）。

核心概念：ContextVar

ContextVar 类似于线程本地存储（Thread-Local Storage），但它是针对异步任务设计的。每个新的 asyncio 任务在创建时都会自动复制当前的上下文，确保其内部设置的上下文变量不会泄漏到其他任务中。

实操步骤与代码示例

我们将演示如何定义一个 ContextVar 来存储请求 ID，并确保两个并发运行的异步任务能够读取到各自正确的 ID。

步骤一：定义 ContextVar

首先，从 contextvars 导入 ContextVar 并定义一个变量，最好设置一个默认值。

import asyncio
import contextvars
import time

# 定义一个上下文变量，用于存储请求ID，并设置默认值
request_id_var = contextvars.ContextVar('request_id', default='N/A')

步骤二：创建异步工作函数

工作函数负责读取当前的上下文变量值。

async def worker_task(task_name):
    """一个模拟工作的异步任务，读取其上下文中的 Request ID"""
    # 1. 任务启动时读取 Request ID
    current_id = request_id_var.get()
    print(f"[{task_name} START] Current Request ID: {current_id}")

    # 模拟 I/O 操作，此时事件循环可能会切换到另一个任务
    await asyncio.sleep(0.05)

    # 2. 任务继续执行，再次读取 Request ID
    after_sleep_id = request_id_var.get()
    print(f"[{task_name} END] After sleep, Request ID: {after_sleep_id}")

步骤三：在不同上下文中运行任务

在主函数中，我们为每个任务设置不同的 request_id_var 值，并观察隔离效果。

async def main():
    print("--- 启动任务 A (Request ID: 1001) ---")
    # 设置上下文变量的值。这个设置仅影响当前任务及其子任务。
    token_a = request_id_var.set(1001)
    task_a = asyncio.create_task(worker_task("Task A"))

    # 注意：我们不需要在设置 Task A 之后立即重置上下文（使用 token_a），
    # 因为 contextvars 在创建新任务时会自动复制当前的上下文环境。

    print("\n--- 启动任务 B (Request ID: 2002) ---")
    # 再次设置上下文变量的值，这会覆盖主线程的上下文，但不会影响 Task A 的副本
    token_b = request_id_var.set(2002)
    task_b = asyncio.create_task(worker_task("Task B"))

    # 等待所有任务完成
    await asyncio.gather(task_a, task_b)

    # 重要的清理步骤：在任务完成后，重置主线程的上下文到原始状态
    request_id_var.reset(token_a)
    request_id_var.reset(token_b)

# 运行主函数
if __name__ == "__main__":
    # 使用 try-except 来确保在非脚本运行时也能正常运行
    try:
        asyncio.run(main())
    except RuntimeError as e:
        # 可能是 Jupyter/IPython 环境，跳过。
        pass

运行结果分析

运行上述代码，你会观察到以下输出：

--- 启动任务 A (Request ID: 1001) ---

--- 启动任务 B (Request ID: 2002) ---
[Task A START] Current Request ID: 1001
[Task B START] Current Request ID: 2002
[Task A END] After sleep, Request ID: 1001
[Task B END] After sleep, Request ID: 2002

关键点： 尽管 Task A 和 Task B 都在同一个事件循环中交错执行，并且主线程在启动 Task A 后又立即为 Task B 更改了上下文变量的值，但两个任务都保持了自己被创建时所继承的 Request ID。这证明了 contextvars 成功提供了任务级别的上下文隔离。