diff --git a/src/lib.rs b/src/lib.rs index 8649aa7f..f1c1848e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ mod serde; mod set; pub mod setref; mod t; +pub mod try_result; mod util; #[cfg(feature = "rayon")] @@ -32,6 +33,7 @@ pub use read_only::ReadOnlyView; pub use set::DashSet; use std::collections::hash_map::RandomState; pub use t::Map; +use try_result::TryResult; cfg_if! { if #[cfg(feature = "raw-api")] { @@ -446,6 +448,61 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: BuildHasher + Clone> DashMap { self._get_mut(key) } + /// Get an immutable reference to an entry in the map, if the shard is not locked. + /// If the shard is locked, the function will return [TryResult::Locked]. + /// + /// # Examples + /// + /// ``` + /// use dashmap::DashMap; + /// use dashmap::try_result::TryResult; + /// + /// let map = DashMap::new(); + /// map.insert("Johnny", 21); + /// + /// assert_eq!(*map.try_get("Johnny").unwrap(), 21); + /// + /// let _result1_locking = map.get_mut("Johnny"); + /// + /// let result2 = map.try_get("Johnny"); + /// assert!(result2.is_locked()); + /// ``` + pub fn try_get(&'a self, key: &Q) -> TryResult> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self._try_get(key) + } + + /// Get a mutable reference to an entry in the map, if the shard is not locked. + /// If the shard is locked, the function will return [TryResult::Locked]. + /// + /// # Examples + /// + /// ``` + /// use dashmap::DashMap; + /// use dashmap::try_result::TryResult; + /// + /// let map = DashMap::new(); + /// map.insert("Johnny", 21); + /// + /// *map.try_get_mut("Johnny").unwrap() += 1; + /// assert_eq!(*map.get("Johnny").unwrap(), 22); + /// + /// let _result1_locking = map.get("Johnny"); + /// + /// let result2 = map.try_get_mut("Johnny"); + /// assert!(result2.is_locked()); + /// ``` + pub fn try_get_mut(&'a self, key: &Q) -> TryResult> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self._try_get_mut(key) + } + /// Remove excess capacity to reduce memory usage. /// /// **Locking behaviour:** May deadlock if called when holding any sort of reference into the map. @@ -640,6 +697,14 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: BuildHasher + Clone> DashMap { pub fn entry(&'a self, key: K) -> Entry<'a, K, V, S> { self._entry(key) } + + /// Advanced entry API that tries to mimic `std::collections::HashMap`. + /// See the documentation on `dashmap::mapref::entry` for more details. + /// + /// Returns None if the shard is currently locked. + pub fn try_entry(&'a self, key: K) -> Option> { + self._try_entry(key) + } } impl<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + BuildHasher + Clone> Map<'a, K, V, S> @@ -667,6 +732,18 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + BuildHasher + Clone> Map<'a, K, V, S> self.shards.get_unchecked(i).write() } + unsafe fn _try_yield_read_shard(&'a self, i: usize) -> Option>> { + debug_assert!(i < self.shards.len()); + + self.shards.get_unchecked(i).try_read() + } + + unsafe fn _try_yield_write_shard(&'a self, i: usize) -> Option>> { + debug_assert!(i < self.shards.len()); + + self.shards.get_unchecked(i).try_write() + } + fn _insert(&self, key: K, value: V) -> Option { let hash = self.hash_usize(&key); @@ -799,6 +876,60 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + BuildHasher + Clone> Map<'a, K, V, S> } } + fn _try_get(&'a self, key: &Q) -> TryResult> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + let hash = self.hash_usize(&key); + + let idx = self.determine_shard(hash); + + let shard = match unsafe { self._try_yield_read_shard(idx) } { + Some(shard) => shard, + None => return TryResult::Locked, + }; + + if let Some((kptr, vptr)) = shard.get_key_value(key) { + unsafe { + let kptr = util::change_lifetime_const(kptr); + + let vptr = &mut *vptr.as_ptr(); + + TryResult::Present(Ref::new(shard, kptr, vptr)) + } + } else { + TryResult::Absent + } + } + + fn _try_get_mut(&'a self, key: &Q) -> TryResult> + where + K: Borrow, + Q: Hash + Eq + ?Sized + { + let hash = self.hash_usize(&key); + + let idx = self.determine_shard(hash); + + let shard = match unsafe { self._try_yield_write_shard(idx) } { + Some(shard) => shard, + None => return TryResult::Locked, + }; + + if let Some((kptr, vptr)) = shard.get_key_value(key) { + unsafe { + let kptr = util::change_lifetime_const(kptr); + + let vptr = &mut *vptr.as_ptr(); + + TryResult::Present(RefMut::new(shard, kptr, vptr)) + } + } else { + TryResult::Absent + } + } + fn _shrink_to_fit(&self) { self.shards.iter().for_each(|s| s.write().shrink_to_fit()); } @@ -866,6 +997,29 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + BuildHasher + Clone> Map<'a, K, V, S> } } + fn _try_entry(&'a self, key: K) -> Option> { + let hash = self.hash_usize(&key); + + let idx = self.determine_shard(hash); + + let shard = match unsafe { self._try_yield_write_shard(idx) } { + Some(shard) => shard, + None => return None, + }; + + if let Some((kptr, vptr)) = shard.get_key_value(&key) { + unsafe { + let kptr = util::change_lifetime_const(kptr); + + let vptr = &mut *vptr.as_ptr(); + + Some(Entry::Occupied(OccupiedEntry::new(shard, key, (kptr, vptr)))) + } + } else { + Some(Entry::Vacant(VacantEntry::new(shard, key))) + } + } + fn _hasher(&self) -> S { self.hasher.clone() } @@ -1090,4 +1244,32 @@ mod tests { let not_in_map = dm.view(&30, |_k, _v| false); assert_eq!(not_in_map, None); } + + #[test] + fn test_try_get() { + { + let map = DashMap::new(); + map.insert("Johnny", 21); + + assert_eq!(*map.try_get("Johnny").unwrap(), 21); + + let _result1_locking = map.get_mut("Johnny"); + + let result2 = map.try_get("Johnny"); + assert!(result2.is_locked()); + } + + { + let map = DashMap::new(); + map.insert("Johnny", 21); + + *map.try_get_mut("Johnny").unwrap() += 1; + assert_eq!(*map.get("Johnny").unwrap(), 22); + + let _result1_locking = map.get("Johnny"); + + let result2 = map.try_get_mut("Johnny"); + assert!(result2.is_locked()); + } + } } diff --git a/src/t.rs b/src/t.rs index 61684619..40259832 100644 --- a/src/t.rs +++ b/src/t.rs @@ -3,6 +3,7 @@ use crate::iter::{Iter, IterMut}; use crate::mapref::entry::Entry; use crate::mapref::one::{Ref, RefMut}; +use crate::try_result::TryResult; use crate::HashMap; use core::borrow::Borrow; use core::hash::{BuildHasher, Hash}; @@ -27,6 +28,16 @@ pub trait Map<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + Clone + BuildHasher> { /// The index must not be out of bounds. unsafe fn _yield_write_shard(&'a self, i: usize) -> RwLockWriteGuard<'a, HashMap>; + /// # Safety + /// + /// The index must not be out of bounds. + unsafe fn _try_yield_read_shard(&'a self, i: usize) -> Option>>; + + /// # Safety + /// + /// The index must not be out of bounds. + unsafe fn _try_yield_write_shard(&'a self, i: usize) -> Option>>; + fn _insert(&self, key: K, value: V) -> Option; fn _remove(&self, key: &Q) -> Option<(K, V)> @@ -62,6 +73,16 @@ pub trait Map<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + Clone + BuildHasher> { K: Borrow, Q: Hash + Eq + ?Sized; + fn _try_get(&'a self, key: &Q) -> TryResult> + where + K: Borrow, + Q: Hash + Eq + ?Sized; + + fn _try_get_mut(&'a self, key: &Q) -> TryResult> + where + K: Borrow, + Q: Hash + Eq + ?Sized; + fn _shrink_to_fit(&self); fn _retain(&self, f: impl FnMut(&K, &mut V) -> bool); @@ -84,6 +105,8 @@ pub trait Map<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + Clone + BuildHasher> { fn _entry(&'a self, key: K) -> Entry<'a, K, V, S>; + fn _try_entry(&'a self, key: K) -> Option>; + fn _hasher(&self) -> S; // provided diff --git a/src/try_result.rs b/src/try_result.rs new file mode 100644 index 00000000..082c6827 --- /dev/null +++ b/src/try_result.rs @@ -0,0 +1,59 @@ + +/// Represents the result of a non-blocking read from a [DashMap](crate::DashMap). +#[derive(Debug)] +pub enum TryResult { + /// The value was present in the map, and the lock for the shard was successfully obtained. + Present(R), + /// The shard wasn't locked, and the value wasn't present in the map. + Absent, + /// The shard was locked. + Locked, +} + +impl TryResult { + /// Returns `true` if the value was present in the map, and the lock for the shard was successfully obtained. + #[inline] + pub fn is_present(&self) -> bool { + match self { + TryResult::Present(_) => true, + _ => false, + } + } + + /// Returns `true` if the shard wasn't locked, and the value wasn't present in the map. + #[inline] + pub fn is_absent(&self) -> bool { + match self { + TryResult::Absent => true, + _ => false, + } + } + + /// Returns `true` if the shard was locked. + #[inline] + pub fn is_locked(&self) -> bool { + match self { + TryResult::Locked => true, + _ => false, + } + } + + /// If `self` is [Present](TryResult::Present), returns the reference to the value in the map. + /// Panics if `self` is not [Present](TryResult::Present). + pub fn unwrap(self) -> R { + match self { + TryResult::Present(r) => r, + TryResult::Locked => panic!("Called unwrap() on TryResult::Locked"), + TryResult::Absent => panic!("Called unwrap() on TryResult::Absent"), + } + } + + /// If `self` is [Present](TryResult::Present), returns the reference to the value in the map. + /// If `self` is not [Present](TryResult::Present), returns `None`. + pub fn try_unwrap(self) -> Option { + match self { + TryResult::Present(r) => Some(r), + _ => None, + } + } +}