Merge pull request #178 from Earthcomputer/try_lock

Add DashMap::try_get, DashMap::try_get_mut, DashMap::try_entry

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<K, V, S> {
         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<Q>(&'a self, key: &Q) -> TryResult<Ref<'a, K, V, S>>
+    where
+        K: Borrow<Q>,
+        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<Q>(&'a self, key: &Q) -> TryResult<RefMut<'a, K, V, S>>
+    where
+        K: Borrow<Q>,
+        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<K, V, S> {
     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<Entry<'a, K, V, S>> {
+        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<RwLockReadGuard<'a, HashMap<K, V, S>>> {
+        debug_assert!(i < self.shards.len());
+
+        self.shards.get_unchecked(i).try_read()
+    }
+
+    unsafe fn _try_yield_write_shard(&'a self, i: usize) -> Option<RwLockWriteGuard<'a, HashMap<K, V, S>>> {
+        debug_assert!(i < self.shards.len());
+
+        self.shards.get_unchecked(i).try_write()
+    }
+
     fn _insert(&self, key: K, value: V) -> Option<V> {
         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<Q>(&'a self, key: &Q) -> TryResult<Ref<'a, K, V, S>>
+    where
+        K: Borrow<Q>,
+        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<Q>(&'a self, key: &Q) -> TryResult<RefMut<'a, K, V, S>>
+    where
+        K: Borrow<Q>,
+        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<Entry<'a, K, V, S>> {
+        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());
+        }
+    }
 }

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<K, V, S>>;
 
+    /// # Safety
+    ///
+    /// The index must not be out of bounds.
+    unsafe fn _try_yield_read_shard(&'a self, i: usize) -> Option<RwLockReadGuard<'a, HashMap<K, V, S>>>;
+
+    /// # Safety
+    ///
+    /// The index must not be out of bounds.
+    unsafe fn _try_yield_write_shard(&'a self, i: usize) -> Option<RwLockWriteGuard<'a, HashMap<K, V, S>>>;
+
     fn _insert(&self, key: K, value: V) -> Option<V>;
 
     fn _remove<Q>(&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>,
         Q: Hash + Eq + ?Sized;
 
+    fn _try_get<Q>(&'a self, key: &Q) -> TryResult<Ref<'a, K, V, S>>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq + ?Sized;
+
+    fn _try_get_mut<Q>(&'a self, key: &Q) -> TryResult<RefMut<'a, K, V, S>>
+    where
+        K: Borrow<Q>,
+        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<Entry<'a, K, V, S>>;
+
     fn _hasher(&self) -> S;
 
     // provided

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<R> {
+    /// 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<R> TryResult<R> {
+    /// 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<R> {
+        match self {
+            TryResult::Present(r) => Some(r),
+            _ => None,
+        }
+    }
+}