Store the LoggingContext in a ContextVar

changelog.d/18871.misc

@@ -0,0 +1 @@
+Store the `LoggingContext` in a `ContextVar` instead of a thread-local variable.

synapse/__init__.py

@@ -86,12 +86,5 @@ def _immutabledict_cb(d: immutabledict) -> Dict[str, Any]:
 
 __version__ = synapse.util.SYNAPSE_VERSION
 
-if bool(os.environ.get("SYNAPSE_TEST_PATCH_LOG_CONTEXTS", False)):
-    # We import here so that we don't have to install a bunch of deps when
-    # running the packaging tox test.
-    from synapse.util.patch_inline_callbacks import do_patch
-
-    do_patch()
-
 
 check_rust_lib_up_to_date()

synapse/app/_base.py

@@ -601,6 +601,12 @@ def run_sighup(*args: Any, **kwargs: Any) -> None:
     hs.get_datastores().main.db_pool.start_profiling()
     hs.get_pusherpool().start()
 
+    # Register background tasks required by this server. This must be done
+    # somewhat manually due to the background tasks not being registered
+    # unless handlers are instantiated.
+    if hs.config.worker.run_background_tasks:
+        hs.start_background_tasks()
+
     # Log when we start the shut down process.
     hs.get_reactor().addSystemEventTrigger(
         "before", "shutdown", logger.info, "Shutting down..."

synapse/logging/context.py

@@ -34,6 +34,7 @@
 import threading
 import typing
 import warnings
+from contextvars import ContextVar
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
@@ -653,13 +654,12 @@ def __exit__(
                 )
 
 
-_thread_local = threading.local()
-_thread_local.current_context = SENTINEL_CONTEXT
+_current_context: ContextVar[LoggingContextOrSentinel] = ContextVar("current_context")
 
 
 def current_context() -> LoggingContextOrSentinel:
     """Get the current logging context from thread local storage"""
-    return getattr(_thread_local, "current_context", SENTINEL_CONTEXT)
+    return _current_context.get(SENTINEL_CONTEXT)
 
 
 def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSentinel:
@@ -680,7 +680,7 @@ def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSe
     if current is not context:
         rusage = get_thread_resource_usage()
         current.stop(rusage)
-        _thread_local.current_context = context
+        _current_context.set(context)
         context.start(rusage)
 
     return current
@@ -796,7 +796,6 @@ def run_in_background(
     CRITICAL error about an unhandled error will be logged without much
     indication about where it came from.
     """
-    current = current_context()
     try:
         res = f(*args, **kwargs)
     except Exception:
@@ -825,23 +824,6 @@ def run_in_background(
         # optimise out the messing about
         return d
 
-    # The function may have reset the context before returning, so
-    # we need to restore it now.
-    ctx = set_current_context(current)
-
-    # The original context will be restored when the deferred
-    # completes, but there is nothing waiting for it, so it will
-    # get leaked into the reactor or some other function which
-    # wasn't expecting it. We therefore need to reset the context
-    # here.
-    #
-    # (If this feels asymmetric, consider it this way: we are
-    # effectively forking a new thread of execution. We are
-    # probably currently within a ``with LoggingContext()`` block,
-    # which is supposed to have a single entry and exit point. But
-    # by spawning off another deferred, we are effectively
-    # adding a new exit point.)
-    d.addBoth(_set_context_cb, ctx)
     return d
 
 
@@ -861,65 +843,20 @@ def run_coroutine_in_background(
     cannot change the log contexts.
     """
 
-    current = current_context()
-    d = defer.ensureDeferred(coroutine)
-
-    # The function may have reset the context before returning, so
-    # we need to restore it now.
-    ctx = set_current_context(current)
-
-    # The original context will be restored when the deferred
-    # completes, but there is nothing waiting for it, so it will
-    # get leaked into the reactor or some other function which
-    # wasn't expecting it. We therefore need to reset the context
-    # here.
-    #
-    # (If this feels asymmetric, consider it this way: we are
-    # effectively forking a new thread of execution. We are
-    # probably currently within a ``with LoggingContext()`` block,
-    # which is supposed to have a single entry and exit point. But
-    # by spawning off another deferred, we are effectively
-    # adding a new exit point.)
-    d.addBoth(_set_context_cb, ctx)
-    return d
+    return defer.ensureDeferred(coroutine)
 
 
 T = TypeVar("T")
 
 
+# TODO: This function is a no-op now and should be removed in a follow-up PR.
 def make_deferred_yieldable(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
-    """Given a deferred, make it follow the Synapse logcontext rules:
-
-    If the deferred has completed, essentially does nothing (just returns another
-    completed deferred with the result/failure).
-
-    If the deferred has not yet completed, resets the logcontext before
-    returning a deferred. Then, when the deferred completes, restores the
-    current logcontext before running callbacks/errbacks.
-
-    (This is more-or-less the opposite operation to run_in_background.)
-    """
-    if deferred.called and not deferred.paused:
-        # it looks like this deferred is ready to run any callbacks we give it
-        # immediately. We may as well optimise out the logcontext faffery.
-        return deferred
-
-    # ok, we can't be sure that a yield won't block, so let's reset the
-    # logcontext, and add a callback to the deferred to restore it.
-    prev_context = set_current_context(SENTINEL_CONTEXT)
-    deferred.addBoth(_set_context_cb, prev_context)
     return deferred
 
 
 ResultT = TypeVar("ResultT")
 
 
-def _set_context_cb(result: ResultT, context: LoggingContextOrSentinel) -> ResultT:
-    """A callback function which just sets the logging context"""
-    set_current_context(context)
-    return result
-
-
 def defer_to_thread(
     reactor: "ISynapseReactor", f: Callable[P, R], *args: P.args, **kwargs: P.kwargs
 ) -> "defer.Deferred[R]":
@@ -931,9 +868,6 @@ def defer_to_thread(
     logcontext (so its CPU usage metrics will get attributed to the current
     logcontext). `f` should preserve the logcontext it is given.
 
-    The result deferred follows the Synapse logcontext rules: you should `yield`
-    on it.
-
     Args:
         reactor: The reactor in whose main thread the Deferred will be invoked,
             and whose threadpool we should use for the function.
@@ -971,9 +905,6 @@ def defer_to_threadpool(
     logcontext (so its CPU usage metrics will get attributed to the current
     logcontext). `f` should preserve the logcontext it is given.
 
-    The result deferred follows the Synapse logcontext rules: you should `yield`
-    on it.
-
     Args:
         reactor: The reactor in whose main thread the Deferred will be invoked.
             Normally this will be hs.get_reactor().
@@ -991,18 +922,6 @@ def defer_to_threadpool(
         A Deferred which fires a callback with the result of `f`, or an
             errback if `f` throws an exception.
     """
-    curr_context = current_context()
-    if not curr_context:
-        logger.warning(
-            "Calling defer_to_threadpool from sentinel context: metrics will be lost"
-        )
-        parent_context = None
-    else:
-        assert isinstance(curr_context, LoggingContext)
-        parent_context = curr_context
-
-    def g() -> R:
-        with LoggingContext(str(curr_context), parent_context=parent_context):
-            return f(*args, **kwargs)
-
-    return make_deferred_yieldable(threads.deferToThreadPool(reactor, threadpool, g))
+    return make_deferred_yieldable(
+        threads.deferToThreadPool(reactor, threadpool, f, *args, **kwargs)
+    )

synapse/metrics/background_process_metrics.py

@@ -223,10 +223,9 @@ def run_as_background_process(
     This should be used to wrap processes which are fired off to run in the
     background, instead of being associated with a particular request.
 
-    It returns a Deferred which completes when the function completes, but it doesn't
-    follow the synapse logcontext rules, which makes it appropriate for passing to
-    clock.looping_call and friends (or for firing-and-forgetting in the middle of a
-    normal synapse async function).
+    It returns a Deferred which completes when the function completes, which makes it
+    appropriate for passing to clock.looping_call and friends (or for
+    firing-and-forgetting in the middle of a normal synapse async function).
 
     Args:
         desc: a description for this background process type
@@ -241,8 +240,6 @@ def run_as_background_process(
 
     Returns:
         Deferred which returns the result of func, or `None` if func raises.
-        Note that the returned Deferred does not follow the synapse logcontext
-        rules.
     """
 
     async def run() -> Optional[R]:

synapse/module_api/__init__.py

@@ -237,10 +237,9 @@ def run_as_background_process(
     This should be used to wrap processes which are fired off to run in the
     background, instead of being associated with a particular request.
 
-    It returns a Deferred which completes when the function completes, but it doesn't
-    follow the synapse logcontext rules, which makes it appropriate for passing to
-    clock.looping_call and friends (or for firing-and-forgetting in the middle of a
-    normal synapse async function).
+    It returns a Deferred which completes when the function completes, which makes it
+    appropriate for passing to clock.looping_call and friends (or for
+    firing-and-forgetting in the middle of a normal synapse async function).
 
     Args:
         desc: a description for this background process type
@@ -255,8 +254,6 @@ def run_as_background_process(
 
     Returns:
         Deferred which returns the result of func, or `None` if func raises.
-        Note that the returned Deferred does not follow the synapse logcontext
-        rules.
     """
 
     logger.warning(
@@ -1375,9 +1372,7 @@ def looping_background_call(
 
         Args:
             f: The function to call repeatedly. f can be either synchronous or
-                asynchronous, and must follow Synapse's logcontext rules.
-                More info about logcontexts is available at
-                https://element-hq.github.io/synapse/latest/log_contexts.html
+                asynchronous.
             msec: How long to wait between calls in milliseconds.
             *args: Positional arguments to pass to function.
             desc: The background task's description. Default to the function's name.
@@ -1431,9 +1426,7 @@ def delayed_background_call(
         Args:
             msec: How long to wait before calling, in milliseconds.
             f: The function to call once. f can be either synchronous or
-                asynchronous, and must follow Synapse's logcontext rules.
-                More info about logcontexts is available at
-                https://element-hq.github.io/synapse/latest/log_contexts.html
+                asynchronous.
             *args: Positional arguments to pass to function.
             desc: The background task's description. Default to the function's name.
             **kwargs: Keyword arguments to pass to function.
@@ -1668,10 +1661,9 @@ def run_as_background_process(
         This should be used to wrap processes which are fired off to run in the
         background, instead of being associated with a particular request.
 
-        It returns a Deferred which completes when the function completes, but it doesn't
-        follow the synapse logcontext rules, which makes it appropriate for passing to
-        clock.looping_call and friends (or for firing-and-forgetting in the middle of a
-        normal synapse async function).
+        It returns a Deferred which completes when the function completes, which makes
+        it appropriate for passing to clock.looping_call and friends (or for
+        firing-and-forgetting in the middle of a normal synapse async function).
 
         Args:
             desc: a description for this background process type
@@ -1686,8 +1678,6 @@ def run_as_background_process(
 
         Returns:
             Deferred which returns the result of func, or `None` if func raises.
-            Note that the returned Deferred does not follow the synapse logcontext
-            rules.
         """
         return _run_as_background_process(
             desc, self.server_name, func, *args, bg_start_span=bg_start_span, **kwargs

synapse/server.py

@@ -366,12 +366,6 @@ def setup(self) -> None:
         self.datastores = Databases(self.DATASTORE_CLASS, self)
         logger.info("Finished setting up.")
 
-        # Register background tasks required by this server. This must be done
-        # somewhat manually due to the background tasks not being registered
-        # unless handlers are instantiated.
-        if self.config.worker.run_background_tasks:
-            self.setup_background_tasks()
-
     def __del__(self) -> None:
         """
         Called when an the homeserver is garbage collected.
@@ -410,7 +404,7 @@ def start_listening(self) -> None:  # noqa: B027 (no-op by design)
         appropriate listeners.
         """
 
-    def setup_background_tasks(self) -> None:
+    def start_background_tasks(self) -> None:
         """
         Some handlers have side effects on instantiation (like registering
         background updates). This function causes them to be fetched, and

synapse/util/async_helpers.py

@@ -84,9 +84,6 @@ def observe(self) -> "defer.Deferred[_T]":
         This returns a brand new deferred that is resolved when the underlying
         deferred is resolved. Interacting with the returned deferred does not
         effect the underlying deferred.
-
-        Note that the returned Deferred doesn't follow the Synapse logcontext rules -
-        you will probably want to `make_deferred_yieldable` it.
         """
         ...
 
@@ -100,11 +97,6 @@ class ObservableDeferred(Generic[_T], AbstractObservableDeferred[_T]):
 
     Cancelling or otherwise resolving an observer will not affect the original
     ObservableDeferred.
-
-    NB that it does not attempt to do anything with logcontexts; in general
-    you should probably make_deferred_yieldable the deferreds
-    returned by `observe`, and ensure that the original deferred runs its
-    callbacks in the sentinel logcontext.
     """
 
     __slots__ = ["_deferred", "_observers", "_result"]
@@ -861,16 +853,12 @@ def stop_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
     """Prevent a `Deferred` from being cancelled by wrapping it in another `Deferred`.
 
     Args:
-        deferred: The `Deferred` to protect against cancellation. Must not follow the
-            Synapse logcontext rules.
+        deferred: The `Deferred` to protect against cancellation.
 
     Returns:
         A new `Deferred`, which will contain the result of the original `Deferred`.
         The new `Deferred` will not propagate cancellation through to the original.
         When cancelled, the new `Deferred` will fail with a `CancelledError`.
-
-        The new `Deferred` will not follow the Synapse logcontext rules and should be
-        wrapped with `make_deferred_yieldable`.
     """
     new_deferred: "defer.Deferred[T]" = defer.Deferred()
     deferred.chainDeferred(new_deferred)
@@ -896,8 +884,7 @@ def delay_cancellation(awaitable: Awaitable[T]) -> Awaitable[T]:
     resolve with a `CancelledError` until the original awaitable resolves.
 
     Args:
-        deferred: The coroutine or `Deferred` to protect against cancellation. May
-            optionally follow the Synapse logcontext rules.
+        deferred: The coroutine or `Deferred` to protect against cancellation.
 
     Returns:
         A new `Deferred`, which will contain the result of the original coroutine or
@@ -906,10 +893,6 @@ def delay_cancellation(awaitable: Awaitable[T]) -> Awaitable[T]:
 
         When cancelled, the new `Deferred` will wait until the original coroutine or
         `Deferred` resolves before failing with a `CancelledError`.
-
-        The new `Deferred` will follow the Synapse logcontext rules if `awaitable`
-        follows the Synapse logcontext rules. Otherwise the new `Deferred` should be
-        wrapped with `make_deferred_yieldable`.
     """
 
     # First, convert the awaitable into a `Deferred`.

synapse/util/caches/deferred_cache.py

@@ -295,9 +295,6 @@ def set(
         *original* `value`, (c) any future calls to `get()` will complete with the
         result from the *new* `value`.
 
-        It is expected that `value` does *not* follow the synapse logcontext rules - ie,
-        if it is incomplete, it runs its callbacks in the sentinel context.
-
         Args:
             key: Key to be set
             value: a deferred which will complete with a result to add to the cache