Fixup logcontexts after replication PR.

Fixes logcontext leaks introduced in #19138.

Author: erikjohnston

synapse/util/async_helpers.py

@@ -1029,7 +1029,8 @@ def __init__(self, clock: Clock) -> None:
 
     def set(self) -> None:
         if not self._deferred.called:
-            self._deferred.callback(None)
+            with PreserveLoggingContext():
+                self._deferred.callback(None)
 
     def clear(self) -> None:
         if self._deferred.called:
@@ -1042,26 +1043,15 @@ async def wait(self, timeout_seconds: float) -> bool:
         if self.is_set():
             return True
 
-        # Create a deferred that gets called in N seconds
-        sleep_deferred: "defer.Deferred[None]" = defer.Deferred()
-        call = self._clock.call_later(
-            timeout_seconds,
-            sleep_deferred.callback,
-            None,
-        )
-
         try:
             await make_deferred_yieldable(
-                defer.DeferredList(
-                    [sleep_deferred, self._deferred],
-                    fireOnOneCallback=True,
-                    fireOnOneErrback=True,
-                    consumeErrors=True,
+                timeout_deferred(
+                    deferred=stop_cancellation(self._deferred),
+                    timeout=timeout_seconds,
+                    clock=self._clock,
                 )
             )
-        finally:
-            # Cancel the sleep if we were woken up
-            if call.active():
-                call.cancel()
+        except defer.TimeoutError:
+            pass
 
         return self.is_set()

synapse/util/background_queue.py

@@ -25,6 +25,8 @@
     TypeVar,
 )
 
+from twisted.internet import defer
+
 from synapse.util.async_helpers import DeferredEvent
 from synapse.util.constants import MILLISECONDS_PER_SECOND
 
@@ -110,6 +112,8 @@ async def _process_queue(self) -> None:
                     item = self._queue.popleft()
                     try:
                         await self._callback(item)
+                    except defer.CancelledError:
+                        raise
                     except Exception:
                         logger.exception("Error processing background queue item")
 

tests/util/test_background_queue.py

@@ -18,11 +18,13 @@
 from twisted.internet.defer import Deferred
 from twisted.internet.testing import MemoryReactor
 
+from synapse.logging.context import PreserveLoggingContext, make_deferred_yieldable
 from synapse.server import HomeServer
+from synapse.util.async_helpers import DoneAwaitable
 from synapse.util.background_queue import BackgroundQueue
 from synapse.util.clock import Clock
 
-from tests.unittest import HomeserverTestCase
+from tests.unittest import HomeserverTestCase, logcontext_clean
 
 
 class BackgroundQueueTests(HomeserverTestCase):
@@ -38,11 +40,14 @@ def prepare(
             timeout_ms=1000,
         )
 
+    @logcontext_clean
     def test_simple_call(self) -> None:
         """Test that items added to the queue are processed."""
         # Register a deferred to be the return value of the callback.
         callback_result_deferred: Deferred[None] = Deferred()
-        self._process_item_mock.side_effect = callback_result_deferred
+        self._process_item_mock.side_effect = lambda _: make_deferred_yieldable(
+            callback_result_deferred
+        )
 
         # Adding an item should cause the callback to be invoked.
         self.queue.add(1)
@@ -57,16 +62,20 @@ def test_simple_call(self) -> None:
 
         # Once the first callback completes, the second item should be
         # processed.
-        callback_result_deferred.callback(None)
+        with PreserveLoggingContext():
+            callback_result_deferred.callback(None)
         self._process_item_mock.assert_called_once_with(2)
 
+    @logcontext_clean
     def test_timeout(self) -> None:
         """Test that the background process wakes up if its idle, and that it
         times out after being idle."""
 
         # Register a deferred to be the return value of the callback.
         callback_result_deferred: Deferred[None] = Deferred()
-        self._process_item_mock.side_effect = callback_result_deferred
+        self._process_item_mock.side_effect = lambda _: make_deferred_yieldable(
+            callback_result_deferred
+        )
 
         # Adding an item should cause the callback to be invoked.
         self.queue.add(1)
@@ -75,7 +84,8 @@ def test_timeout(self) -> None:
         self._process_item_mock.reset_mock()
 
         # Let the callback complete.
-        callback_result_deferred.callback(None)
+        with PreserveLoggingContext():
+            callback_result_deferred.callback(None)
 
         # Advance the clock by less than the timeout, and add another item.
         self.reactor.advance(0.5)
@@ -84,12 +94,15 @@ def test_timeout(self) -> None:
 
         # The callback should be invoked again.
         callback_result_deferred = Deferred()
-        self._process_item_mock.side_effect = callback_result_deferred
+        self._process_item_mock.side_effect = lambda _: make_deferred_yieldable(
+            callback_result_deferred
+        )
         self._process_item_mock.assert_called_once_with(2)
         self._process_item_mock.reset_mock()
 
         # Let the callback complete.
-        callback_result_deferred.callback(None)
+        with PreserveLoggingContext():
+            callback_result_deferred.callback(None)
 
         # Advance the clock by more than the timeout.
         self.reactor.advance(1.5)