matrix-synapse 1.145.0rc2__cp310-abi3-manylinux_2_28_x86_64.whl

synapse/logging/context.py

@@ -0,0 +1,1209 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2019 The Matrix.org Foundation C.I.C.
+# Copyright 2014-2016 OpenMarket Ltd
+# Copyright (C) 2023 New Vector, Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# See the GNU Affero General Public License for more details:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+#
+# Originally licensed under the Apache License, Version 2.0:
+# <http://www.apache.org/licenses/LICENSE-2.0>.
+#
+# [This file includes modifications made by New Vector Limited]
+#
+#
+
+"""Thread-local-alike tracking of log contexts within synapse
+
+This module provides objects and utilities for tracking contexts through
+synapse code, so that log lines can include a request identifier, and so that
+CPU and database activity can be accounted for against the request that caused
+them.
+
+See doc/log_contexts.rst for details on how this works.
+"""
+
+import logging
+import threading
+import typing
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Literal,
+    TypeVar,
+    Union,
+    overload,
+)
+
+import attr
+from typing_extensions import ParamSpec
+
+from twisted.internet import defer, threads
+from twisted.python.threadpool import ThreadPool
+
+from synapse.logging.loggers import ExplicitlyConfiguredLogger
+from synapse.util.stringutils import random_string_insecure_fast
+
+if TYPE_CHECKING:
+    from synapse.logging.scopecontextmanager import _LogContextScope
+    from synapse.types import ISynapseReactor
+
+logger = logging.getLogger(__name__)
+
+original_logger_class = logging.getLoggerClass()
+logging.setLoggerClass(ExplicitlyConfiguredLogger)
+logcontext_debug_logger = logging.getLogger("synapse.logging.context.debug")
+"""
+A logger for debugging when the logcontext switches.
+
+Because this is very noisy and probably something only developers want to see when
+debugging logcontext problems, we want people to explictly opt-in before seeing anything
+in the logs. Requires explicitly setting `synapse.logging.context.debug` in the logging
+configuration and does not inherit the log level from the parent logger.
+"""
+# Restore the original logger class
+logging.setLoggerClass(original_logger_class)
+
+try:
+    import resource
+
+    # Python doesn't ship with a definition of RUSAGE_THREAD but it's defined
+    # to be 1 on linux so we hard code it.
+    RUSAGE_THREAD = 1
+
+    # If the system doesn't support RUSAGE_THREAD then this should throw an
+    # exception.
+    resource.getrusage(RUSAGE_THREAD)
+
+    is_thread_resource_usage_supported = True
+
+    def get_thread_resource_usage() -> "resource.struct_rusage | None":
+        return resource.getrusage(RUSAGE_THREAD)
+
+except Exception:
+    # If the system doesn't support resource.getrusage(RUSAGE_THREAD) then we
+    # won't track resource usage.
+    is_thread_resource_usage_supported = False
+
+    def get_thread_resource_usage() -> "resource.struct_rusage | None":
+        return None
+
+
+# a hook which can be set during testing to assert that we aren't abusing logcontexts.
+def logcontext_error(msg: str) -> None:
+    logger.warning(msg)
+
+
+# get an id for the current thread.
+#
+# threading.get_ident doesn't actually return an OS-level tid, and annoyingly,
+# on Linux it actually returns the same value either side of a fork() call. However
+# we only fork in one place, so it's not worth the hoop-jumping to get a real tid.
+#
+get_thread_id = threading.get_ident
+
+
+class ContextResourceUsage:
+    """Object for tracking the resources used by a log context
+
+    Attributes:
+        ru_utime (float): user CPU time (in seconds)
+        ru_stime (float): system CPU time (in seconds)
+        db_txn_count (int): number of database transactions done
+        db_sched_duration_sec (float): amount of time spent waiting for a
+            database connection
+        db_txn_duration_sec (float): amount of time spent doing database
+            transactions (excluding scheduling time)
+        evt_db_fetch_count (int): number of events requested from the database
+    """
+
+    __slots__ = [
+        "ru_stime",
+        "ru_utime",
+        "db_txn_count",
+        "db_txn_duration_sec",
+        "db_sched_duration_sec",
+        "evt_db_fetch_count",
+    ]
+
+    def __init__(self, copy_from: "ContextResourceUsage | None" = None) -> None:
+        """Create a new ContextResourceUsage
+
+        Args:
+            copy_from: if not None, an object to copy stats from
+        """
+        if copy_from is None:
+            self.reset()
+        else:
+            # FIXME: mypy can't infer the types set via reset() above, so specify explicitly for now
+            self.ru_utime: float = copy_from.ru_utime
+            self.ru_stime: float = copy_from.ru_stime
+            self.db_txn_count: int = copy_from.db_txn_count
+
+            self.db_txn_duration_sec: float = copy_from.db_txn_duration_sec
+            self.db_sched_duration_sec: float = copy_from.db_sched_duration_sec
+            self.evt_db_fetch_count: int = copy_from.evt_db_fetch_count
+
+    def copy(self) -> "ContextResourceUsage":
+        return ContextResourceUsage(copy_from=self)
+
+    def reset(self) -> None:
+        self.ru_stime = 0.0
+        self.ru_utime = 0.0
+        self.db_txn_count = 0
+
+        self.db_txn_duration_sec = 0.0
+        self.db_sched_duration_sec = 0.0
+        self.evt_db_fetch_count = 0
+
+    def __repr__(self) -> str:
+        return (
+            "<ContextResourceUsage ru_stime='%r', ru_utime='%r', "
+            "db_txn_count='%r', db_txn_duration_sec='%r', "
+            "db_sched_duration_sec='%r', evt_db_fetch_count='%r'>"
+        ) % (
+            self.ru_stime,
+            self.ru_utime,
+            self.db_txn_count,
+            self.db_txn_duration_sec,
+            self.db_sched_duration_sec,
+            self.evt_db_fetch_count,
+        )
+
+    def __iadd__(self, other: "ContextResourceUsage") -> "ContextResourceUsage":
+        """Add another ContextResourceUsage's stats to this one's.
+
+        Args:
+            other: the other resource usage object
+        """
+        self.ru_utime += other.ru_utime
+        self.ru_stime += other.ru_stime
+        self.db_txn_count += other.db_txn_count
+        self.db_txn_duration_sec += other.db_txn_duration_sec
+        self.db_sched_duration_sec += other.db_sched_duration_sec
+        self.evt_db_fetch_count += other.evt_db_fetch_count
+        return self
+
+    def __isub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage":
+        self.ru_utime -= other.ru_utime
+        self.ru_stime -= other.ru_stime
+        self.db_txn_count -= other.db_txn_count
+        self.db_txn_duration_sec -= other.db_txn_duration_sec
+        self.db_sched_duration_sec -= other.db_sched_duration_sec
+        self.evt_db_fetch_count -= other.evt_db_fetch_count
+        return self
+
+    def __add__(self, other: "ContextResourceUsage") -> "ContextResourceUsage":
+        res = ContextResourceUsage(copy_from=self)
+        res += other
+        return res
+
+    def __sub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage":
+        res = ContextResourceUsage(copy_from=self)
+        res -= other
+        return res
+
+
+@attr.s(slots=True, auto_attribs=True)
+class ContextRequest:
+    """
+    A bundle of attributes from the SynapseRequest object.
+
+    This exists to:
+
+    * Avoid a cycle between LoggingContext and SynapseRequest.
+    * Be a single variable that can be passed from parent LoggingContexts to
+      their children.
+    """
+
+    request_id: str
+    ip_address: str
+    site_tag: str
+    requester: str | None
+    authenticated_entity: str | None
+    method: str
+    url: str
+    protocol: str
+    user_agent: str
+
+
+LoggingContextOrSentinel = Union["LoggingContext", "_Sentinel"]
+
+
+class _Sentinel:
+    """
+    Sentinel to represent the root context
+
+    This should only be used for tasks outside of Synapse like when we yield control
+    back to the Twisted reactor (event loop) so we don't leak the current logging
+    context to other tasks that are scheduled next in the event loop.
+
+    Nothing from the Synapse homeserver should be logged with the sentinel context. i.e.
+    we should always know which server the logs are coming from.
+    """
+
+    __slots__ = [
+        "previous_context",
+        "finished",
+        "scope",
+        "server_name",
+        "request",
+        "tag",
+    ]
+
+    def __init__(self) -> None:
+        # Minimal set for compatibility with LoggingContext
+        self.previous_context = None
+        self.finished = False
+        self.server_name = "unknown_server_from_sentinel_context"
+        self.request = None
+        self.scope = None
+        self.tag = None
+
+    def __str__(self) -> str:
+        return "sentinel"
+
+    def start(self, rusage: "resource.struct_rusage | None") -> None:
+        pass
+
+    def stop(self, rusage: "resource.struct_rusage | None") -> None:
+        pass
+
+    def add_database_transaction(self, duration_sec: float) -> None:
+        pass
+
+    def add_database_scheduled(self, sched_sec: float) -> None:
+        pass
+
+    def record_event_fetch(self, event_count: int) -> None:
+        pass
+
+    def __bool__(self) -> Literal[False]:
+        return False
+
+
+SENTINEL_CONTEXT = _Sentinel()
+
+
+class LoggingContext:
+    """Additional context for log formatting. Contexts are scoped within a
+    "with" block.
+
+    If a parent is given when creating a new context, then:
+        - logging fields are copied from the parent to the new context on entry
+        - when the new context exits, the cpu usage stats are copied from the
+          child to the parent
+
+    Args:
+        name: Name for the context for logging.
+        server_name: The name of the server this context is associated with
+            (`config.server.server_name` or `hs.hostname`)
+        parent_context (LoggingContext|None): The parent of the new context
+        request: Synapse Request Context object. Useful to associate all the logs
+            happening to a given request.
+
+    """
+
+    __slots__ = [
+        "previous_context",
+        "name",
+        "server_name",
+        "parent_context",
+        "_resource_usage",
+        "usage_start",
+        "main_thread",
+        "finished",
+        "request",
+        "tag",
+        "scope",
+    ]
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        server_name: str,
+        parent_context: "LoggingContext | None" = None,
+        request: ContextRequest | None = None,
+    ) -> None:
+        self.previous_context = current_context()
+
+        # track the resources used by this context so far
+        self._resource_usage = ContextResourceUsage()
+
+        # The thread resource usage when the logcontext became active. None
+        # if the context is not currently active.
+        self.usage_start: resource.struct_rusage | None = None
+
+        self.name = name
+        self.server_name = server_name
+        self.main_thread = get_thread_id()
+        self.request = None
+        self.tag = ""
+        self.scope: "_LogContextScope" | None = None
+
+        # keep track of whether we have hit the __exit__ block for this context
+        # (suggesting that the the thing that created the context thinks it should
+        # be finished, and that re-activating it would suggest an error).
+        self.finished = False
+
+        self.parent_context = parent_context
+
+        # Inherit some fields from the parent context
+        if self.parent_context is not None:
+            # which request this corresponds to
+            self.request = self.parent_context.request
+
+            # we also track the current scope:
+            self.scope = self.parent_context.scope
+
+        if request is not None:
+            # the request param overrides the request from the parent context
+            self.request = request
+
+    def __str__(self) -> str:
+        return self.name
+
+    def __enter__(self) -> "LoggingContext":
+        """Enters this logging context into thread local storage"""
+        logcontext_debug_logger.debug("LoggingContext(%s).__enter__", self.name)
+        old_context = set_current_context(self)
+        if self.previous_context != old_context:
+            logcontext_error(
+                "Expected previous context %r, found %r"
+                % (
+                    self.previous_context,
+                    old_context,
+                )
+            )
+        return self
+
+    def __exit__(
+        self,
+        type: type[BaseException] | None,
+        value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        """Restore the logging context in thread local storage to the state it
+        was before this context was entered.
+        Returns:
+            None to avoid suppressing any exceptions that were thrown.
+        """
+        logcontext_debug_logger.debug(
+            "LoggingContext(%s).__exit__ --> %s", self.name, self.previous_context
+        )
+        current = set_current_context(self.previous_context)
+        if current is not self:
+            if current is SENTINEL_CONTEXT:
+                logcontext_error("Expected logging context %s was lost" % (self,))
+            else:
+                logcontext_error(
+                    "Expected logging context %s but found %s" % (self, current)
+                )
+
+        # the fact that we are here suggests that the caller thinks that everything
+        # is done and dusted for this logcontext, and further activity will not get
+        # recorded against the correct metrics.
+        self.finished = True
+
+    def start(self, rusage: "resource.struct_rusage | None") -> None:
+        """
+        Record that this logcontext is currently running.
+
+        This should not be called directly: use set_current_context
+
+        Args:
+            rusage: the resources used by the current thread, at the point of
+                switching to this logcontext. May be None if this platform doesn't
+                support getrusuage.
+        """
+        if get_thread_id() != self.main_thread:
+            logcontext_error("Started logcontext %s on different thread" % (self,))
+            return
+
+        if self.finished:
+            logcontext_error("Re-starting finished log context %s" % (self,))
+
+        # If we haven't already started record the thread resource usage so
+        # far
+        if self.usage_start:
+            logcontext_error("Re-starting already-active log context %s" % (self,))
+        else:
+            self.usage_start = rusage
+
+    def stop(self, rusage: "resource.struct_rusage | None") -> None:
+        """
+        Record that this logcontext is no longer running.
+
+        This should not be called directly: use set_current_context
+
+        Args:
+            rusage: the resources used by the current thread, at the point of
+                switching away from this logcontext. May be None if this platform
+                doesn't support getrusuage.
+        """
+
+        try:
+            if get_thread_id() != self.main_thread:
+                logcontext_error("Stopped logcontext %s on different thread" % (self,))
+                return
+
+            if not rusage:
+                return
+
+            # Record the cpu used since we started
+            if not self.usage_start:
+                logcontext_error(
+                    "Called stop on logcontext %s without recording a start rusage"
+                    % (self,)
+                )
+                return
+
+            utime_delta, stime_delta = self._get_cputime(rusage)
+            self.add_cputime(utime_delta, stime_delta)
+        finally:
+            self.usage_start = None
+
+    def get_resource_usage(self) -> ContextResourceUsage:
+        """Get resources used by this logcontext so far.
+
+        Returns:
+            A *copy* of the object tracking resource usage so far
+        """
+        # we always return a copy, for consistency
+        res = self._resource_usage.copy()
+
+        # If we are on the correct thread and we're currently running then we
+        # can include resource usage so far.
+        is_main_thread = get_thread_id() == self.main_thread
+        if self.usage_start and is_main_thread:
+            rusage = get_thread_resource_usage()
+            assert rusage is not None
+            utime_delta, stime_delta = self._get_cputime(rusage)
+            res.ru_utime += utime_delta
+            res.ru_stime += stime_delta
+
+        return res
+
+    def _get_cputime(self, current: "resource.struct_rusage") -> tuple[float, float]:
+        """Get the cpu usage time between start() and the given rusage
+
+        Args:
+            rusage: the current resource usage
+
+        Returns: tuple[float, float]: seconds in user mode, seconds in system mode
+        """
+        assert self.usage_start is not None
+
+        utime_delta = current.ru_utime - self.usage_start.ru_utime
+        stime_delta = current.ru_stime - self.usage_start.ru_stime
+
+        # sanity check
+        if utime_delta < 0:
+            logger.error(
+                "utime went backwards! %f < %f",
+                current.ru_utime,
+                self.usage_start.ru_utime,
+            )
+            utime_delta = 0
+
+        if stime_delta < 0:
+            logger.error(
+                "stime went backwards! %f < %f",
+                current.ru_stime,
+                self.usage_start.ru_stime,
+            )
+            stime_delta = 0
+
+        return utime_delta, stime_delta
+
+    def add_cputime(self, utime_delta: float, stime_delta: float) -> None:
+        """Update the CPU time usage of this context (and any parents, recursively).
+
+        Args:
+            utime_delta: additional user time, in seconds, spent in this context.
+            stime_delta: additional system time, in seconds, spent in this context.
+        """
+        self._resource_usage.ru_utime += utime_delta
+        self._resource_usage.ru_stime += stime_delta
+        if self.parent_context:
+            self.parent_context.add_cputime(utime_delta, stime_delta)
+
+    def add_database_transaction(self, duration_sec: float) -> None:
+        """Record the use of a database transaction and the length of time it took.
+
+        Args:
+            duration_sec: The number of seconds the database transaction took.
+        """
+        if duration_sec < 0:
+            raise ValueError("DB txn time can only be non-negative")
+        self._resource_usage.db_txn_count += 1
+        self._resource_usage.db_txn_duration_sec += duration_sec
+        if self.parent_context:
+            self.parent_context.add_database_transaction(duration_sec)
+
+    def add_database_scheduled(self, sched_sec: float) -> None:
+        """Record a use of the database pool
+
+        Args:
+            sched_sec: number of seconds it took us to get a connection
+        """
+        if sched_sec < 0:
+            raise ValueError("DB scheduling time can only be non-negative")
+        self._resource_usage.db_sched_duration_sec += sched_sec
+        if self.parent_context:
+            self.parent_context.add_database_scheduled(sched_sec)
+
+    def record_event_fetch(self, event_count: int) -> None:
+        """Record a number of events being fetched from the db
+
+        Args:
+            event_count: number of events being fetched
+        """
+        self._resource_usage.evt_db_fetch_count += event_count
+        if self.parent_context:
+            self.parent_context.record_event_fetch(event_count)
+
+
+class LoggingContextFilter(logging.Filter):
+    """Logging filter that adds values from the current logging context to each
+    record.
+    """
+
+    def __init__(
+        self,
+        # `request` is here for backwards compatibility since we previously recommended
+        # people manually configure `LoggingContextFilter` like the following.
+        #
+        # ```yaml
+        # filters:
+        #   context:
+        #       (): synapse.logging.context.LoggingContextFilter
+        #       request: ""
+        # ```
+        #
+        # TODO: Since we now configure `LoggingContextFilter` automatically since #8051
+        # (2020-08-11), we could consider removing this useless parameter. This would
+        # require people to remove their own manual configuration of
+        # `LoggingContextFilter` as it would cause `TypeError: Filter.__init__() got an
+        # unexpected keyword argument 'request'` -> `ValueError: Unable to configure
+        # filter 'context'`
+        request: str = "",
+    ):
+        self._default_request = request
+
+    def filter(self, record: logging.LogRecord) -> Literal[True]:
+        """
+        Add each field from the logging context to the record.
+
+        Please be mindful of 3rd-party code outside of Synapse (like in the case of
+        Synapse Pro for small hosts) as this is running as a global log record filter.
+        Other code may have set their own attributes on the record and the log record
+        may not be relevant to Synapse at all so we should not mangle it.
+
+        We can have some defaults but we should avoid overwriting existing attributes on
+        any log record unless we actually have a Synapse logcontext (not just the
+        default sentinel logcontext).
+
+        Returns:
+            True to include the record in the log output.
+        """
+        context = current_context()
+        # type-ignore: `context` should never be `None`, but if it somehow ends up
+        # being, then we end up in a death spiral of infinite loops, so let's check, for
+        # robustness' sake.
+        #
+        # Add some default values to avoid log formatting errors.
+        if context is None:
+            record.request = self._default_request  # type: ignore[unreachable]
+
+            # Avoid overwriting an existing `server_name` on the record. This is running in
+            # the context of a global log record filter so there may be 3rd-party code that
+            # adds their own `server_name` and we don't want to interfere with that
+            # (clobber).
+            if not hasattr(record, "server_name"):
+                record.server_name = "unknown_server_from_no_logcontext"
+
+        # Otherwise, in the normal, expected case, fill in the log record attributes
+        # from the logcontext.
+        else:
+
+            def safe_set(attr: str, value: Any) -> None:
+                """
+                Only write the attribute if it hasn't already been set or we actually have
+                a Synapse logcontext (indicating that this log record is relevant to
+                Synapse).
+                """
+                if context is not SENTINEL_CONTEXT or not hasattr(record, attr):
+                    setattr(record, attr, value)
+
+            safe_set("server_name", context.server_name)
+
+            # Logging is interested in the request ID. Note that for backwards
+            # compatibility this is stored as the "request" on the record.
+            safe_set("request", str(context))
+
+            # Add some data from the HTTP request.
+            request = context.request
+            # The sentinel logcontext has no request so if we get past this point, we
+            # know we have some actual Synapse logcontext and don't need to worry about
+            # using `safe_set`. We'll consider this an optimization since this is a
+            # pretty hot-path.
+            if request is None:
+                return True
+
+            record.ip_address = request.ip_address
+            record.site_tag = request.site_tag
+            record.requester = request.requester
+            record.authenticated_entity = request.authenticated_entity
+            record.method = request.method
+            record.url = request.url
+            record.protocol = request.protocol
+            record.user_agent = request.user_agent
+
+        return True
+
+
+class PreserveLoggingContext:
+    """
+    Context manager which replaces the logging context
+
+    The previous logging context is restored on exit.
+
+    `make_deferred_yieldable` is pretty equivalent to using `with
+    PreserveLoggingContext():` (using the default sentinel context), i.e. it clears the
+    logcontext before awaiting (and so before execution passes back to the reactor) and
+    restores the old context once the awaitable completes (execution passes from the
+    reactor back to the code).
+    """
+
+    __slots__ = ["_old_context", "_new_context", "_instance_id"]
+
+    def __init__(
+        self, new_context: LoggingContextOrSentinel = SENTINEL_CONTEXT
+    ) -> None:
+        self._new_context = new_context
+        self._instance_id = random_string_insecure_fast(5)
+
+    def __enter__(self) -> None:
+        logcontext_debug_logger.debug(
+            "PreserveLoggingContext(%s).__enter__ %s --> %s",
+            self._instance_id,
+            current_context(),
+            self._new_context,
+        )
+        self._old_context = set_current_context(self._new_context)
+
+    def __exit__(
+        self,
+        type: type[BaseException] | None,
+        value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        logcontext_debug_logger.debug(
+            "PreserveLoggingContext(%s).__exit %s --> %s",
+            self._instance_id,
+            current_context(),
+            self._old_context,
+        )
+        context = set_current_context(self._old_context)
+
+        if context != self._new_context:
+            if not context:
+                logcontext_error(
+                    "Expected logging context %s was lost" % (self._new_context,)
+                )
+            else:
+                logcontext_error(
+                    "Expected logging context %s but found %s"
+                    % (
+                        self._new_context,
+                        context,
+                    )
+                )
+
+
+_thread_local = threading.local()
+_thread_local.current_context = SENTINEL_CONTEXT
+
+
+def current_context() -> LoggingContextOrSentinel:
+    """Get the current logging context from thread local storage"""
+    return getattr(_thread_local, "current_context", SENTINEL_CONTEXT)
+
+
+def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSentinel:
+    """Set the current logging context in thread local storage
+    Args:
+        context: The context to activate.
+
+    Returns:
+        The context that was previously active
+    """
+    # everything blows up if we allow current_context to be set to None, so sanity-check
+    # that now.
+    if context is None:
+        raise TypeError("'context' argument may not be None")
+
+    current = current_context()
+
+    if current is not context:
+        rusage = get_thread_resource_usage()
+        current.stop(rusage)
+        _thread_local.current_context = context
+        context.start(rusage)
+
+    return current
+
+
+def nested_logging_context(suffix: str) -> LoggingContext:
+    """Creates a new logging context as a child of another.
+
+    The nested logging context will have a 'name' made up of the parent context's
+    name, plus the given suffix.
+
+    CPU/db usage stats will be added to the parent context's on exit.
+
+    Normal usage looks like:
+
+        with nested_logging_context(suffix):
+            # ... do stuff
+
+    Args:
+        suffix: suffix to add to the parent context's 'name'.
+
+    Returns:
+        A new logging context.
+    """
+    curr_context = current_context()
+    if not curr_context:
+        logger.warning(
+            "Starting nested logging context from sentinel context: metrics will be lost"
+        )
+        parent_context = None
+        server_name = "unknown_server_from_sentinel_context"
+    else:
+        assert isinstance(curr_context, LoggingContext)
+        parent_context = curr_context
+        server_name = parent_context.server_name
+    prefix = str(curr_context)
+    return LoggingContext(
+        name=prefix + "-" + suffix,
+        server_name=server_name,
+        parent_context=parent_context,
+    )
+
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+async def _unwrap_awaitable(awaitable: Awaitable[R]) -> R:
+    """Unwraps an arbitrary awaitable by awaiting it."""
+    return await awaitable
+
+
+@overload
+def preserve_fn(
+    f: Callable[P, Awaitable[R]],
+) -> Callable[P, "defer.Deferred[R]"]:
+    # The `type: ignore[misc]` above suppresses
+    # "Overloaded function signatures 1 and 2 overlap with incompatible return types"
+    ...
+
+
+@overload
+def preserve_fn(f: Callable[P, R]) -> Callable[P, "defer.Deferred[R]"]: ...
+
+
+def preserve_fn(
+    f: Callable[P, R] | Callable[P, Awaitable[R]],
+) -> Callable[P, "defer.Deferred[R]"]:
+    """Function decorator which wraps the function with run_in_background"""
+
+    def g(*args: P.args, **kwargs: P.kwargs) -> "defer.Deferred[R]":
+        return run_in_background(f, *args, **kwargs)
+
+    return g
+
+
+@overload
+def run_in_background(
+    f: Callable[P, Awaitable[R]], *args: P.args, **kwargs: P.kwargs
+) -> "defer.Deferred[R]":
+    # The `type: ignore[misc]` above suppresses
+    # "Overloaded function signatures 1 and 2 overlap with incompatible return types"
+    ...
+
+
+@overload
+def run_in_background(
+    f: Callable[P, R], *args: P.args, **kwargs: P.kwargs
+) -> "defer.Deferred[R]": ...
+
+
+def run_in_background(
+    f: Callable[P, R] | Callable[P, Awaitable[R]],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> "defer.Deferred[R]":
+    """Calls a function, ensuring that the current context is restored after
+    return from the function, and that the sentinel context is set once the
+    deferred returned by the function completes.
+
+    To explain how the log contexts work here:
+     - When `run_in_background` is called, the calling logcontext is stored
+       ("original"), we kick off the background task in the current context, and we
+       restore that original context before returning.
+     - For a completed deferred, that's the end of the story.
+     - For an incomplete deferred, when the background task finishes, we don't want to
+       leak our context into the reactor which would erroneously get attached to the
+       next operation picked up by the event loop. We add a callback to the deferred
+       which will clear the logging context after it finishes and yields control back to
+       the reactor.
+
+    Useful for wrapping functions that return a deferred or coroutine, which you don't
+    yield or await on (for instance because you want to pass it to
+    deferred.gatherResults()).
+
+    If f returns a Coroutine object, it will be wrapped into a Deferred (which will have
+    the side effect of executing the coroutine).
+
+    Note that if you completely discard the result, you should make sure that
+    `f` doesn't raise any deferred exceptions, otherwise a scary-looking
+    CRITICAL error about an unhandled error will be logged without much
+    indication about where it came from.
+
+    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.
+    """
+    instance_id = random_string_insecure_fast(5)
+    calling_context = current_context()
+    logcontext_debug_logger.debug(
+        "run_in_background(%s): called with logcontext=%s", instance_id, calling_context
+    )
+    try:
+        # (kick off the task in the current context)
+        res = f(*args, **kwargs)
+    except Exception:
+        # the assumption here is that the caller doesn't want to be disturbed
+        # by synchronous exceptions, so let's turn them into Failures.
+        return defer.fail()
+
+    # `res` may be a coroutine, `Deferred`, some other kind of awaitable, or a plain
+    # value. Convert it to a `Deferred`.
+    #
+    # Wrapping the value in a deferred has the side effect of executing the coroutine,
+    # if it is one. If it's already a deferred, then we can just use that.
+    d: "defer.Deferred[R]"
+    if isinstance(res, typing.Coroutine):
+        # Wrap the coroutine in a `Deferred`.
+        d = defer.ensureDeferred(res)
+    elif isinstance(res, defer.Deferred):
+        d = res
+    elif isinstance(res, Awaitable):
+        # `res` is probably some kind of completed awaitable, such as a `DoneAwaitable`
+        # or `Future` from `make_awaitable`.
+        d = defer.ensureDeferred(_unwrap_awaitable(res))
+    else:
+        # `res` is a plain value. Wrap it in a `Deferred`.
+        d = defer.succeed(res)
+
+    # The deferred has already completed
+    if d.called and not d.paused:
+        # If the function messes with logcontexts, we can assume it follows the Synapse
+        # logcontext rules (Rules for functions returning awaitables: "If the awaitable
+        # is already complete, the function returns with the same logcontext it started
+        # with."). If the function doesn't touch logcontexts at all, we can also assume
+        # the logcontext is unchanged.
+        #
+        # Either way, the function should have maintained the calling logcontext, so we
+        # can avoid messing with it further. Additionally, if the deferred has already
+        # completed, then it would be a mistake to then add a deferred callback (below)
+        # to reset the logcontext to the sentinel logcontext as that would run
+        # immediately (remember our goal is to maintain the calling logcontext when we
+        # return).
+        if current_context() != calling_context:
+            logcontext_error(
+                "run_in_background(%s): deferred already completed but the function did not maintain the calling logcontext %s (found %s)"
+                % (
+                    instance_id,
+                    calling_context,
+                    current_context(),
+                )
+            )
+        else:
+            logcontext_debug_logger.debug(
+                "run_in_background(%s): deferred already completed (maintained the calling logcontext %s)",
+                instance_id,
+                calling_context,
+            )
+        return d
+
+    # Since the function we called may follow the Synapse logcontext rules (Rules for
+    # functions returning awaitables: "If the awaitable is incomplete, the function
+    # clears the logcontext before returning"), the function may have reset the
+    # logcontext before returning, so we need to restore the calling logcontext now
+    # before we return ourselves.
+    #
+    # Our goal is to have the caller logcontext unchanged after firing off the
+    # background task and returning.
+    logcontext_debug_logger.debug(
+        "run_in_background(%s): restoring calling logcontext %s",
+        instance_id,
+        calling_context,
+    )
+    set_current_context(calling_context)
+
+    # If the function we called is playing nice and following the Synapse logcontext
+    # rules, it will restore original calling logcontext when the deferred completes;
+    # but there is nothing waiting for it, so it will get leaked into the reactor (which
+    # would then get picked up by the next thing the reactor does). We therefore need to
+    # reset the logcontext here (set the `sentinel` logcontext) before yielding control
+    # back to the reactor.
+    #
+    # (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.)
+    if logcontext_debug_logger.isEnabledFor(logging.DEBUG):
+
+        def _log_set_context_cb(
+            result: ResultT, context: LoggingContextOrSentinel
+        ) -> ResultT:
+            logcontext_debug_logger.debug(
+                "run_in_background(%s): resetting logcontext to %s",
+                instance_id,
+                context,
+            )
+            set_current_context(context)
+            return result
+
+        d.addBoth(_log_set_context_cb, SENTINEL_CONTEXT)
+    else:
+        d.addBoth(_set_context_cb, SENTINEL_CONTEXT)
+
+    return d
+
+
+def run_coroutine_in_background(
+    coroutine: typing.Coroutine[Any, Any, R],
+) -> "defer.Deferred[R]":
+    """Run the coroutine, ensuring that the current context is restored after
+    return from the function, and that the sentinel context is set once the
+    deferred returned by the function completes.
+
+    Useful for wrapping coroutines that you don't yield or await on (for
+    instance because you want to pass it to deferred.gatherResults()).
+
+    This is a special case of `run_in_background` where we can accept a coroutine
+    directly rather than a function. We can do this because coroutines do not continue
+    running once they have yielded.
+
+    This is an ergonomic helper so we can do this:
+    ```python
+    run_coroutine_in_background(func1(arg1))
+    ```
+    Rather than having to do this:
+    ```python
+    run_in_background(lambda: func1(arg1))
+    ```
+    """
+    return run_in_background(lambda: coroutine)
+
+
+T = TypeVar("T")
+
+
+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
+      incomplete deferred. Then, when the deferred completes, restores the current
+      logcontext before running callbacks/errbacks.
+
+    This means the resultant deferred can be awaited without leaking the current
+    logcontext to the reactor (which would then get erroneously picked up by the next
+    thing the reactor does), and also means that the logcontext is preserved when the
+    deferred completes.
+
+    (This is more-or-less the opposite operation to run_in_background in terms of how it
+    handles log contexts.)
+
+    Pretty much equivalent to using `with PreserveLoggingContext():`, i.e. it clears the
+    logcontext before awaiting (and so before execution passes back to the reactor) and
+    restores the old context once the awaitable completes (execution passes from the
+    reactor back to the code).
+    """
+    instance_id = random_string_insecure_fast(5)
+    logcontext_debug_logger.debug(
+        "make_deferred_yieldable(%s): called with logcontext=%s",
+        instance_id,
+        current_context(),
+    )
+
+    # The deferred has already completed
+    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.
+        logcontext_debug_logger.debug(
+            "make_deferred_yieldable(%s): deferred already completed and the function should have maintained the logcontext",
+            instance_id,
+        )
+        return deferred
+
+    # Our goal is to have the caller logcontext unchanged after they yield/await the
+    # returned deferred.
+    #
+    # When the caller yield/await's the returned deferred, it may yield
+    # control back to the reactor. To avoid leaking the current logcontext to the
+    # reactor (which would then get erroneously picked up by the next thing the reactor
+    # does) while the deferred runs in the reactor event loop, we reset the logcontext
+    # and add a callback to the deferred to restore it so the caller's logcontext is
+    # active when the deferred completes.
+
+    logcontext_debug_logger.debug(
+        "make_deferred_yieldable(%s): resetting logcontext to %s",
+        instance_id,
+        SENTINEL_CONTEXT,
+    )
+    calling_context = set_current_context(SENTINEL_CONTEXT)
+
+    if logcontext_debug_logger.isEnabledFor(logging.DEBUG):
+
+        def _log_set_context_cb(
+            result: ResultT, context: LoggingContextOrSentinel
+        ) -> ResultT:
+            logcontext_debug_logger.debug(
+                "make_deferred_yieldable(%s): restoring calling logcontext to %s",
+                instance_id,
+                context,
+            )
+            set_current_context(context)
+            return result
+
+        deferred.addBoth(_log_set_context_cb, calling_context)
+    else:
+        deferred.addBoth(_set_context_cb, calling_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]":
+    """
+    Calls the function `f` using a thread from the reactor's default threadpool and
+    returns the result as a Deferred.
+
+    Creates a new logcontext for `f`, which is created as a child of the current
+    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.
+
+            Normally this will be hs.get_reactor().
+
+        f: The function to call.
+
+        args: positional arguments to pass to f.
+
+        kwargs: keyword arguments to pass to f.
+
+    Returns:
+        A Deferred which fires a callback with the result of `f`, or an
+            errback if `f` throws an exception.
+    """
+    return defer_to_threadpool(reactor, reactor.getThreadPool(), f, *args, **kwargs)
+
+
+def defer_to_threadpool(
+    reactor: "ISynapseReactor",
+    threadpool: ThreadPool,
+    f: Callable[P, R],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> "defer.Deferred[R]":
+    """
+    A wrapper for twisted.internet.threads.deferToThreadpool, which handles
+    logcontexts correctly.
+
+    Calls the function `f` using a thread from the given threadpool and returns
+    the result as a Deferred.
+
+    Creates a new logcontext for `f`, which is created as a child of the current
+    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().
+
+        threadpool: The threadpool to use for running `f`. Normally this will be
+            hs.get_reactor().getThreadPool().
+
+        f: The function to call.
+
+        args: positional arguments to pass to f.
+
+        kwargs: keyword arguments to pass to f.
+
+    Returns:
+        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
+        server_name = "unknown_server_from_sentinel_context"
+    else:
+        assert isinstance(curr_context, LoggingContext)
+        parent_context = curr_context
+        server_name = parent_context.server_name
+
+    def g() -> R:
+        with LoggingContext(
+            name=str(curr_context),
+            server_name=server_name,
+            parent_context=parent_context,
+        ):
+            return f(*args, **kwargs)
+
+    return make_deferred_yieldable(threads.deferToThreadPool(reactor, threadpool, g))