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 index 6ea0608e1f4..d75c9f09776 100644 --- 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 { - 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) -> Result<()> { diff --git a/crates/matrix-sdk-common/src/store_locks.rs b/crates/matrix-sdk-common/src/store_locks.rs index 6ee4de48a0f..13350c1fa25 100644 --- 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>>, + leases: Arc>>, } 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>, + 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 index b773d819d94..6b177ea2641 100644 --- 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 { - 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)) } }