From fc9dbae23b3bfcbcfcbaa05410a890cea0a706cf Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Thu, 9 Mar 2023 16:06:37 +0000
Subject: [PATCH] fixup! Rewrite `wait_if_user_pending` to fix races

---
 .../src/session_manager/sessions.rs              | 16 +++++++++++-----
 crates/matrix-sdk-crypto/src/store/mod.rs        |  4 ++++
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
index 3c5ebc6a3..89e9a3759 100644
--- a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
+++ b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
@@ -394,7 +394,7 @@ impl SessionManager {
 
 #[cfg(test)]
 mod tests {
-    use std::{collections::BTreeMap, iter, sync::Arc};
+    use std::{collections::BTreeMap, iter, ops::Deref, sync::Arc};
 
     use dashmap::DashMap;
     use matrix_sdk_common::locks::Mutex;
@@ -547,7 +547,13 @@ mod tests {
 
         // ... and start off an attempt to get the missing sessions. This should block
         // for now.
-        let missing_sessions_future = manager.get_missing_sessions(iter::once(bob.user_id()));
+        let missing_sessions_task = {
+            let manager = manager.clone();
+            let bob_user_id = bob.user_id.clone();
+            tokio::spawn(async move {
+                manager.get_missing_sessions(iter::once(bob_user_id.deref())).await
+            })
+        };
 
         // the initial keys query completes, and we start another
         let response_json = json!({ "device_keys": { manager.account.user_id(): {}}});
@@ -567,10 +573,10 @@ mod tests {
             KeysQueryResponse::try_from_http_response(response_from_file(&response_json)).unwrap();
         identity_manager.receive_keys_query_response(&key_query_txn_id, &response).await.unwrap();
 
-        // the missing_sessions_future should now finally complete, with a claim
+        // the missing_sessions_task should now finally complete, with a claim
         // including bob's device
-        let (_, keys_claim_request) = missing_sessions_future.await.unwrap().unwrap();
-        //info!("Key claim: {:?}", keys_claim_request);
+        let (_, keys_claim_request) = missing_sessions_task.await.unwrap().unwrap().unwrap();
+        info!("Key claim request: {:?}", keys_claim_request.one_time_keys);
         let bob_key_claims = keys_claim_request.one_time_keys.get(bob.user_id()).unwrap();
         assert!(bob_key_claims.contains_key(bob_device.device_id()));
     }
diff --git a/crates/matrix-sdk-crypto/src/store/mod.rs b/crates/matrix-sdk-crypto/src/store/mod.rs
index c8cd0e68b..9e3d478e6 100644
--- a/crates/matrix-sdk-crypto/src/store/mod.rs
+++ b/crates/matrix-sdk-crypto/src/store/mod.rs
@@ -190,12 +190,15 @@ impl UsersForKeyQuery {
                 // the TaskAwaitingKeyQuery has been dropped, so it probably timed out and the
                 // caller went away. We can remove it from our list whether or not it's for this
                 // user.
+                trace!("removing expired waiting task");
                 return false;
             };
             if waiter.user == user && waiter.sequence_number.wrapping_sub(query_sequence) <= 0 {
+                trace!(?user, query_sequence, waiter_sequence=waiter.sequence_number, "removing completed waiting task");
                 waiter.completed.store(true, Ordering::Relaxed);
                 false
             } else {
+                trace!(?user, query_sequence, waiter_user=?waiter.user, waiter_sequence=waiter.sequence_number, "retaining still-waiting task");
                 true
             }
         });
@@ -234,6 +237,7 @@ impl UsersForKeyQuery {
         match self.user_map.get(user) {
             None => None,
             Some(&sequence_number) => {
+                trace!(?user, sequence_number, "registering new waiting task");
                 let waiter = Arc::new(KeysQueryWaiter {
                     sequence_number,
                     user: user.to_owned(),