Python contextvars 异步上下文：从线程局部到协程隔离

1
import threading
2
import time
3

4
# 创建线程局部存储
5
local_data = threading.local()
6

7
def worker(thread_id: str):
8
    # 每个线程只能看到自己存储的值
9
    local_data.request_id = thread_id
10
    time.sleep(0.1)  # 模拟处理
11
    print(f"线程 {threading.current_thread().name}: "
12
          f"request_id = {local_data.request_id}")
13

14
threads = []
15
for i in range(3):
16
    t = threading.Thread(target=worker, args=(f"req-{i}",))
17
    threads.append(t)
18
    t.start()
19
for t in threads:
20
    t.join()
21
# 输出：
22
# Thread-1: request_id = req-0
23
# Thread-2: request_id = req-1
24
# Thread-3: request_id = req-2

1
import asyncio
2

3
local_data = threading.local()
4

5
async def handle_request(request_id: str):
6
    local_data.request_id = request_id
7
    # 模拟异步操作（主动让出控制权）
8
    await asyncio.sleep(0.1)
9
    # 醒来后，local_data.request_id 可能已经不是你的了！
10
    print(f"请求 {request_id}: 读取 request_id = {local_data.request_id}")
11

12
async def main():
13
    # 三个协程在同一个线程中交替执行
14
    await asyncio.gather(
15
        handle_request("req-A"),
16
        handle_request("req-B"),
17
        handle_request("req-C"),
18
    )
19

20
asyncio.run(main())
21
# 可能的输出（全乱了）：
22
# 请求 req-A: 读取 request_id = req-C
23
# 请求 req-B: 读取 request_id = req-C
24
# 请求 req-C: 读取 request_id = req-C

1
from contextvars import ContextVar
2

3
# 声明一个上下文变量，default 是默认值
4
request_id: ContextVar[str] = ContextVar("request_id", default="no-request")
5

6
# 读取当前值
7
print(request_id.get())  # no-request
8

9
# 设置值（会返回一个 Token，用于后续恢复）
10
token = request_id.set("abc-123")
11
print(request_id.get())  # abc-123
12

13
# 恢复到设置之前的状态
14
request_id.reset(token)
15
print(request_id.get())  # no-request

1
import asyncio
2
from contextvars import ContextVar
3

4
request_id = ContextVar("request_id", default="no-request")
5

6
async def handle_request(request_id_value: str):
7
    # 在任务内部设置上下文变量
8
    request_id.set(request_id_value)
9
    print(f"[开始] 任务 {request_id_value}: "
10
          f"当前 request_id = {request_id.get()}")
11

12
    await asyncio.sleep(0.1)  # 让出控制权
13

14
    # 醒来后，读取到的仍然是自己的值！
15
    print(f"[结束] 任务 {request_id_value}: "
16
          f"当前 request_id = {request_id.get()}")
17

18
async def main():
19
    await asyncio.gather(
20
        handle_request("req-A"),
21
        handle_request("req-B"),
22
        handle_request("req-C"),
23
    )
24

25
asyncio.run(main())
26
# 输出（完全隔离）：
27
# [开始] 任务 req-A: 当前 request_id = req-A
28
# [开始] 任务 req-B: 当前 request_id = req-B
29
# [开始] 任务 req-C: 当前 request_id = req-C
30
# [结束] 任务 req-A: 当前 request_id = req-A
31
# [结束] 任务 req-B: 当前 request_id = req-B
32
# [结束] 任务 req-C: 当前 request_id = req-C

1
import asyncio
2
from contextvars import ContextVar
3

4
trace_id = ContextVar("trace_id", default="root")
5

6
async def child_task():
7
    await asyncio.sleep(0.05)
8
    print(f"子任务读取 trace_id = {trace_id.get()}")
9

10
async def main():
11
    trace_id.set("parent-trace")
12
    print(f"父任务设置 trace_id = {trace_id.get()}")
13

14
    # create_task 时，会在此刻复制上下文快照
15
    task = asyncio.create_task(child_task())
16

17
    # 之后父任务的修改不影响子任务
18
    trace_id.set("new-parent-trace")
19

20
    await task
21
    print(f"父任务最终 trace_id = {trace_id.get()}")
22

23
asyncio.run(main())
24
# 输出：
25
# 父任务设置 trace_id = parent-trace
26
# 子任务读取 trace_id = parent-trace
27
# 父任务最终 trace_id = new-parent-trace

1
import asyncio
2
from contextvars import ContextVar
3

4
# ⚠️ 危险：存储可变对象
5
request_context = ContextVar("request_context", default={})
6

7
async def dangerous_task(task_id: str):
8
    ctx = request_context.get()
9
    ctx["task"] = task_id  # 修改共享的 dict！
10
    print(f"任务 {task_id}: {request_context.get()}")
11
    await asyncio.sleep(0.1)
12
    print(f"任务 {task_id} (醒来后): {request_context.get()}")
13

14
async def main():
15
    await asyncio.gather(
16
        dangerous_task("A"),
17
        dangerous_task("B"),
18
    )
19

20
asyncio.run(main())
21
# 可能的输出（污染了）：
22
# 任务 A: {'task': 'A'}
23
# 任务 B: {'task': 'B'}
24
# 任务 A (醒来后): {'task': 'B'}   ← 被任务 B 修改了！
25
# 任务 B (醒来后): {'task': 'B'}

1
import asyncio
2
from contextvars import ContextVar
3

4
request_context = ContextVar("request_context", default=None)
5

6
async def safe_task(task_id: str):
7
    # 每次 set 时创建全新的 dict
8
    request_context.set({"task": task_id, "timestamp": asyncio.get_event_loop().time()})
9
    ctx = request_context.get()
10
    print(f"任务 {task_id}: {ctx}")
11
    await asyncio.sleep(0.1)
12
    ctx = request_context.get()
13
    print(f"任务 {task_id} (醒来后): {ctx}")
14

15
async def main():
16
    await asyncio.gather(
17
        safe_task("A"),
18
        safe_task("B"),
19
    )
20

21
asyncio.run(main())
22
# 输出（完全隔离）：
23
# 任务 A: {'task': 'A', 'timestamp': 123456.1}
24
# 任务 B: {'task': 'B', 'timestamp': 123456.1}
25
# 任务 A (醒来后): {'task': 'A', 'timestamp': 123456.1}
26
# 任务 B (醒来后): {'task': 'B', 'timestamp': 123456.1}

1
import asyncio
2
import uuid
3
import logging
4
from contextvars import ContextVar
5
from functools import wraps
6

7
# 全局上下文变量
8
trace_id_ctx: ContextVar[str] = ContextVar("trace_id", default="")
9
user_id_ctx: ContextVar[str] = ContextVar("user_id", default="anonymous")
10

11
# 自定义日志格式，自动注入 trace_id
12
class TraceFormatter(logging.Formatter):
13
    def format(self, record):
14
        record.trace_id = trace_id_ctx.get()
15
        record.user_id = user_id_ctx.get()
16
        return super().format(record)
17

18
# 配置日志
19
logger = logging.getLogger("api")
20
handler = logging.StreamHandler()
21
handler.setFormatter(TraceFormatter(
22
    "[%(trace_id)s] [%(user_id)s] %(levelname)s: %(message)s"
23
))
24
logger.addHandler(handler)
25
logger.setLevel(logging.INFO)
26

27
def with_trace(func):
28
    """装饰器：为异步函数自动创建 trace_id"""
29
    @wraps(func)
30
    async def wrapper(*args, **kwargs):
31
        token_trace = trace_id_ctx.set(str(uuid.uuid4())[:8])
32
        try:
33
            result = await func(*args, **kwargs)
34
            return result
35
        finally:
36
            trace_id_ctx.reset(token_trace)
37
    return wrapper
38

39
@with_trace
40
async def handle_user_request(user_id: str):
41
    user_id_ctx.set(user_id)
42
    logger.info(f"收到用户请求")
43

44
    await validate_user(user_id)
45
    await fetch_user_data(user_id)
46
    await send_notification(user_id)
47

48
    logger.info(f"请求处理完成")
49
    return {"status": "ok"}
50

51
async def validate_user(user_id: str):
52
    logger.info(f"验证用户 {user_id}")
53
    await asyncio.sleep(0.05)
54

55
async def fetch_user_data(user_id: str):
56
    logger.info(f"查询用户数据 {user_id}")
57
    await asyncio.sleep(0.05)
58

59
async def send_notification(user_id: str):
60
    logger.info(f"发送通知给 {user_id}")
61
    await asyncio.sleep(0.05)
62

63
async def main():
64
    # 模拟三个并发请求
65
    await asyncio.gather(
66
        handle_user_request("alice"),
67
        handle_user_request("bob"),
68
        handle_user_request("charlie"),
69
    )
70

71
asyncio.run(main())
72
# 输出（每条日志都带有正确的 trace_id 和 user_id）：
73
# [a1b2c3d4] [alice] INFO: 收到用户请求
74
# [e5f6g7h8] [bob] INFO: 收到用户请求
75
# [i9j0k1l2] [charlie] INFO: 收到用户请求
76
# [a1b2c3d4] [alice] INFO: 验证用户 alice
77
# [e5f6g7h8] [bob] INFO: 验证用户 bob
78
# [i9j0k1l2] [charlie] INFO: 验证用户 charlie
79
# ...

1
from fastapi import FastAPI, Request
2
from contextvars import ContextVar
3
import uuid
4

5
trace_id_ctx: ContextVar[str] = ContextVar("trace_id", default="")
6

7
app = FastAPI()
8

9
@app.middleware("http")
10
async def trace_middleware(request: Request, call_next):
11
    # 为每个 HTTP 请求生成 trace_id
12
    token = trace_id_ctx.set(str(uuid.uuid4())[:8])
13
    try:
14
        response = await call_next(request)
15
        response.headers["X-Trace-Id"] = trace_id_ctx.get()
16
        return response
17
    finally:
18
        trace_id_ctx.reset(token)
19

20
@app.get("/users/{user_id}")
21
async def get_user(user_id: str):
22
    # 直接读取，无需从 request 对象传递
23
    logger.info(f"查询用户: {user_id}")
24
    # ... 业务逻辑
25
    return {"user_id": user_id, "trace_id": trace_id_ctx.get()}

1
import asyncio
2
from contextvars import ContextVar
3
from dataclasses import dataclass
4

5
@dataclass
6
class DBConnection:
7
    """模拟数据库连接"""
8
    conn_id: str
9
    in_transaction: bool = False
10
    transaction_level: int = 0
11

12
# 每个协程独立的数据库连接上下文
13
db_conn_ctx: ContextVar[DBConnection | None] = ContextVar("db_conn", default=None)
14

15
class DBPool:
16
    """模拟连接池"""
17
    def __init__(self):
18
        self._pool_counter = 0
19

20
    async def acquire(self) -> DBConnection:
21
        self._pool_counter += 1
22
        conn = DBConnection(conn_id=f"conn-{self._pool_counter}")
23
        await asyncio.sleep(0.01)  # 模拟获取连接的 IO
24
        return conn
25

26
    async def release(self, conn: DBConnection):
27
        await asyncio.sleep(0.01)
28
        print(f"  释放连接 {conn.conn_id}")
29

30
pool = DBPool()
31

32
async def with_db(func):
33
    """装饰器：自动管理数据库连接生命周期"""
34
    @wraps(func)
35
    async def wrapper(*args, **kwargs):
36
        conn = await pool.acquire()
37
        token = db_conn_ctx.set(conn)
38
        try:
39
            return await func(*args, **kwargs)
40
        finally:
41
            db_conn_ctx.reset(token)
42
            await pool.release(conn)
43
    return wrapper
44

45
async def begin_transaction():
46
    """开启事务"""
47
    conn = db_conn_ctx.get()
48
    if conn is None:
49
        raise RuntimeError("没有可用的数据库连接")
50
    conn.in_transaction = True
51
    conn.transaction_level += 1
52
    print(f"  事务开启 [{conn.conn_id}] (嵌套层级: {conn.transaction_level})")
53

54
async def commit():
55
    """提交事务"""
56
    conn = db_conn_ctx.get()
57
    conn.in_transaction = False
58
    print(f"  事务提交 [{conn.conn_id}]")
59

60
async def execute_query(sql: str):
61
    """执行查询"""
62
    conn = db_conn_ctx.get()
63
    if conn is None:
64
        raise RuntimeError("没有可用的数据库连接")
65
    tx_status = " [事务中]" if conn.in_transaction else ""
66
    print(f"  执行: {sql} [{conn.conn_id}]{tx_status}")
67
    await asyncio.sleep(0.02)
68

69
@with_db
70
async def create_order(user_id: str, item: str):
71
    await begin_transaction()
72
    await execute_query(f"INSERT INTO orders (user_id, item) VALUES ('{user_id}', '{item}')")
73
    await execute_query(f"UPDATE inventory SET stock = stock - 1 WHERE item = '{item}'")
74
    await commit()
75
    print(f"  ✓ 订单创建完成: {user_id} -> {item}")
76

77
async def main():
78
    await asyncio.gather(
79
        create_order("alice", "机械键盘"),
80
        create_order("bob", "显示器"),
81
        create_order("charlie", "鼠标"),
82
    )
83

84
asyncio.run(main())
85
# 输出（每个请求使用独立的连接）：
86
#   事务开启 [conn-1] (嵌套层级: 1)
87
#   事务开启 [conn-2] (嵌套层级: 1)
88
#   事务开启 [conn-3] (嵌套层级: 1)
89
#   执行: INSERT INTO orders (user_id, item) VALUES ('alice', '机械键盘') [conn-1] [事务中]
90
#   执行: INSERT INTO orders (user_id, item) VALUES ('bob', '显示器') [conn-2] [事务中]
91
#   执行: INSERT INTO orders (user_id, item) VALUES ('charlie', '鼠标') [conn-3] [事务中]
92
#   ...
93
#   事务提交 [conn-1]
94
#   ✓ 订单创建完成: alice -> 机械键盘
95
#   释放连接 conn-1
96
#   ...

1
import asyncio
2
from contextvars import ContextVar
3
from dataclasses import dataclass, field
4

5
@dataclass
6
class TenantConfig:
7
    """租户配置"""
8
    tenant_id: str
9
    db_schema: str
10
    feature_flags: dict = field(default_factory=dict)
11
    rate_limit: int = 1000
12

13
# 租户上下文
14
tenant_ctx: ContextVar[TenantConfig] = ContextVar("tenant", default=TenantConfig(
15
    tenant_id="default",
16
    db_schema="public",
17
    feature_flags={},
18
    rate_limit=1000,
19
))
20

21
# 租户配置数据库（模拟）
22
TENANT_REGISTRY = {
23
    "acme": TenantConfig("acme", "acme_schema", {"dark_mode": True}, 500),
24
    "globex": TenantConfig("globex", "globex_schema", {"dark_mode": False}, 2000),
25
    "initech": TenantConfig("initech", "initech_schema", {"dark_mode": True}, 100),
26
}
27

28
async def handle_tenant_request(tenant_id: str, action: str):
29
    config = TENANT_REGISTRY.get(tenant_id)
30
    if config is None:
31
        config = tenant_ctx.get()  # fallback 到默认租户
32

33
    token = tenant_ctx.set(config)
34
    try:
35
        current = tenant_ctx.get()
36
        print(f"[{current.tenant_id}] schema={current.db_schema}, "
37
              f"dark_mode={current.feature_flags.get('dark_mode')}, "
38
              f"rate_limit={current.rate_limit}")
39

40
        # 模拟处理
41
        await asyncio.sleep(0.05)
42
        print(f"[{current.tenant_id}] {action} 处理完成")
43
    finally:
44
        tenant_ctx.reset(token)
45

46
async def main():
47
    await asyncio.gather(
48
        handle_tenant_request("acme", "导出数据"),
49
        handle_tenant_request("globex", "导入数据"),
50
        handle_tenant_request("initech", "生成报表"),
51
    )
52

53
asyncio.run(main())
54
# 输出（每个租户独立配置）：
55
# [acme] schema=acme_schema, dark_mode=True, rate_limit=500
56
# [globex] schema=globex_schema, dark_mode=False, rate_limit=2000
57
# [initech] schema=initech_schema, dark_mode=True, rate_limit=100
58
# [acme] 导出数据 处理完成
59
# [globex] 导入数据 处理完成
60
# [initech] 生成报表 处理完成

1
import asyncio
2
from concurrent.futures import ThreadPoolExecutor
3
from contextvars import ContextVar
4

5
request_id = ContextVar("request_id", default="")
6

7
def sync_worker():
8
    # ⚠️ 在线程中读取不到异步上下文的值
9
    # request_id.get() 返回默认值 ""
10
    return f"线程内 request_id = {request_id.get()}"
11

12
async def async_handler():
13
    request_id.set("async-trace-123")
14

15
    loop = asyncio.get_event_loop()
16
    with ThreadPoolExecutor() as executor:
17
        result = await loop.run_in_executor(executor, sync_worker)
18

19
    print(result)  # 线程内 request_id =  (空！)
20
    print(f"异步上下文: {request_id.get()}")  # async-trace-123
21

22
asyncio.run(async_handler())

1
# ❌ 危险：default 是可变对象时，所有任务共享同一个实例
2
from contextvars import ContextVar
3

4
bad_default = ContextVar("bad", default={"key": "value"})
5

6
async def task_a():
7
    d = bad_default.get()
8
    d["task"] = "a"
9
    print(d)
10

11
async def task_b():
12
    d = bad_default.get()
13
    d["task"] = "b"
14
    print(d)
15

16
# 两个任务会修改同一个 dict 对象