mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-29 15:31:05 +08:00
fix: propagate contextvars through thread_pool_exec (#16247)
## Problem
`thread_pool_exec()` dispatches work via `loop.run_in_executor()`, which
submits the callable with a plain `executor.submit(func, *args)` and
does **not** copy the caller's `contextvars.Context`. So a `ContextVar`
set in the async caller is not visible inside the function running in
the worker thread.
This differs from `asyncio.to_thread()`, which runs the callable inside
a copied context. `run_in_executor()` has never propagated context
(verified on Python 3.12 and 3.13) — so this is a pre-existing gap in
the helper, **not** a regression or a Python-version compatibility
issue.
Concretely, any code that sets a `ContextVar` in async code and reads it
inside a function dispatched via `thread_pool_exec` (request tracing,
per-task state, Langfuse trace propagation, etc.) silently loses that
context.
## Fix
Copy the current context before submitting and run the callable inside
it with `ctx.run()`, matching what `asyncio.to_thread()` does:
```python
async def thread_pool_exec(func, *args, **kwargs):
loop = asyncio.get_running_loop()
ctx = contextvars.copy_context()
if kwargs:
inner = functools.partial(func, *args, **kwargs)
return await loop.run_in_executor(_thread_pool_executor(), ctx.run, inner)
return await loop.run_in_executor(_thread_pool_executor(), ctx.run, func, *args)
```
This explicitly **adds** ContextVar propagation to the helper (it does
not restore any prior behavior). Backward-compatible.
## Tests
`TestThreadPoolExec` covers propagation, the kwargs path, per-call
isolation and the unset-default case.
> Note: the branch name still contains `python313` for historical
reasons; the change is unrelated to any Python version.
This commit is contained in:
@@ -16,6 +16,7 @@
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextvars
|
||||
import functools
|
||||
import hashlib
|
||||
import logging
|
||||
@@ -250,8 +251,13 @@ def _thread_pool_executor():
|
||||
|
||||
|
||||
async def thread_pool_exec(func, *args, **kwargs):
|
||||
# loop.run_in_executor() submits the callable without propagating the caller's
|
||||
# contextvars (unlike asyncio.to_thread, which copies the context). Copy the
|
||||
# current context and run the callable inside it so ContextVars set by the
|
||||
# caller (e.g. tracing / per-request state) are visible in the worker thread.
|
||||
loop = asyncio.get_running_loop()
|
||||
ctx = contextvars.copy_context()
|
||||
if kwargs:
|
||||
func = functools.partial(func, *args, **kwargs)
|
||||
return await loop.run_in_executor(_thread_pool_executor(), func)
|
||||
return await loop.run_in_executor(_thread_pool_executor(), func, *args)
|
||||
inner = functools.partial(func, *args, **kwargs)
|
||||
return await loop.run_in_executor(_thread_pool_executor(), ctx.run, inner)
|
||||
return await loop.run_in_executor(_thread_pool_executor(), ctx.run, func, *args)
|
||||
|
||||
Reference in New Issue
Block a user