AberSheeran
Aber Sheeran

ContextVars 详解

起笔自
所属文集: Python-Package
共计 2012 个字符
落笔于

ContextVars 是 Python3.7 以及之后版本里的标准库。

应用场景

Context managers that have state should use Context Variables instead of threading.local() to prevent their state from bleeding to other code unexpectedly, when used in concurrent code.

从官方文档的说明来看,它是用于解决并发代码里的长生命周期变量的管理问题的。这里的“并发代码”就是指使用了 async/await 的协程代码。

那么,什么是“长生命周期变量”?如果你使用过 bottle/flask 当中的一个就能知道 request 这个变量能在框架处理请求的大部分过程中被调用,并且在不同的请求之间,它的值是互相隔离的。它虽然被创建在整个调用过程的开始,却不需要显式的层层传递给每个函数。它看起来像是一个全局变量,但却不是程序结束才被销毁、并且也不是全局共享的——它对于每个请求来说都是唯一无二的变量。

在传统的多线程编程里,使用 threading.local() 就可以做到这一点,但在 async/await 这种多协程编程里,往往一个线程内要跑许多个协程,使用 threading.local() 无法达到这个目的。但协程又号称微线程,对标的是用户级线程,那么自然也需要一种方式来隔离多个协程之间的长生命周期变量,于是 ContextVars 应运而生。

实现原理

ContextVars 是怎么做到协程间隔离变量的呢?

https://github.com/MagicStack/contextvars 是同一个作者为低版本 Python 写的兼容实现,可以看出来它只是继承了 Mapping 然后实现了 Copy 方法而已,唯一的隔离在 copy_context 里。

在它的相关 PEP 文档——PEP-567里提到,在 asyncio 几个关键函数里增加了 copy_context 方法,然后再调用 context.run 去执行回调函数。

简单粗暴但很有效的做法。

这与我们直接使用此标准库提供的接口 context.run 跑普通函数没有区别,asyncio 只不过是把需要程序员手动 copy_context 的行为放到源码里执行了。

与 asyncio 一起使用

import asyncio
import contextvars

client_addr_var = contextvars.ContextVar('client_addr')

def render_goodbye():
    # 从 client_addr_var 获取当前连接地址
    client_addr = client_addr_var.get()
    return f'Good bye, client @ {client_addr}\n'.encode()

async def handle_request(reader, writer):
    addr = writer.transport.get_extra_info('socket').getpeername()
    # 在设置完之后,就可以使用 client_addr_var.get() 获取设置的值
    token = client_addr_var.set(addr)

    try:
        while True:
            line = await reader.readline()
            print(line)
            if not line.strip():
                break
            writer.write(line)

        writer.write(render_goodbye())
        writer.close()
    finally:
        # 使用 token 重置值到上次 set 前
        client_addr_var.reset(token)

async def main():
    srv = await asyncio.start_server(
        handle_request, '127.0.0.1', 8081)

    async with srv:
        await srv.serve_forever()

asyncio.run(main())
如果你觉得本文值得,不妨赏杯茶
pip 安装与脚本安装
PDM 使用精要