Python 并发与 asyncio 14.2 上下文变量

14.2 上下文变量

想象一下，我们正在使用一个基于每请求一个线程的 Web 服务器。当有请求到达时，我们可能需要记录关于发起请求的用户的一些信息，比如用户 ID、访问令牌等。我们可能会想把这些数据全局存储起来，但这其实不太妥当。最大的问题是，我们需要维护线程与数据之间的映射关系，并且还得处理竞争条件的锁。而这个问题可以通过一种叫 线程局部变量 的概念来解决。线程局部变量是一种只属于某个特定线程的全局状态。一旦在线程里设置了数据，只有那个线程才能看到它，避免了复杂的线程-数据映射问题和竞态条件。虽然我们不深入展开线程局部变量的细节，但你可以参考 threading 模块文档中的相关内容：线程局部数据(https://docs.python.org/3/library/threading.html#thread-local-data)。

当然，在 asyncio 应用中，通常只有一个线程，所以任何设置为线程局部的数据都可以在整个应用中随意访问。而 PEP-567 提出了一个新概念——上下文变量，用来解决单线程并发模型下的“线程局部变量”问题。上下文变量和线程局部变量类似，但区别在于它是作用于一个特定的任务，而不是一个线程。这意味着，如果一个任务创建了一个上下文变量，那么该任务内的任何内层协程或子任务都能访问这个变量。其他外部的任务则完全看不到也修改不了这个变量。这样一来，我们就可以轻松地追踪某个特定任务的专属状态，而无需将其作为显式参数传递。

来看个实际例子：我们来写一个简单的服务器，监听来自连接客户端的数据。我们用一个上下文变量来记录已连接用户的地址，当用户发送消息时，我们不仅打印出消息内容，还显示发送者的地址。

import asynciofrom asyncio import StreamReader, StreamWriterfrom contextvars import ContextVarclass Server:    user_address = ContextVar('user_address')  # ❶    def __init__(self, host: str, port: int):        self.host = host        self.port = port    async def start_server(self):        server = await asyncio.start_server(self._client_connected,                                            self.host, self.port)        await server.serve_forever()    def _client_connected(self, reader: StreamReader, writer: StreamWriter):        self.user_address.set(writer.get_extra_info('peername'))  # ❷        asyncio.create_task(self.listen_for_messages(reader))    async def listen_for_messages(self, reader: StreamReader):        while data := await reader.readline():            print(f'Got message {data} from {self.user_address.get()}')  # ❸async def main():    server = Server('127.0.0.1', 9000)    await server.start_server()asyncio.run(main())

在上面的例子中，我们首先创建了一个 ContextVar 实例来存放用户地址信息。上下文变量需要一个字符串名称，所以我们给了一个有意义的名字 user_address，主要是为了调试方便。然后在 _client_connected 回调里，我们把这个客户端的地址设置到上下文变量里。这样，从这个父任务派生出来的任何任务都能访问到我们设置的数据，在这里就是负责监听客户端消息的任务。

在 listen_for_messages 协程里，我们监听客户端的数据，拿到后就一边打印出来，一边把之前存好的上下文变量里的地址一并显示。当你运行这个程序，并用多个客户端连接发送消息时，你应该能看到类似下面的输出：

Got message b'Hello!\r\n' from ('127.0.0.1', 50036)Got message b'Okay!\r\n' from ('127.0.0.1', 50038)

注意地址里的端口号不一样，说明我们确实收到了来自两个不同客户端的消息。尽管我们只创建了一个上下文变量，但依然能分别获取到每个客户端的独立数据。这种方式提供了一种非常干净的方法来在任务间传递数据，而无需显式地将数据作为参数传入。