Faster redis replication handling

changelog.d/19138.misc

@@ -0,0 +1 @@
+Minor speed up of processing of inbound replication.

synapse/replication/tcp/handler.py

@@ -20,7 +20,6 @@
 #
 #
 import logging
-from collections import deque
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -71,6 +70,7 @@
     DeviceListsStream,
     ThreadSubscriptionsStream,
 )
+from synapse.util.async_helpers import BackgroundQueue
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -115,8 +115,8 @@
 
 
 # the type of the entries in _command_queues_by_stream
-_StreamCommandQueue = deque[
-    tuple[Union[RdataCommand, PositionCommand], IReplicationConnection]
+_StreamCommandQueueItem = tuple[
+    Union[RdataCommand, PositionCommand], IReplicationConnection
 ]
 
 
@@ -265,7 +265,12 @@ def __init__(self, hs: "HomeServer"):
         # for each stream, a queue of commands that are awaiting processing, and the
         # connection that they arrived on.
         self._command_queues_by_stream = {
-            stream_name: _StreamCommandQueue() for stream_name in self._streams
+            stream_name: BackgroundQueue[_StreamCommandQueueItem](
+                hs,
+                "process-replication-data",
+                self._unsafe_process,
+            )
+            for stream_name in self._streams
         }
 
         # For each connection, the incoming stream names that have received a POSITION
@@ -349,38 +354,13 @@ def _add_command_to_stream_queue(
             logger.error("Got %s for unknown stream: %s", cmd.NAME, stream_name)
             return
 
-        queue.append((cmd, conn))
-
-        # if we're already processing this stream, there's nothing more to do:
-        # the new entry on the queue will get picked up in due course
-        if stream_name in self._processing_streams:
-            return
-
-        # fire off a background process to start processing the queue.
-        self.hs.run_as_background_process(
-            "process-replication-data",
-            self._unsafe_process_queue,
-            stream_name,
-        )
-
-    async def _unsafe_process_queue(self, stream_name: str) -> None:
-        """Processes the command queue for the given stream, until it is empty
+        queue.add((cmd, conn))
 
-        Does not check if there is already a thread processing the queue, hence "unsafe"
-        """
-        assert stream_name not in self._processing_streams
-
-        self._processing_streams.add(stream_name)
-        try:
-            queue = self._command_queues_by_stream.get(stream_name)
-            while queue:
-                cmd, conn = queue.popleft()
-                try:
-                    await self._process_command(cmd, conn, stream_name)
-                except Exception:
-                    logger.exception("Failed to handle command %s", cmd)
-        finally:
-            self._processing_streams.discard(stream_name)
+    async def _unsafe_process(self, item: _StreamCommandQueueItem) -> None:
+        """Process a single command from the stream queue"""
+        cmd, conn = item
+        stream_name = cmd.stream_name
+        await self._process_command(cmd, conn, stream_name)
 
     async def _process_command(
         self,

synapse/util/async_helpers.py

@@ -28,6 +28,7 @@
 import typing
 from contextlib import asynccontextmanager
 from typing import (
+    TYPE_CHECKING,
     Any,
     AsyncContextManager,
     AsyncIterator,
@@ -61,6 +62,9 @@
 )
 from synapse.util.clock import Clock
 
+if TYPE_CHECKING:
+    from synapse.server import HomeServer
+
 logger = logging.getLogger(__name__)
 
 _T = TypeVar("_T")
@@ -1065,3 +1069,94 @@ async def wait(self, timeout_seconds: float) -> bool:
                 call.cancel()
 
         return self.is_set()
+
+
+class BackgroundQueue(Generic[T]):
+    """A single-producer single-consumer async queue processing items in the
+    background.
+
+    This is optimised for the case where we receive many items, but processing
+    each one takes a short amount of time. In this case we don't want to pay the
+    overhead of a new background process each time. Instead, we spawn a
+    background process that will wait for new items to arrive.
+
+    If the background process has been idle for a while, it will exit, and a new
+    background process will be spawned when new items arrive.
+
+    Args:
+        hs: The homeserver.
+        name: The name of the background process.
+        callback: The async callback to process each item.
+        timeout_ms: The time in milliseconds to wait for new items before
+            exiting the background process.
+    """
+
+    def __init__(
+        self,
+        hs: "HomeServer",
+        name: str,
+        callback: Callable[[T], Awaitable[None]],
+        timeout_ms: int = 1000,
+    ) -> None:
+        self._hs = hs
+        self._name = name
+        self._callback = callback
+        self._timeout_ms = timeout_ms
+
+        # The queue of items to process.
+        self._queue: collections.deque[T] = collections.deque()
+
+        # Indicates if a background process is running, and if so whether there
+        # is new data in the queue.
+        self._event: Optional[DeferredEvent] = None
+
+    def add(self, item: T) -> None:
+        """Add an item into the queue."""
+
+        self._queue.append(item)
+        if self._event is None:
+            self._hs.run_as_background_process(self._name, self._process_queue)
+        else:
+            self._event.set()
+
+    async def _process_queue(self) -> None:
+        """Process items in the queue until it is empty."""
+
+        # Make sure we're the only background process.
+        if self._event is not None:
+            self._event.set()
+            return
+
+        self._event = DeferredEvent(self._hs.get_clock())
+
+        try:
+            while True:
+                # Clear the event before checking the queue.
+                self._event.clear()
+
+                while self._queue:
+                    item = self._queue.popleft()
+                    try:
+                        await self._callback(item)
+                    except Exception:
+                        logger.exception("Error processing background queue item")
+
+                # Wait for new data to arrive, timing out after a while to avoid
+                # keeping the background process alive forever.
+                #
+                # New data may have arrived and been processed while we were
+                # pulling from the queue, so this may return that there is new
+                # data immediately even though there isn't. That's fine, we'll
+                # just loop round, clear the event, recheck the queue, and then
+                # wait here again.
+                new_data = await self._event.wait(
+                    timeout_seconds=self._timeout_ms / 1000
+                )
+                if not new_data:
+                    # Timed out waiting for new data, so exit the loop
+                    break
+        finally:
+            self._event = None
+
+    def __len__(self) -> int:
+        return len(self._queue)