Python:将上下文(contextvars.Context)复制到单独的线程

Question

就目前而言，我已经找到了很多有关 contextvars 模块如何与 asyncio 配合使用的示例，但没有找到有关如何与线程配合使用的示例(asyncio.get_event_loop().run_in_executor、threading.Thread 等)。

我的问题是，如何将上下文传递给单独的线程？下面您可以看到一个不起作用的代码片段(python 3.9.8)。

import typing
import asyncio
import contextvars
import concurrent.futures


class CustomThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor):
    def submit(
        self,
        function: typing.Callable,
        *args,
        **kwargs
    ) -> concurrent.futures.Future:
        context = contextvars.copy_context()
        return super().submit(
            context.run,
            functools.partial(function, *args, **kwargs)
        )


def function():
    print(var.get())


async def main():
    await asyncio.get_event_loop().run_in_executor(None, function)


if __name__ == '__main__':
    var = contextvars.ContextVar('variable')
    var.set('Message.')

    asyncio.get_event_loop().set_default_executor(CustomThreadPoolExecutor)
    asyncio.run(main())

Answer 1

您可以使用采用 copy_context.items() 的包装函数，设置它们并调用您的函数。 functools.partial 将帮助您创建包装函数以传递给 run_in_executor。这是我的装饰器的工作测试:

def test_run_in_thread_pool_executor():
    def init(func, ctx_vars, *args, **kwargs):
        for var, value in ctx_vars:
            var.set(value)
        return func(*args, **kwargs)

    @async_add_headers('streaming')
    async def wrapper(f):
        loop = asyncio.get_event_loop()
        ctx = contextvars.copy_context()
        executor = futures.ThreadPoolExecutor(max_workers=5)
        return await loop.run_in_executor(executor, functools.partial(init, f, ctx.items()))

    @add_headers('client')
    def foo():
        assert caller_context_var.get() == 'streaming'

    async def main_test():
        await wrapper(foo)

    asyncio.run(main_test())

这里 add_headers 和 async_add_headers 按照调用函数的顺序更改了一些上下文变量。 caller_context_var.get() 将等于没有 init 函数的 'client'。

不幸的是，它仅适用于 ThreadPoolExecutor，不适用于 ProcessPoolExecutor，因为 Context 对象不可picklable。查看亲戚PEP 567 section 。还有执行器的示例:

executor = ThreadPoolExecutor()
current_context = contextvars.copy_context()

executor.submit(current_context.run, some_function)