Safe async event cache

changelog.d/13308.bugfix

@@ -0,0 +1 @@
+Fix async get event cache invalidation logic. Contributed by Nick @ Beeper (@fizzadar).

synapse/storage/_base.py

@@ -96,6 +96,10 @@ def _attempt_to_invalidate_cache(
         cache doesn't exist. Mainly used for invalidating caches on workers,
         where they may not have the cache.
 
+        Note that this function does not invalidate any remote caches, only the
+        local in-memory ones. Any remote invalidation must be performed before
+        calling this.
+
         Args:
             cache_name
             key: Entry to invalidate. If None then invalidates the entire
@@ -112,7 +116,10 @@ def _attempt_to_invalidate_cache(
         if key is None:
             cache.invalidate_all()
         else:
-            cache.invalidate(tuple(key))
+            # Prefer any local-only invalidation method. Invalidating any non-local
+            # cache must be be done before this.
+            invalidate_method = getattr(cache, "invalidate_local", cache.invalidate)
+            invalidate_method(tuple(key))
 
 
 def db_to_json(db_content: Union[memoryview, bytes, bytearray, str]) -> Any:

synapse/storage/database.py

@@ -23,6 +23,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Awaitable,
     Callable,
     Collection,
     Dict,
@@ -33,6 +34,7 @@
     Tuple,
     Type,
     TypeVar,
+    Union,
     cast,
     overload,
 )
@@ -57,7 +59,7 @@
 from synapse.storage.background_updates import BackgroundUpdater
 from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, Sqlite3Engine
 from synapse.storage.types import Connection, Cursor
-from synapse.util.async_helpers import delay_cancellation, maybe_awaitable
+from synapse.util.async_helpers import delay_cancellation
 from synapse.util.iterutils import batch_iter
 
 if TYPE_CHECKING:
@@ -208,7 +210,9 @@ def __getattr__(self, name: str) -> Any:
 
 
 # The type of entry which goes on our after_callbacks and exception_callbacks lists.
-_CallbackListEntry = Tuple[Callable[..., object], Tuple[object, ...], Dict[str, object]]
+_CallbackListEntry = Tuple[
+    Callable[..., Union[object, Awaitable]], Tuple[object, ...], Dict[str, object]
+]
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -796,6 +800,29 @@ async def runInteraction(
             The result of func
         """
 
+        async def _run_callbacks(callbacks: List[_CallbackListEntry]):
+            """
+            This function takes a list of mixed sync/async callbacks and executes
+            the async ones first and then the sync callbacks.
+
+            We do this with the assumption that async functions call out to external
+            systems (e.g. to invalidate a cache) and the sync functions make these
+            changes on any local in-memory caches/similar, and thus must be second.
+            """
+
+            sync_callbacks: List[_CallbackListEntry] = []
+
+            for cb, args, kwargs in callbacks:
+                if inspect.iscoroutinefunction(cb):
+                    awaitable = cb(*args, **kwargs)
+                    assert isinstance(awaitable, Awaitable)
+                    await awaitable
+                else:
+                    sync_callbacks.append((cb, args, kwargs))
+
+            for cb, args, kwargs in sync_callbacks:
+                cb(*args, **kwargs)
+
         async def _runInteraction() -> R:
             after_callbacks: List[_CallbackListEntry] = []
             exception_callbacks: List[_CallbackListEntry] = []
@@ -817,15 +844,10 @@ async def _runInteraction() -> R:
                         **kwargs,
                     )
 
-                for after_callback, after_args, after_kwargs in after_callbacks:
-                    await maybe_awaitable(after_callback(*after_args, **after_kwargs))
-
+                await _run_callbacks(after_callbacks)
                 return cast(R, result)
             except Exception:
-                for exception_callback, after_args, after_kwargs in exception_callbacks:
-                    await maybe_awaitable(
-                        exception_callback(*after_args, **after_kwargs)
-                    )
+                await _run_callbacks(exception_callbacks)
                 raise
 
         # To handle cancellation, we ensure that `after_callback`s and

synapse/storage/databases/main/censor_events.py

@@ -194,7 +194,7 @@ def delete_expired_event_txn(txn: LoggingTransaction) -> None:
             # changed its content in the database. We can't call
             # self._invalidate_cache_and_stream because self.get_event_cache isn't of the
             # right type.
-            txn.call_after(self._get_event_cache.invalidate, (event.event_id,))
+            self.invalidate_get_event_cache_after_txn(txn, event.event_id)
             # Send that invalidation to replication so that other workers also invalidate
             # the event cache.
             self._send_invalidation_to_replication(

synapse/storage/databases/main/events.py

@@ -1293,7 +1293,7 @@ def _update_room_depths_txn(
         depth_updates: Dict[str, int] = {}
         for event, context in events_and_contexts:
             # Remove the any existing cache entries for the event_ids
-            txn.call_after(self.store._invalidate_get_event_cache, event.event_id)
+            self.store.invalidate_get_event_cache_after_txn(txn, event.event_id)
             # Then update the `stream_ordering` position to mark the latest
             # event as the front of the room. This should not be done for
             # backfilled events because backfilled events have negative
@@ -1684,7 +1684,7 @@ def _store_redaction(self, txn: LoggingTransaction, event: EventBase) -> None:
         _invalidate_caches_for_event.
         """
         assert event.redacts is not None
-        txn.call_after(self.store._invalidate_get_event_cache, event.redacts)
+        self.store.invalidate_get_event_cache_after_txn(txn, event.redacts)
         txn.call_after(self.store.get_relations_for_event.invalidate, (event.redacts,))
         txn.call_after(self.store.get_applicable_edit.invalidate, (event.redacts,))
 

synapse/storage/databases/main/events_worker.py

@@ -712,17 +712,41 @@ async def get_missing_events_from_db() -> Dict[str, EventCacheEntry]:
 
         return event_entry_map
 
-    async def _invalidate_get_event_cache(self, event_id: str) -> None:
-        # First we invalidate the asynchronous cache instance. This may include
-        # out-of-process caches such as Redis/memcache. Once complete we can
-        # invalidate any in memory cache. The ordering is important here to
-        # ensure we don't pull in any remote invalid value after we invalidate
-        # the in-memory cache.
+    def invalidate_get_event_cache_after_txn(
+        self, txn: LoggingTransaction, event_id: str
+    ) -> None:
+        """
+        Prepares a database transaction to invalidate the get event cache for a given
+        event ID when executed successfully. This is achieved by attaching two callbacks
+        to the transaction, one to invalidate the async cache and one for the in memory
+        sync cache (importantly called in that order).
+
+        Arguments:
+            txn: the database transaction to attach the callbacks to
+            event_id: the event ID to be invalidated from caches
+        """
+
+        txn.call_after(self._invalidate_async_get_event_cache, event_id)
+        txn.call_after(self._invalidate_local_get_event_cache, event_id)
+
+    async def _invalidate_async_get_event_cache(self, event_id: str) -> None:
+        """
+        Invalidates an event in the asyncronous get event cache, which may be remote.
+
+        Arguments:
+            event_id: the event ID to invalidate
+        """
+
         await self._get_event_cache.invalidate((event_id,))
-        self._event_ref.pop(event_id, None)
-        self._current_event_fetches.pop(event_id, None)
 
     def _invalidate_local_get_event_cache(self, event_id: str) -> None:
+        """
+        Invalidates an event in local in-memory get event caches.
+
+        Arguments:
+            event_id: the event ID to invalidate
+        """
+
         self._get_event_cache.invalidate_local((event_id,))
         self._event_ref.pop(event_id, None)
         self._current_event_fetches.pop(event_id, None)

synapse/storage/databases/main/purge_events.py

@@ -302,7 +302,7 @@ def _purge_history_txn(
                 self._invalidate_cache_and_stream(
                     txn, self.have_seen_event, (room_id, event_id)
                 )
-                txn.call_after(self._invalidate_get_event_cache, event_id)
+                self.invalidate_get_event_cache_after_txn(txn, event_id)
 
         logger.info("[purge] done")