refactor: Use a common code for try_take_leased_lock.

This code is shared by all `MemoryStore` implementations.
matrix-org · Nov 6, 2024 · 94bd421 · 94bd421
1 parent 7b3eb0b
commit 94bd421
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 111 deletions.
diff --git a/crates/matrix-sdk-base/src/event_cache_store/memory_store.rs b/crates/matrix-sdk-base/src/event_cache_store/memory_store.rs
@@ -12,15 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{
-    collections::{hash_map::Entry, HashMap},
-    num::NonZeroUsize,
-    sync::RwLock as StdRwLock,
-    time::{Duration, Instant},
-};
+use std::{collections::HashMap, num::NonZeroUsize, sync::RwLock as StdRwLock, time::Instant};
 
 use async_trait::async_trait;
-use matrix_sdk_common::ring_buffer::RingBuffer;
+use matrix_sdk_common::{
+    ring_buffer::RingBuffer, store_locks::memory_store_helper::try_take_leased_lock,
+};
 use ruma::{MxcUri, OwnedMxcUri};
 
 use super::{EventCacheStore, EventCacheStoreError, Result};
@@ -66,44 +63,7 @@ impl EventCacheStore for MemoryStore {
         key: &str,
         holder: &str,
     ) -> Result<bool, Self::Error> {
-        let now = Instant::now();
-        let expiration = now + Duration::from_millis(lease_duration_ms.into());
-
-        match self.leases.write().unwrap().entry(key.to_owned()) {
-            // There is an existing holder.
-            Entry::Occupied(mut entry) => {
-                let (current_holder, current_expiration) = entry.get_mut();
-
-                if current_holder == holder {
-                    // We had the lease before, extend it.
-                    *current_expiration = expiration;
-
-                    Ok(true)
-                } else {
-                    // We didn't have it.
-                    if *current_expiration < now {
-                        // Steal it!
-                        *current_holder = holder.to_owned();
-                        *current_expiration = expiration;
-
-                        Ok(true)
-                    } else {
-                        // We tried our best.
-                        Ok(false)
-                    }
-                }
-            }
-
-            // There is no holder, easy.
-            Entry::Vacant(entry) => {
-                entry.insert((
-                    holder.to_owned(),
-                    Instant::now() + Duration::from_millis(lease_duration_ms.into()),
-                ));
-
-                Ok(true)
-            }
-        }
+        Ok(try_take_leased_lock(&self.leases, lease_duration_ms, key, holder))
     }
 
     async fn add_media_content(&self, request: &MediaRequest, data: Vec<u8>) -> Result<()> {

diff --git a/crates/matrix-sdk-common/src/store_locks.rs b/crates/matrix-sdk-common/src/store_locks.rs
@@ -338,7 +338,7 @@ pub enum LockStoreError {
 mod tests {
     use std::{
         collections::HashMap,
-        sync::{atomic, Arc, Mutex},
+        sync::{atomic, Arc, RwLock},
         time::Instant,
     };
 
@@ -350,47 +350,18 @@ mod tests {
     };
 
     use super::{
-        BackingStore, CrossProcessStoreLock, CrossProcessStoreLockGuard, LockStoreError,
-        EXTEND_LEASE_EVERY_MS,
+        memory_store_helper::try_take_leased_lock, BackingStore, CrossProcessStoreLock,
+        CrossProcessStoreLockGuard, LockStoreError, EXTEND_LEASE_EVERY_MS,
     };
 
     #[derive(Clone, Default)]
     struct TestStore {
-        leases: Arc<Mutex<HashMap<String, (String, Instant)>>>,
+        leases: Arc<RwLock<HashMap<String, (String, Instant)>>>,
     }
 
     impl TestStore {
         fn try_take_leased_lock(&self, lease_duration_ms: u32, key: &str, holder: &str) -> bool {
-            let now = Instant::now();
-            let expiration = now + Duration::from_millis(lease_duration_ms.into());
-            let mut leases = self.leases.lock().unwrap();
-            if let Some(prev) = leases.get_mut(key) {
-                if prev.0 == holder {
-                    // We had the lease before, extend it.
-                    prev.1 = expiration;
-                    true
-                } else {
-                    // We didn't have it.
-                    if prev.1 < now {
-                        // Steal it!
-                        prev.0 = holder.to_owned();
-                        prev.1 = expiration;
-                        true
-                    } else {
-                        // We tried our best.
-                        false
-                    }
-                }
-            } else {
-                leases.insert(
-                    key.to_owned(),
-                    (
-                        holder.to_owned(),
-                        Instant::now() + Duration::from_millis(lease_duration_ms.into()),
-                    ),
-                );
-                true
-            }
+            try_take_leased_lock(&self.leases, lease_duration_ms, key, holder)
         }
     }
 
@@ -525,3 +496,59 @@ mod tests {
         Ok(())
     }
 }
+
+/// Some code that is shared by almost all `MemoryStore` implementations out
+/// there.
+pub mod memory_store_helper {
+    use std::{
+        collections::{hash_map::Entry, HashMap},
+        sync::RwLock,
+        time::{Duration, Instant},
+    };
+
+    pub fn try_take_leased_lock(
+        leases: &RwLock<HashMap<String, (String, Instant)>>,
+        lease_duration_ms: u32,
+        key: &str,
+        holder: &str,
+    ) -> bool {
+        let now = Instant::now();
+        let expiration = now + Duration::from_millis(lease_duration_ms.into());
+
+        match leases.write().unwrap().entry(key.to_owned()) {
+            // There is an existing holder.
+            Entry::Occupied(mut entry) => {
+                let (current_holder, current_expiration) = entry.get_mut();
+
+                if current_holder == holder {
+                    // We had the lease before, extend it.
+                    *current_expiration = expiration;
+
+                    true
+                } else {
+                    // We didn't have it.
+                    if *current_expiration < now {
+                        // Steal it!
+                        *current_holder = holder.to_owned();
+                        *current_expiration = expiration;
+
+                        true
+                    } else {
+                        // We tried our best.
+                        false
+                    }
+                }
+            }
+
+            // There is no holder, easy.
+            Entry::Vacant(entry) => {
+                entry.insert((
+                    holder.to_owned(),
+                    Instant::now() + Duration::from_millis(lease_duration_ms.into()),
+                ));
+
+                true
+            }
+        }
+    }
+}
diff --git a/crates/matrix-sdk-crypto/src/store/memorystore.rs b/crates/matrix-sdk-crypto/src/store/memorystore.rs
@@ -13,13 +13,14 @@
 // limitations under the License.
 
 use std::{
-    collections::{hash_map::Entry, BTreeMap, HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     convert::Infallible,
     sync::RwLock as StdRwLock,
-    time::{Duration, Instant},
+    time::Instant,
 };
 
 use async_trait::async_trait;
+use matrix_sdk_common::store_locks::memory_store_helper::try_take_leased_lock;
 use ruma::{
     events::secret::request::SecretName, DeviceId, OwnedDeviceId, OwnedRoomId, OwnedTransactionId,
     OwnedUserId, RoomId, TransactionId, UserId,
@@ -631,36 +632,7 @@ impl CryptoStore for MemoryStore {
         key: &str,
         holder: &str,
     ) -> Result<bool> {
-        let now = Instant::now();
-        let expiration = now + Duration::from_millis(lease_duration_ms.into());
-        match self.leases.write().unwrap().entry(key.to_owned()) {
-            Entry::Occupied(mut o) => {
-                let prev = o.get_mut();
-                if prev.0 == holder {
-                    // We had the lease before, extend it.
-                    prev.1 = expiration;
-                    Ok(true)
-                } else {
-                    // We didn't have it.
-                    if prev.1 < now {
-                        // Steal it!
-                        prev.0 = holder.to_owned();
-                        prev.1 = expiration;
-                        Ok(true)
-                    } else {
-                        // We tried our best.
-                        Ok(false)
-                    }
-                }
-            }
-            Entry::Vacant(v) => {
-                v.insert((
-                    holder.to_owned(),
-                    Instant::now() + Duration::from_millis(lease_duration_ms.into()),
-                ));
-                Ok(true)
-            }
-        }
+        Ok(try_take_leased_lock(&self.leases, lease_duration_ms, key, holder))
     }
 }