Python使用contextvars模块传递上下文的底层原理

前言

contextvars:是Python提供的用于存放上下文信息的模块，支持asyncio，可以将上下文信息无感地在不同的协程方法中传递。contextvars模块主要有两个类：ContextVar和Context，Context可以是一个map，map的键是ContextVar。不同方法中的上下文传递实际上是通过拷贝Context来实现的。本文主要介绍contextvars模块的基本用法、底层实现、写时拷贝以及浅拷贝需要注意的事项。

contextvars基本用法

如下方代码示例和运行结果所示，contextvars有以下特点：

• contextvar无需显式地在函数参数中透传
• 子方法中改变contextvar的值，不对父方法中该contextvar的值造成影响
• contextvar需要写成全局变量，参考

代码示例：

import asyncio
import contextvars


ctx = contextvars.ContextVar('trace')
ctx.set("begin")


async def fun():
    ctx.set(ctx.get() + "|fun")
    print("ctx:", ctx.get())


async def main():
    ctx.set(ctx.get()+"|main")
    print("befor call fun: ctx", ctx.get())
    await fun()
    print("after call fun: ctx", ctx.get())

print("befor call main: ctx", ctx.get())
asyncio.get_event_loop().run_until_complete(main())
print("after call main: ctx", ctx.get())

运行结果：

befor call main: ctx begin
befor call fun: ctx begin|main
ctx: begin|main|fun
after call fun: ctx begin|main|fun
after call main: ctx begin

asyncio支持context传递的底层实现

前文提到，ContextVar实际上只是Context中的一个key，如果有多个ContextVar，Context中将会有多个这样的key，而ContextVar在父子方法中传递则是通过拷贝Context来实现的。
asyncio 通过Loop.call_soon()、Loop.call_later()、 和 Loop.call_at() 来调度协程。 asyncio.Task 用 call_soon() 运行一个包装过的协程.
asyncio为了支持contextvars，对Loop.call_{at,later,soon} 和 Future.add_done_callback() 做了修改，支持传入一个Context，这个Context默认为当前的Context：

def call_soon(self, callback, *args, context=None):
    if context is None:
        context = contextvars.copy_context()

    # ... some time later
    context.run(callback, *args)

asyncio中的Task对象也会维护一份它创建时的Context:

class Task:
    def __init__(self, coro):
        ...
        # Get the current context snapshot.
        self._context = contextvars.copy_context()
        self._loop.call_soon(self._step, context=self._context)

    def _step(self, exc=None):
        ...
        # Every advance of the wrapped coroutine is done in
        # the task's context.
        self._loop.call_soon(self._step, context=self._context)
        ...

写时拷贝

读到这里你可能会想，如果每个协程都拷贝了一份Context，会不会造成内存资源浪费？
实际上Context的拷贝和Linux中的进程fork运用了同一种技术：写时拷贝。即只有在子方法对Context进行写操作时，才会执行拷贝，这样那些不需要修改Context的协程其实只拥有了一个指向父协程Context的一个指针而已，并不会造成资源浪费。

浅拷贝的坑

前面说过，子方法会拷贝一份（如果不修改也可能不拷贝）父方法的Context，而这里的拷贝实际上是浅拷贝，也就是说当ContextVar是一个复杂对象时，子方法对ContextVar的值进行修改会对父方法产生影响（无论是执行了写时拷贝还是没执行），因为他们指向的实际上时同一个对象。
考虑下面代码：

import asyncio
import contextvars


ctx = contextvars.ContextVar('dict')
ctx2 = contextvars.ContextVar("number")
ctx2.set(0)
ctx.set({})


async def fun():
    ctx2.set(1)  # 此时会进行写时拷贝
    test_dict = ctx.get()
    test_dict["fun"] = "xxx"
    print("fun: ctx", ctx.get(), ctx2.get())


def main():
    print("befor call fun: ctx", ctx.get(), ctx2.get())
    context = contextvars.copy_context()
    asyncio.get_event_loop().run_until_complete(fun())
    print("after call fun: ctx", ctx.get(), ctx2.get())


main()

运行结果:

befor call fun: ctx {} 0
fun: ctx {'fun': 'xxx'} 1
after call fun: ctx {'fun': 'xxx'} 0

虽然fun()中拷贝了一份Context，写ctx1对main中的ctx1并未产生影响，但是fun()中拷贝的ctx的value是dict类型，只是进行了浅拷贝，所以对它修改会影响main中的ctx1。
实际运用时一定要小心浅拷贝的坑。

普通函数中传递context

前文介绍的都是在协程间传递上下文，由于asyncio对部分方法做了修改，所以可以无感知地传递Context。但是如果想要在普通函数间传递，则需要用户层面手动地调用copy_context()拷贝一份Context，然后再调用context.run()在Context中运行子方法，这样可确保父子方法的Context互不影响，见下方代码：

import contextvars


ctx = contextvars.ContextVar("number")
ctx.set(0)


def fun():
    ctx.set(1)
    print("fun: ctx", ctx.get())


def main():
    print("befor call fun: ctx", ctx.get())
    context = contextvars.copy_context()
    context.run(fun)
    print("after call fun: ctx", ctx.get())


main()

运行结果:

befor call fun: ctx 0
fun: ctx 1
after call fun: ctx 0

参考

• contextvars
• PEP 567