arbitrary_lock/
lib.rs

1#[cfg(test)]
2pub mod tests;
3
4use dashmap::mapref::entry::Entry;
5use dashmap::DashMap;
6use std::hash::Hash;
7use std::mem::ManuallyDrop;
8use std::ops::Deref;
9use std::sync::Arc;
10
11struct EntryInner<K: Clone + Eq + Hash, L: Default> {
12  key: K,
13  lock: L,
14}
15
16pub struct ArbitraryLockEntry<K: Clone + Eq + Hash, L: Default> {
17  // We use ManuallyDrop because we need to drop this before we drop the ArbitraryLockEntry (usually fields are dropped after the containing struct). See the Drop impl for more details. Idea by https://stackoverflow.com/a/41056727/6249022.
18  inner: ManuallyDrop<Arc<EntryInner<K, L>>>,
19  map: ArbitraryLock<K, L>,
20}
21
22impl<K: Clone + Eq + Hash, L: Default> Drop for ArbitraryLockEntry<K, L> {
23  fn drop(&mut self) {
24    let key = self.inner.key.clone();
25    // We must drop `inner` Arc before we attempt to drop the map's Arc.
26    unsafe {
27      ManuallyDrop::drop(&mut self.inner);
28    };
29    let Entry::Occupied(mut e) = self.map.map.entry(key) else {
30      // Already returned. This is possible because some other thread may have dropped the map entry after we droppepd our `inner` Arc but before we were able to acquire a lock on the map entry.
31      return;
32    };
33    let arc = e.get_mut().take().unwrap();
34    // `try_unwrap` is the appropriate function, not `Arc::into_inner`; there is no race condition because we're holdding onto the map entry lock, and `into_inner` will drop the map's Arc even if it isn't actually the last reference which is not what we want.
35    match Arc::try_unwrap(arc) {
36      Ok(_) => {
37        // We dropped the map's Arc, which means there is no other reference anywhere and we must remove the map entry (and not leave it as a None and leak memory).
38        e.remove();
39      }
40      Err(arc) => {
41        // We couldn't drop the map's Arc, so do not delete the map entry and restore its value to a Some containing the Arc.
42        *e.get_mut() = Some(arc);
43      }
44    };
45  }
46}
47
48impl<K: Clone + Eq + Hash, L: Default> Deref for ArbitraryLockEntry<K, L> {
49  type Target = L;
50
51  fn deref(&self) -> &Self::Target {
52    &self.inner.lock
53  }
54}
55
56/// This can be cheaply cloned.
57/// Provide your desired lock implementation for the `L` generic parameter. Options include `tokio::sync::RwLock`, `parking_lot::Mutex`, and `std::sync::Exclusive`.
58pub struct ArbitraryLock<K: Clone + Eq + Hash, L: Default> {
59  map: Arc<DashMap<K, Option<Arc<EntryInner<K, L>>>, ahash::RandomState>>,
60}
61
62impl<K: Clone + Hash + Eq, L: Default> ArbitraryLock<K, L> {
63  pub fn new() -> Self {
64    Self {
65      map: Default::default(),
66    }
67  }
68
69  pub fn get(&self, k: K) -> ArbitraryLockEntry<K, L> {
70    let inner = self
71      .map
72      .entry(k.clone())
73      .or_insert_with(|| {
74        Some(Arc::new(EntryInner {
75          key: k,
76          lock: L::default(),
77        }))
78      })
79      .clone()
80      // If the entry exists, it must be `Some`.
81      .unwrap();
82    ArbitraryLockEntry {
83      inner: ManuallyDrop::new(inner),
84      map: self.clone(),
85    }
86  }
87
88  pub fn len(&self) -> usize {
89    self.map.len()
90  }
91
92  pub fn is_empty(&self) -> bool {
93    self.map.is_empty()
94  }
95
96  pub fn capacity(&self) -> usize {
97    self.map.capacity()
98  }
99
100  pub fn shrink_to_fit(&self) {
101    self.map.shrink_to_fit();
102  }
103}
104
105impl<K: Clone + Hash + Eq, L: Default> Clone for ArbitraryLock<K, L> {
106  fn clone(&self) -> Self {
107    Self {
108      map: self.map.clone(),
109    }
110  }
111}
arbitrary_lock/lib.rs

arbitrary_lock/
lib.rs
