From 83513b75f7b7d63d6ce2e1565568d8051ba658f8 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erikj@element.io>
Date: Wed, 30 Oct 2024 10:51:04 +0000
Subject: Speed up sliding sync by computing extensions in parallel (#17884)

The main change here is to add a helper function
`gather_optional_coroutines`, which works in a similar way as
`yieldable_gather_results` but takes a set of coroutines rather than a
function
---
 synapse/logging/context.py | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

(limited to 'synapse/logging/context.py')

diff --git a/synapse/logging/context.py b/synapse/logging/context.py
index ae2b3d11c0..8a2dfeba13 100644
--- a/synapse/logging/context.py
+++ b/synapse/logging/context.py
@@ -37,6 +37,7 @@ import warnings
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    Any,
     Awaitable,
     Callable,
     Optional,
@@ -850,6 +851,45 @@ def run_in_background(
     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 run until called, and so calling an async function without awaiting
+    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
+
+
 T = TypeVar("T")
 
 
-- 
cgit 1.5.1