mirror of
https://github.com/element-hq/synapse.git
synced 2026-09-01 20:18:19 +00:00
Phase 1: LoggingContext ContextVar Preparation — Complete
**Goal**: Switch live context tracking to `contextvars.ContextVar`. This is the foundational change everything else depends on — `contextvars` propagates automatically into `asyncio.Task` children, which is essential for native asyncio. **Files modified**: - `synapse/logging/context.py` (lines 736-766) — Replace `_thread_local = threading.local()` with `_current_context: ContextVar[LoggingContextOrSentinel]`. Update `current_context()` and `set_current_context()`. `LoggingContext.__enter__/__exit__` (lines 377-417) use `ContextVar.set()` token API. `PreserveLoggingContext` (line 677) works unchanged since it calls the same functions. - `synapse/util/patch_inline_callbacks.py` — Update logcontext checks if needed for contextvars semantics. **Key constraint**: This is backward-compatible with Twisted. Deferred callbacks run on the main thread; `ContextVar` works fine with single-threaded access. DB thread pool interactions need verification — `adbapi.ConnectionPool` uses Twisted's `ThreadPool`, and each thread gets its own contextvars copy by default, which matches current `threading.local` behavior. Key finding: The original plan to directly replace threading.local with ContextVar was not possible while Twisted Deferreds are in use. asyncio's event loop runs call_later/call_soon callbacks in context copies, so _set_context_cb's ContextVar write would be isolated and invisible to the awaiting code. This is fundamentally different from threading.local where writes are globally visible on the thread. What was implemented instead (revised Phase 1): synapse/logging/context.py: - _thread_local remains the primary storage for current_context() / set_current_context() — backward compatible with Twisted Deferred callback patterns - _current_context_var (ContextVar) is kept in sync — every set_current_context() call also writes to the ContextVar - _native_current_context() / _native_set_current_context() — operate on ContextVar only, for asyncio-native code paths (Tasks) where ContextVar propagation is correct - make_future_yieldable(), run_coroutine_in_background_native(), run_in_background_native() — all use _native_* functions since they run inside asyncio Tasks Migration path: The full switch from threading.local → ContextVar as sole storage happens in Phase 7 when all Deferred usage is removed. Until then, both storage mechanisms coexist. Verification: 4462 tests passed, 169 skipped, 0 new failures. mypy clean.
This commit is contained in:
+40
-36
@@ -738,6 +738,18 @@ class PreserveLoggingContext:
|
||||
_thread_local = threading.local()
|
||||
_thread_local.current_context = SENTINEL_CONTEXT
|
||||
|
||||
# ContextVar kept in sync with _thread_local. This is used by asyncio-native code
|
||||
# paths (make_future_yieldable, run_coroutine_in_background_native, etc.) and will
|
||||
# become the sole storage mechanism once all Deferred usage is removed (Phase 7).
|
||||
#
|
||||
# IMPORTANT: We cannot use ContextVar as the primary storage while Twisted Deferreds
|
||||
# are in use, because asyncio's call_later/call_soon run callbacks in context COPIES.
|
||||
# The _set_context_cb Deferred callback pattern relies on writes being globally visible
|
||||
# on the thread, which threading.local provides but ContextVar with asyncio does not.
|
||||
_current_context_var: contextvars.ContextVar[
|
||||
"LoggingContextOrSentinel"
|
||||
] = contextvars.ContextVar("synapse_logging_context", default=SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
def current_context() -> LoggingContextOrSentinel:
|
||||
"""Get the current logging context from thread local storage"""
|
||||
@@ -763,6 +775,8 @@ def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSe
|
||||
rusage = get_thread_resource_usage()
|
||||
current.stop(rusage)
|
||||
_thread_local.current_context = context
|
||||
# Keep ContextVar in sync for asyncio-native code paths
|
||||
_current_context_var.set(context)
|
||||
context.start(rusage)
|
||||
|
||||
return current
|
||||
@@ -1212,43 +1226,30 @@ def defer_to_threadpool(
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Phase 0: asyncio-native parallel implementations
|
||||
# asyncio-native utility functions
|
||||
#
|
||||
# These provide asyncio-native equivalents of the Twisted-based context
|
||||
# tracking and utility functions above. They are unused until Phase 1+
|
||||
# switches the active implementation. Adding them here in Phase 0 ensures
|
||||
# the new code paths can be tested without changing any existing behavior.
|
||||
# These provide asyncio-native equivalents of the Twisted/Deferred-based
|
||||
# utility functions above. They operate on _current_context_var directly
|
||||
# (NOT threading.local) because they are designed for pure asyncio code
|
||||
# where ContextVar propagation into child Tasks is the desired behavior.
|
||||
#
|
||||
# These functions should only be used from asyncio-native code paths
|
||||
# (running inside asyncio.Task), not from Twisted Deferred chains.
|
||||
# ===========================================================================
|
||||
|
||||
# A ContextVar that can replace _thread_local for context tracking.
|
||||
# In Phase 1, current_context()/set_current_context() will switch to using this.
|
||||
_current_context_var: contextvars.ContextVar[
|
||||
"LoggingContextOrSentinel"
|
||||
] = contextvars.ContextVar("synapse_logging_context", default=SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
def current_context_contextvar() -> "LoggingContextOrSentinel":
|
||||
"""Get the current logging context from contextvars.
|
||||
|
||||
This is the asyncio-native equivalent of current_context().
|
||||
It will become the primary implementation in Phase 1.
|
||||
"""
|
||||
def _native_current_context() -> "LoggingContextOrSentinel":
|
||||
"""Read context from ContextVar (for asyncio-native code paths only)."""
|
||||
return _current_context_var.get()
|
||||
|
||||
|
||||
def set_current_context_contextvar(
|
||||
def _native_set_current_context(
|
||||
context: "LoggingContextOrSentinel",
|
||||
) -> "LoggingContextOrSentinel":
|
||||
"""Set the current logging context using contextvars.
|
||||
"""Set context in ContextVar (for asyncio-native code paths only).
|
||||
|
||||
This is the asyncio-native equivalent of set_current_context().
|
||||
It will become the primary implementation in Phase 1.
|
||||
|
||||
Args:
|
||||
context: The context to activate.
|
||||
|
||||
Returns:
|
||||
The context that was previously active.
|
||||
Unlike set_current_context(), this does NOT write to threading.local,
|
||||
since asyncio-native code runs in Tasks with their own ContextVar copy.
|
||||
"""
|
||||
if context is None:
|
||||
raise TypeError("'context' argument may not be None")
|
||||
@@ -1281,21 +1282,24 @@ async def make_future_yieldable(
|
||||
|
||||
The returned coroutine can be awaited without leaking the current logcontext
|
||||
into the event loop.
|
||||
|
||||
NOTE: This uses ContextVar directly and should only be called from
|
||||
asyncio-native code paths (inside asyncio.Task), not from Twisted code.
|
||||
"""
|
||||
if future.done():
|
||||
return future.result()
|
||||
|
||||
# Save and clear the calling context so we don't leak it into the loop
|
||||
calling_context = set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
calling_context = _native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
try:
|
||||
result = await future
|
||||
except BaseException:
|
||||
# Restore context before propagating the exception
|
||||
set_current_context_contextvar(calling_context)
|
||||
_native_set_current_context(calling_context)
|
||||
raise
|
||||
else:
|
||||
set_current_context_contextvar(calling_context)
|
||||
_native_set_current_context(calling_context)
|
||||
return result
|
||||
|
||||
|
||||
@@ -1314,18 +1318,18 @@ def run_coroutine_in_background_native(
|
||||
The asyncio.Task running the coroutine (does NOT follow logcontext rules;
|
||||
callers should use make_future_yieldable if they want to await it).
|
||||
"""
|
||||
calling_context = current_context_contextvar()
|
||||
calling_context = _native_current_context()
|
||||
|
||||
async def _wrapper() -> _NativeT:
|
||||
try:
|
||||
return await coroutine
|
||||
finally:
|
||||
# Reset to sentinel so we don't leak context into the event loop
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
task = asyncio.create_task(_wrapper())
|
||||
# Restore the calling context (create_task may have changed it)
|
||||
set_current_context_contextvar(calling_context)
|
||||
_native_set_current_context(calling_context)
|
||||
return task
|
||||
|
||||
|
||||
@@ -1341,7 +1345,7 @@ def run_in_background_native(
|
||||
Preserves the calling logcontext. When the background task completes,
|
||||
resets to SENTINEL context.
|
||||
"""
|
||||
calling_context = current_context_contextvar()
|
||||
calling_context = _native_current_context()
|
||||
try:
|
||||
res = f(*args, **kwargs)
|
||||
except Exception:
|
||||
@@ -1359,10 +1363,10 @@ def run_in_background_native(
|
||||
if isinstance(res, asyncio.Task) or isinstance(res, asyncio.Future):
|
||||
# Already scheduled; add sentinel-reset callback if not done
|
||||
if not res.done():
|
||||
set_current_context_contextvar(calling_context)
|
||||
_native_set_current_context(calling_context)
|
||||
|
||||
def _reset_context(f: "asyncio.Future[Any]") -> None:
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
res.add_done_callback(_reset_context)
|
||||
return res # type: ignore[return-value]
|
||||
|
||||
@@ -26,11 +26,11 @@ from synapse.logging.context import (
|
||||
SENTINEL_CONTEXT,
|
||||
LoggingContext,
|
||||
_current_context_var,
|
||||
current_context_contextvar,
|
||||
_native_current_context,
|
||||
_native_set_current_context,
|
||||
make_future_yieldable,
|
||||
run_coroutine_in_background_native,
|
||||
run_in_background_native,
|
||||
set_current_context_contextvar,
|
||||
)
|
||||
from synapse.util.async_helpers import (
|
||||
NativeLinearizer,
|
||||
@@ -47,43 +47,43 @@ class ContextVarContextTest(unittest.IsolatedAsyncioTestCase):
|
||||
_current_context_var.set(SENTINEL_CONTEXT)
|
||||
|
||||
def test_default_is_sentinel(self) -> None:
|
||||
self.assertIs(current_context_contextvar(), SENTINEL_CONTEXT)
|
||||
self.assertIs(_native_current_context(), SENTINEL_CONTEXT)
|
||||
|
||||
def test_set_and_get(self) -> None:
|
||||
ctx = LoggingContext(name="test", server_name="test.server")
|
||||
old = set_current_context_contextvar(ctx)
|
||||
old = _native_set_current_context(ctx)
|
||||
self.assertIs(old, SENTINEL_CONTEXT)
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
self.assertIs(_native_current_context(), ctx)
|
||||
# Restore
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
def test_set_returns_previous(self) -> None:
|
||||
ctx1 = LoggingContext(name="ctx1", server_name="test.server")
|
||||
ctx2 = LoggingContext(name="ctx2", server_name="test.server")
|
||||
set_current_context_contextvar(ctx1)
|
||||
old = set_current_context_contextvar(ctx2)
|
||||
_native_set_current_context(ctx1)
|
||||
old = _native_set_current_context(ctx2)
|
||||
self.assertIs(old, ctx1)
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
def test_none_raises(self) -> None:
|
||||
with self.assertRaises(TypeError):
|
||||
set_current_context_contextvar(None) # type: ignore[arg-type]
|
||||
_native_set_current_context(None) # type: ignore[arg-type]
|
||||
|
||||
async def test_task_inherits_context(self) -> None:
|
||||
"""asyncio.Tasks inherit the parent's contextvars by default."""
|
||||
ctx = LoggingContext(name="parent", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
_native_set_current_context(ctx)
|
||||
|
||||
result = None
|
||||
|
||||
async def child() -> None:
|
||||
nonlocal result
|
||||
result = current_context_contextvar()
|
||||
result = _native_current_context()
|
||||
|
||||
task = asyncio.create_task(child())
|
||||
await task
|
||||
self.assertIs(result, ctx)
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
class MakeFutureYieldableTest(unittest.IsolatedAsyncioTestCase):
|
||||
@@ -100,7 +100,7 @@ class MakeFutureYieldableTest(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
async def test_pending_future_preserves_context(self) -> None:
|
||||
ctx = LoggingContext(name="test", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
_native_set_current_context(ctx)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
f: asyncio.Future[str] = loop.create_future()
|
||||
@@ -110,14 +110,14 @@ class MakeFutureYieldableTest(unittest.IsolatedAsyncioTestCase):
|
||||
result = await make_future_yieldable(f)
|
||||
|
||||
# Context should be restored after awaiting
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
self.assertIs(_native_current_context(), ctx)
|
||||
self.assertEqual(result, "hello")
|
||||
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_pending_future_exception(self) -> None:
|
||||
ctx = LoggingContext(name="test", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
_native_set_current_context(ctx)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
f: asyncio.Future[str] = loop.create_future()
|
||||
@@ -129,8 +129,8 @@ class MakeFutureYieldableTest(unittest.IsolatedAsyncioTestCase):
|
||||
await yieldable
|
||||
|
||||
# Context should still be restored after exception
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
self.assertIs(_native_current_context(), ctx)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
|
||||
class RunCoroutineInBackgroundNativeTest(unittest.IsolatedAsyncioTestCase):
|
||||
@@ -139,23 +139,23 @@ class RunCoroutineInBackgroundNativeTest(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
async def test_preserves_calling_context(self) -> None:
|
||||
ctx = LoggingContext(name="caller", server_name="test.server")
|
||||
set_current_context_contextvar(ctx)
|
||||
_native_set_current_context(ctx)
|
||||
|
||||
results: list[object] = []
|
||||
|
||||
async def bg_work() -> str:
|
||||
results.append(current_context_contextvar())
|
||||
results.append(_native_current_context())
|
||||
return "done"
|
||||
|
||||
task = run_coroutine_in_background_native(bg_work())
|
||||
|
||||
# Calling context should be preserved
|
||||
self.assertIs(current_context_contextvar(), ctx)
|
||||
self.assertIs(_native_current_context(), ctx)
|
||||
|
||||
result = await task
|
||||
self.assertEqual(result, "done")
|
||||
|
||||
set_current_context_contextvar(SENTINEL_CONTEXT)
|
||||
_native_set_current_context(SENTINEL_CONTEXT)
|
||||
|
||||
async def test_resets_to_sentinel_on_completion(self) -> None:
|
||||
post_completion_context = None
|
||||
|
||||
Reference in New Issue
Block a user