Sliding Sync: Notify and sync when one-time keys or fallback keys are claimed/uploaded

changelog.d/17820.bugfix

@@ -0,0 +1 @@
+Fix sync not being notified when device one-time keys or fallback keys are uploaded/claimed.

synapse/handlers/e2e_keys.py

@@ -70,6 +70,7 @@ def __init__(self, hs: "HomeServer"):
         self.is_mine = hs.is_mine
         self.clock = hs.get_clock()
         self._worker_lock_handler = hs.get_worker_locks_handler()
+        self._notifier = hs.get_notifier()
 
         federation_registry = hs.get_federation_registry()
 
@@ -615,7 +616,7 @@ async def claim_local_one_time_keys(
         3. Attempt to fetch fallback keys from the database.
 
         Args:
-            local_query: An iterable of tuples of (user ID, device ID, algorithm).
+            local_query: An iterable of tuples of (user ID, device ID, algorithm, number of keys to claim).
             always_include_fallback_keys: True to always include fallback keys.
 
         Returns:
@@ -629,6 +630,7 @@ async def claim_local_one_time_keys(
         ]
 
         otk_results, not_found = await self.store.claim_e2e_one_time_keys(local_query)
+        self._notifier.notify_one_time_keys_changed(otk_results.keys())
 
         # If the application services have not provided any keys via the C-S
         # API, query it directly for one-time keys.
@@ -639,6 +641,7 @@ async def claim_local_one_time_keys(
                 appservice_results,
                 not_found,
             ) = await self._appservice_handler.claim_e2e_one_time_keys(not_found)
+            self._notifier.notify_one_time_keys_changed(appservice_results.keys())
         else:
             appservice_results = {}
 
@@ -693,6 +696,7 @@ async def claim_local_one_time_keys(
         # For each user that does not have a one-time keys available, see if
         # there is a fallback key.
         fallback_results = await self.store.claim_e2e_fallback_keys(fallback_query)
+        self._notifier.notify_one_time_keys_changed(fallback_results.keys())
 
         # Return the results in order, each item from the input query should
         # only appear once in the combined list.
@@ -833,6 +837,7 @@ async def upload_keys_for_user(
             await self._upload_one_time_keys_for_user(
                 user_id, device_id, time_now, one_time_keys
             )
+            self._notifier.notify_one_time_keys_changed([user_id])
         else:
             log_kv(
                 {"message": "Did not update one_time_keys", "reason": "no keys given"}
@@ -849,6 +854,7 @@ async def upload_keys_for_user(
                 }
             )
             await self.store.set_e2e_fallback_keys(user_id, device_id, fallback_keys)
+            self._notifier.notify_one_time_keys_changed([user_id])
         elif fallback_keys:
             log_kv({"message": "Did not update fallback_keys", "reason": "not a dict"})
         else:

synapse/notifier.py

@@ -590,6 +590,8 @@ def on_new_replication_data(self) -> None:
         without waking up any of the normal user event streams"""
         self.notify_replication()
 
+    # FIXME: Should this be renamed to `wait_for_activity`? This listens for new events
+    # and when one-time keys are claimed which doesn't correspond to an event.
     async def wait_for_events(
         self,
         user_id: str,
@@ -900,6 +902,43 @@ def notify_lock_released(
         for cb in self._lock_released_callback:
             cb(instance_name, lock_name, lock_key)
 
+    def notify_one_time_keys_changed(
+        self,
+        users: Union[StrCollection, Collection[UserID]],
+    ) -> None:
+        """
+        Used by handlers to inform the notifier that a one-time key has been
+        claimed or uploaded
+        """
+        # Bail early if there is nothing to do
+        if not users:
+            return
+
+        time_now_ms = self.clock.time_msec()
+        current_token = self.event_sources.get_current_token()
+        listeners: List["Deferred[StreamToken]"] = []
+        for user in users:
+            user_stream = self.user_to_user_stream.get(str(user))
+            if user_stream is None:
+                continue
+
+            try:
+                listeners.extend(
+                    user_stream.update_and_fetch_deferreds(current_token, time_now_ms)
+                )
+            except Exception:
+                logger.exception("Failed to notify listener")
+
+        # We resolve all these deferreds in one go so that we only need to
+        # call `PreserveLoggingContext` once, as it has a bunch of overhead
+        # (to calculate performance stats)
+        with PreserveLoggingContext():
+            for listener in listeners:
+                listener.callback(current_token)
+
+        # FIXME: How can we poke the replication so that other workers also see the
+        # one-time key change
+
 
 @attr.s(auto_attribs=True)
 class ReplicationNotifier: