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:
Matthew Hodgson
2026-03-21 14:25:08 +00:00
parent 2dce74958f
commit 24724a810e
2 changed files with 63 additions and 59 deletions
+40 -36
View File
@@ -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]
+23 -23
View File
@@ -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