crypto(perf): don't hold the cache lock while waiting on a user key query (#2806)

* crypto(fix): don't hold the cache lock while waiting on a user key query

Fixes #2802. The lock was only useful to sync the database and the in-memory cache for the users awaiting a key query request.
So it's possible to slightly tweak the API by moving the method from `SyncedKeyQueryManager` to non-synced `KeyQueryManager`, and require a
`StoreCacheGuard` (i.e. the owned lock, so we can manually drop it when we feel like so).

I've looked at all the other methods, and they do require the cache for writing into it and the store.
At the limit we could also move `SyncedKeyQueryManager::users_for_key_query`
into `KeyQueryManager`, but the lock in there is hold for a very short-time, so it shouldn't be an issue.

* Add test for the key query deadlock while waiting for the response.

* Update crates/matrix-sdk-crypto/src/machine.rs

Co-authored-by: Richard van der Hoff <[email protected]>

Author: bnjbvr

crates/matrix-sdk-crypto/src/machine.rs

@@ -408,6 +408,9 @@ impl OlmMachine {
     /// This can be useful if we need the results from [`get_identity`] or
     /// [`get_user_devices`] to be as up-to-date as possible.
     ///
+    /// Note that this request won't be awaited by other calls waiting for a
+    /// user's or device's keys, since this is an out-of-band query.
+    ///
     /// # Arguments
     ///
     /// * `users` - list of users whose keys should be queried
@@ -1606,10 +1609,8 @@ impl OlmMachine {
             self.inner
                 .identity_manager
                 .key_query_manager
-                .synced(&cache)
-                .await?
-                .wait_if_user_key_query_pending(timeout, user_id)
-                .await;
+                .wait_if_user_key_query_pending(cache, timeout, user_id)
+                .await?;
         }
         Ok(())
     }
@@ -4029,4 +4030,52 @@ pub(crate) mod tests {
             "Our identity should not be verified when there's a mismatch in the cross-signing keys"
         );
     }
+
+    #[async_test]
+    async fn test_wait_on_key_query_doesnt_block_store() {
+        // Waiting for a key query shouldn't delay other write attempts to the store.
+        // This test will end immediately if it works, and times out after a few seconds
+        // if it failed.
+
+        let machine = OlmMachine::new(bob_id(), bob_device_id()).await;
+
+        // Mark Alice as a tracked user, so it gets into the groups of users for which
+        // we need to query keys.
+        machine.update_tracked_users([alice_id()]).await.unwrap();
+
+        // Start a background task that will wait for the key query to finish silently
+        // in the background.
+        let machine_cloned = machine.clone();
+        let wait = tokio::spawn(async move {
+            let machine = machine_cloned;
+            let user_devices =
+                machine.get_user_devices(alice_id(), Some(Duration::from_secs(10))).await.unwrap();
+            assert!(user_devices.devices().next().is_some());
+        });
+
+        // Let the background task work first.
+        tokio::task::yield_now().await;
+
+        // Create a key upload request and process it back immediately.
+        let requests = machine.bootstrap_cross_signing(false).await.unwrap();
+
+        let req = requests.upload_keys_req.expect("upload keys request should be there");
+        let response = keys_upload_response();
+        let mark_request_as_sent = machine.mark_request_as_sent(&req.request_id, &response);
+        tokio::time::timeout(Duration::from_secs(5), mark_request_as_sent)
+            .await
+            .expect("no timeout")
+            .expect("the underlying request has been marked as sent");
+
+        // Answer the key query, so the background task completes immediately?
+        let response = keys_query_response();
+        let key_queries = machine.inner.identity_manager.users_for_key_query().await.unwrap();
+
+        for (id, _) in key_queries {
+            machine.mark_request_as_sent(&id, &response).await.unwrap();
+        }
+
+        // The waiting should successfully complete.
+        wait.await.unwrap();
+    }
 }

crates/matrix-sdk-crypto/src/session_manager/sessions.rs

@@ -188,10 +188,8 @@ impl SessionManager {
                 .key_request_machine
                 .identity_manager()
                 .key_query_manager
-                .synced(&cache)
+                .wait_if_user_key_query_pending(cache, Self::KEYS_QUERY_WAIT_TIME, user_id)
                 .await?
-                .wait_if_user_key_query_pending(Self::KEYS_QUERY_WAIT_TIME, user_id)
-                .await
             {
                 WasPending => self.store.get_readonly_devices_filtered(user_id).await?,
                 _ => user_devices,

crates/matrix-sdk-crypto/src/store/mod.rs

@@ -160,56 +160,40 @@ impl KeyQueryManager {
 
         Ok(())
     }
-}
-
-pub(crate) struct SyncedKeyQueryManager<'a> {
-    cache: &'a StoreCache,
-    manager: &'a KeyQueryManager,
-}
-
-impl<'a> SyncedKeyQueryManager<'a> {
-    /// Add entries to the list of users being tracked for device changes
-    ///
-    /// Any users not already on the list are flagged as awaiting a key query.
-    /// Users that were already in the list are unaffected.
-    pub async fn update_tracked_users(&self, users: impl Iterator<Item = &UserId>) -> Result<()> {
-        let mut store_updates = Vec::new();
-        let mut key_query_lock = self.manager.users_for_key_query.lock().await;
-
-        {
-            let mut tracked_users = self.cache.tracked_users.write().unwrap();
-            for user_id in users {
-                if tracked_users.insert(user_id.to_owned()) {
-                    key_query_lock.insert_user(user_id);
-                    store_updates.push((user_id, true))
-                }
-            }
-        }
-
-        self.cache.store.save_tracked_users(&store_updates).await
-    }
 
     /// Wait for a `/keys/query` response to be received if one is expected for
     /// the given user.
     ///
     /// If the given timeout elapses, the method will stop waiting and return
-    /// `UserKeyQueryResult::TimeoutExpired`
+    /// `UserKeyQueryResult::TimeoutExpired`.
+    ///
+    /// Requires a [`StoreCacheGuard`] to make sure the users for which a key
+    /// query is pending are up to date, but doesn't hold on to it
+    /// thereafter: the lock is short-lived in this case.
     pub async fn wait_if_user_key_query_pending(
         &self,
+        cache: StoreCacheGuard,
         timeout_duration: Duration,
         user: &UserId,
-    ) -> UserKeyQueryResult {
-        let mut users_for_key_query = self.manager.users_for_key_query.lock().await;
+    ) -> Result<UserKeyQueryResult> {
+        {
+            // Drop the cache early, so we don't keep it while waiting (since writing the
+            // results requires to write in the cache, thus take another lock).
+            self.ensure_sync_tracked_users(&cache).await?;
+            drop(cache);
+        }
+
+        let mut users_for_key_query = self.users_for_key_query.lock().await;
         let Some(waiter) = users_for_key_query.maybe_register_waiting_task(user) else {
-            return UserKeyQueryResult::WasNotPending;
+            return Ok(UserKeyQueryResult::WasNotPending);
         };
 
         let wait_for_completion = async {
             while !waiter.completed.load(Ordering::Relaxed) {
                 // Register for being notified before releasing the mutex, so
                 // it's impossible to miss a wakeup between the last check for
                 // whether we should wait, and starting to wait.
-                let mut notified = pin!(self.manager.users_for_key_query_notify.notified());
+                let mut notified = pin!(self.users_for_key_query_notify.notified());
                 notified.as_mut().enable();
                 drop(users_for_key_query);
 
@@ -219,7 +203,7 @@ impl<'a> SyncedKeyQueryManager<'a> {
                 // Reclaim the lock before checking the flag to avoid races
                 // when two notifications happen right after each other and the
                 // second one sets the flag we want to wait for.
-                users_for_key_query = self.manager.users_for_key_query.lock().await;
+                users_for_key_query = self.users_for_key_query.lock().await;
             }
         };
 
@@ -231,10 +215,38 @@ impl<'a> SyncedKeyQueryManager<'a> {
                     not finish yet, some devices might be missing."
                 );
 
-                UserKeyQueryResult::TimeoutExpired
+                Ok(UserKeyQueryResult::TimeoutExpired)
+            }
+            _ => Ok(UserKeyQueryResult::WasPending),
+        }
+    }
+}
+
+pub(crate) struct SyncedKeyQueryManager<'a> {
+    cache: &'a StoreCache,
+    manager: &'a KeyQueryManager,
+}
+
+impl<'a> SyncedKeyQueryManager<'a> {
+    /// Add entries to the list of users being tracked for device changes
+    ///
+    /// Any users not already on the list are flagged as awaiting a key query.
+    /// Users that were already in the list are unaffected.
+    pub async fn update_tracked_users(&self, users: impl Iterator<Item = &UserId>) -> Result<()> {
+        let mut store_updates = Vec::new();
+        let mut key_query_lock = self.manager.users_for_key_query.lock().await;
+
+        {
+            let mut tracked_users = self.cache.tracked_users.write().unwrap();
+            for user_id in users {
+                if tracked_users.insert(user_id.to_owned()) {
+                    key_query_lock.insert_user(user_id);
+                    store_updates.push((user_id, true))
+                }
             }
-            _ => UserKeyQueryResult::WasPending,
         }
+
+        self.cache.store.save_tracked_users(&store_updates).await
     }
 
     /// Process notifications that users have changed devices.