一系列Python asyncio tasks调用中的contextvars

socrates
501 阅读4分钟

在Python中，contextvars某种程度上类似于线程的thread local，但主要用于协程。但是在一系列asyncio tasks调用中，上下文本身是否会保持一致呢？让我们来探究一下。

contextvars

让我们简要回顾一下Python中的contextvars。contextvars可以在一系列函数调用中维护上下文，并且原生支持asyncio。例如

import asyncio
import contextvars

var = contextvars.ContextVar('var', default={})

async def sub2():
    print(f'in sub2, {var.get()=}')
    var.set('sub1 set')

async def sub1():
    print(f'in sub1, {var.get()=}')
    var.set('sub1 set')
    await sub2()

async def main():
    var.set('main set')
    await sub1()
    print(f'in main, {var.get()=}')

asyncio.run(main())

上面程序的输出是

in sub1, var.get()='main set'
in sub2, var.get()='sub1 set'
in main, var.get()='sub1 set'

一个上下文可能包含一组contextvars.ContextVar对象，而contextvars.ContextVar对象的值可以在上下文中set和get。

contextvars有许多使用场景。例如：

  • 在Web server中，将每个请求的Request/Response对象作为contextvars，这样在请求处理时的任何函数中都可以访问Request/Response对象。
  • 在任务队列中，将每个任务的元数据（如任务ID、任务状态等）作为contextvars，这样在任务处理调用的任何函数都可以访问任务元数据。
  • 在日志系统中，将跟踪信息（如用户ID、请求ID等）作为contextvars，这样可以在日志中识别同一 trace/span 中的内容。

contextvars的一致性

contextvars可以避免在函数调用的系列中把相同的参数一直传来传去。但是在asyncio Task的调用系列中，上下文本身是否会保持一致呢？

这实际上取决于asyncio Tasks之间是如何相互调用的。让我们看一个例子：

import asyncio
import contextvars

var = contextvars.ContextVar('var', default={})

async def coro_func(level):
    var.set(level)
    print(f'{level=}, {var.get()=}')
    if level:
        await asyncio.create_task(coro_func(level-1))
    print(f'{level=}, {var.get()=}')

asyncio.run(coro_func(2))

这个例子用asyncio Tasks封装了直接async/await调用。一般来说，asyncio Task可以给协程一些附加的能力，比如取消任务等。但是当我们用Tasks来相互调用，上下文在Tasks之间并不保持一致：

level=2, var.get()=2
level=1, var.get()=1
level=0, var.get()=0
level=0, var.get()=0
level=1, var.get()=1
level=2, var.get()=2

当执行从内部Task返回时，内部Task的上下文也结束了，外部Task的上下文被恢复。上下文在Tasks之间没有保持一致。原因是当创建Task时，Task的上下文是从当前上下文复制的新上下文对象。对新上下文对象的任何写操作都不会影响外部上下文对象。因此，如果我们使用一些列Tasks，默认的任务创建行为将阻断上下文在Tasks之间保持的一致性。

在Python 3.11之后，create_task引入了context参数来将特定上下文传递给新创建的Task。我们可以这样的修改上面的例子：

import asyncio
import contextvars

var = contextvars.ContextVar('var', default={})

async def coro_func(level):
    var.set(level)
    print(f'{level=}, {var.get()=}')
    if level:
        context = asyncio.current_task().get_context() if asyncio.current_task() else None
        await asyncio.create_task(coro_func(level-1), context=context)
    print(f'{level=}, {var.get()=}')

asyncio.run(coro_func(2))

上面的输出将变为

level=2, var.get()=2
level=1, var.get()=1
level=0, var.get()=0
level=0, var.get()=0
level=1, var.get()=0
level=2, var.get()=0

这里Tasks之间的上下文对象是相同的。内部Task的修改反映在外部Task中。

使用context参数传递上下文对象的机制也适用于：

  • asyncio.call_soon, asyncio.call_laster and asyncio.call_at
  • asyncio.Handle
  • asyncio.Runner
  • TaskGroup.create_task

这种机制提供了一种方法，让用户决定是否在Tasks之间保持上下文一致。但是这个细节很容易被忽视，因为Task的创建通常隐藏在框架或库代码中。

starlette中间件和anyio

著名的Python Web框架fastapi建立在ASGI框架starlette之上，fastapi简单地使用了starlette的中间件机制。starlette中间件的每一层都作为由anyio处理的函数调用链执行。starlette使用anyio.create_task_group()来运行一层中间件。当使用asyncio时，anyio简单地将anyio.TaskGroup的执行包装在asyncio.create_task中（在4.4.0版本中没有传入上下文）。

因此，starlette中间件的上下文在中间件的各层之间不保持一致。如果你想在中间层的中间件中更改contextvars.ContextVar，这个更改不会反映在外层的中间件中。在starlette中间件中使用contextvars时，你应该注意这一点。

关于协程 vs 线程的一点思考

当我们说"contextvars某种程度上类似于线程的thread local，但主要用于协程"时，我们认为协程只是另一种类型的线程。但实际上这是不对的，协程不是线程。线程有自己的生命周期和自己的栈空间，所以线程局部存储实际上属于一个线程。但协程只是一个可以暂停和恢复的函数，由循环调度。因此，协程的上下文不属于协程本身。在协程链中，比如Tasks的调用链，我们需要注意上下文的一致性是由用户手动维护的，而不是由系统自动维护的。

avatar
文章
阅读
粉丝