feat: add try_reserve method (#221)

src/lib.rs

@@ -51,6 +51,14 @@ cfg_if! {
 
 pub(crate) type HashMap<K, V, S> = hashbrown::HashMap<K, SharedValue<V>, S>;
 
+// Temporary reimplementation of [`std::collections::TryReserveError`]
+// util [`std::collections::TryReserveError`] stabilises.
+// We cannot easily create `std::collections` error type from `hashbrown` error type
+// without access to `TryReserveError::kind` method.
+#[non_exhaustive]
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct TryReserveError {}
+
 fn default_shard_amount() -> usize {
     (std::thread::available_parallelism().map_or(1, usize::from) * 4).next_power_of_two()
 }
@@ -216,7 +224,7 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: BuildHasher + Clone> DashMap<K, V, S> {
 
     /// Creates a new DashMap with a specified hasher and shard amount
     ///
-    /// shard_amount should greater than 0 and be a power of two.
+    /// shard_amount should be greater than 0 and a power of two.
     /// If a shard_amount which is not a power of two is provided, the function will panic.
     ///
     /// # Examples
@@ -797,6 +805,21 @@ impl<'a, K: 'a + Eq + Hash, V: 'a, S: BuildHasher + Clone> DashMap<K, V, S> {
     pub fn try_entry(&'a self, key: K) -> Option<Entry<'a, K, V, S>> {
         self._try_entry(key)
     }
+
+    /// Advanced entry API that tries to mimic `std::collections::HashMap::try_reserve`.
+    /// Tries to reserve capacity for at least `shard * additional`
+    /// and may reserve more space to avoid frequent reallocations.
+    ///
+    /// # Errors
+    ///
+    /// If the capacity overflows, or the allocator reports a failure, then an error is returned.
+    // TODO: return std::collections::TryReserveError once std::collections::TryReserveErrorKind stabilises.
+    pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        for shard in self.shards.iter() {
+            shard.write().try_reserve(additional).map_err(|_| TryReserveError {})?;
+        }
+        Ok(())
+    }
 }
 
 impl<'a, K: 'a + Eq + Hash, V: 'a, S: 'a + BuildHasher + Clone> Map<'a, K, V, S>
@@ -1367,4 +1390,26 @@ mod tests {
             assert!(result2.is_locked());
         }
     }
+
+    #[test]
+    fn test_try_reserve() {
+        let mut map: DashMap<i32, i32> = DashMap::new();
+        // DashMap is empty and doesn't allocate memory
+        assert_eq!(map.capacity(), 0);
+
+        map.try_reserve(10).unwrap();
+
+        // And now map can hold at least 10 elements
+        assert!(map.capacity() >= 10);
+    }
+
+    #[test]
+    fn test_try_reserve_errors() {
+        let mut map: DashMap<i32, i32> = DashMap::new();
+
+        match map.try_reserve(usize::MAX) {
+            Err(_) => {}
+            _ => panic!("should have raised CapacityOverflow error"),
+        }
+    }
 }