fix: presence stream stalling intermittently (#20090)

This is a fix for presence updates silently stalling when a `/sync`
request is cancelled mid-write, causing a stream ID to be leaked into
`_unfinished_ids` and permanently pinning the persisted stream position.

Fixes https://github.com/element-hq/synapse/issues/19800

### Pull Request Checklist

<!-- Please read
https://element-hq.github.io/synapse/latest/development/contributing_guide.html
before submitting your pull request -->

* [x] Pull request is based on the develop branch
* [x] Pull request includes a [changelog
file](https://element-hq.github.io/synapse/latest/development/contributing_guide.html#changelog).
The entry should:
- Be a short description of your change which makes sense to users.
"Fixed a bug that prevented receiving messages from other servers."
instead of "Moved X method from `EventStore` to `EventWorkerStore`.".
  - Use markdown where necessary, mostly for `code blocks`.
  - End with either a period (.) or an exclamation mark (!).
  - Start with a capital letter.
- Feel free to credit yourself, by adding a sentence "Contributed by
@github_username." or "Contributed by [Your Name]." to the end of the
entry.
* [x] [Code
style](https://element-hq.github.io/synapse/latest/code_style.html) is
correct (run the
[linters](https://element-hq.github.io/synapse/latest/development/contributing_guide.html#run-the-linters))

changelog.d/20090.bugfix

@@ -0,0 +1 @@
+Fix a bug where presence updates could stop being sent to clients (the presence stream position becoming stuck) if a `/sync` request was cancelled while a presence write was allocating a stream ID. Contributed by @FrenchGithubUser @Famedly.

synapse/storage/util/id_generators.py

@@ -906,14 +906,44 @@ class _MultiWriterCtxManager:
     stream_ids: list[int] = attr.Factory(list)
 
     async def __aenter__(self) -> int | list[int]:
+        def _load(txn: LoggingTransaction) -> list[int]:
+            ids = self.id_gen._load_next_mult_id_txn(txn, self.multiple_ids or 1)
+            # Record the allocated IDs on the context manager as a side effect
+            # (rather than only via the return value), so that if this coroutine
+            # is cancelled after the transaction has committed we still know
+            # which IDs to release below.
+            self.stream_ids = ids
+            return ids
+
         # It's safe to run this in autocommit mode as fetching values from a
         # sequence ignores transaction semantics anyway.
-        self.stream_ids = await self.id_gen._db.runInteraction(
-            "_load_next_mult_id",
-            self.id_gen._load_next_mult_id_txn,
-            self.multiple_ids or 1,
-            db_autocommit=True,
-        )
+        try:
+            await self.id_gen._db.runInteraction(
+                "_load_next_mult_id",
+                _load,
+                db_autocommit=True,
+            )
+        except BaseException:
+            # We catch `BaseException` rather than `Exception`,
+            # because request cancellation surfaces here as exceptions that are
+            # not `Exception` subclasses: `asyncio.CancelledError`
+            # and `GeneratorExit` (raised when a paused coroutine is garbage
+            # collected).
+            #
+            # If we're interrupted (e.g. the enclosing request was cancelled)
+            # after the transaction allocated the IDs but before we returned,
+            # then `__aexit__` will never run, because Python only invokes it
+            # once `__aenter__` has returned. The allocated IDs would then be
+            # leaked into `_unfinished_ids` forever, permanently pinning the
+            # persisted stream position and, e.g., wedging presence.
+            #
+            # So mark them as finished here to unblock the position. This mirrors
+            # what `__aexit__` does on the failure path (marking the IDs finished
+            # and notifying replication, but not persisting a new position).
+            if self.stream_ids:
+                self.id_gen._mark_ids_as_finished(self.stream_ids)
+                self.notifier.notify_replication()
+            raise
 
         if self.multiple_ids is None:
             return self.stream_ids[0] * self.id_gen._return_factor

tests/storage/test_id_generators.py

@@ -19,8 +19,12 @@
 #
 #
 
+from unittest import mock
+
+from twisted.internet.defer import CancelledError, Deferred, ensureDeferred
 from twisted.internet.testing import MemoryReactor
 
+from synapse.logging.context import LoggingContext, make_deferred_yieldable
 from synapse.server import HomeServer
 from synapse.storage.database import (
     DatabasePool,
@@ -225,6 +229,90 @@ class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase):
         self.assertEqual(id_gen.get_positions(), {"master": 8})
         self.assertEqual(id_gen.get_current_token_for_writer("master"), 8)
 
+    def test_cancelled_enter_does_not_wedge_position(self) -> None:
+        """Reproduces presence getting stuck.
+
+        If the `get_next()` async context manager is cancelled while
+        `__aenter__` is allocating a stream ID, the DB interaction that runs the
+        sequence has already added the ID to `_unfinished_ids`, but `__aexit__`
+        is never called (Python only invokes `__aexit__` if `__aenter__`
+        returned). The abandoned ID is therefore leaked into `_unfinished_ids`
+        forever, which permanently pins the persisted stream position: new rows
+        keep getting higher IDs, but `get_current_token()` can never advance past
+        `leaked_id - 1` until the process restarts.
+
+        This mirrors a `/sync` request being cancelled part-way through
+        persisting a presence update. `/sync` became `@cancellable` in #19499,
+        and on a monolith the presence write in `PresenceStore.update_presence`
+        is awaited inside that cancellable request scope.
+        """
+        # Prefill table with 7 rows written by 'master'; position starts at 7.
+        self._insert_rows("master", 7)
+
+        id_gen = self._create_id_generator()
+        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
+
+        # We model the cancellation at the seam it actually happens in
+        # production: `__aenter__` awaits `runInteraction("_load_next_mult_id")`,
+        # whose transaction runs in a thread pool and so *always* completes -
+        # allocating stream ID 8 and adding it to `_unfinished_ids` - but the
+        # awaiting coroutine is handed a `CancelledError` because the enclosing
+        # `/sync` request was cancelled. We reproduce that by letting the real
+        # interaction run (applying its side effects) and then failing the
+        # awaited deferred with `CancelledError`.
+        cancel_enter: "Deferred[None]" = Deferred()
+        original_run_interaction = id_gen._db.runInteraction
+
+        async def blocking_run_interaction(desc, func, *args, **kwargs):  # type: ignore[no-untyped-def]
+            result = await original_run_interaction(desc, func, *args, **kwargs)
+            if desc == "_load_next_mult_id":
+                # Stream ID 8 is now allocated and recorded in `_unfinished_ids`.
+                # Deliver the cancellation here, exactly as a cancelled `/sync`
+                # would land it on this `await`.
+                await make_deferred_yieldable(cancel_enter)
+            return result
+
+        async def presence_like_write() -> None:
+            # Mirrors `PresenceStore.update_presence`: allocate an ID and
+            # "persist" under the context manager.
+            with LoggingContext(name="sync", server_name=self.hs.hostname):
+                async with id_gen.get_next():
+                    pass
+
+        with mock.patch.object(
+            id_gen._db, "runInteraction", new=blocking_run_interaction
+        ):
+            write = ensureDeferred(presence_like_write())
+
+            # The write is now blocked inside `__aenter__`, i.e. after stream ID
+            # 8 has been allocated and added to `_unfinished_ids`.
+            self.assertNoResult(write)
+
+            # The client goes away and the `/sync` request is cancelled.
+            cancel_enter.errback(CancelledError())
+
+            # The cancellation must surface as a `CancelledError`.
+            self.get_failure(write, CancelledError)
+
+        # The cancelled write never persisted a row for ID 8, so the generator
+        # must not let that abandoned ID wedge the position. A subsequent
+        # *successful* write should be able to advance the persisted token.
+        async def _successful_write() -> None:
+            async with id_gen.get_next():
+                pass
+
+        self.get_success(_successful_write())
+
+        # On the buggy code the token is still stuck at 7 (ID 8 is leaked in
+        # `_unfinished_ids`, blocking everything behind it). Once the leak is
+        # fixed, the token advances to 9: ID 8 was allocated (and abandoned) by
+        # the cancelled write, so the successful write above takes ID 9.
+        self.assertEqual(
+            id_gen.get_current_token_for_writer("master"),
+            9,
+            "presence stream position is wedged by the cancelled allocation",
+        )
+
     def test_out_of_order_finish(self) -> None:
         """Test that IDs persisted out of order are correctly handled"""