1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

#[cfg(test)]
pub mod tests;

use dashmap::mapref::entry::Entry;
use dashmap::DashMap;
use std::hash::Hash;
use std::mem::ManuallyDrop;
use std::ops::Deref;
use std::sync::Arc;

struct EntryInner<K: Clone + Eq + Hash, L: Default> {
  key: K,
  lock: L,
}

pub struct ArbitraryLockEntry<K: Clone + Eq + Hash, L: Default> {
  // 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.
  inner: ManuallyDrop<Arc<EntryInner<K, L>>>,
  map: ArbitraryLock<K, L>,
}

impl<K: Clone + Eq + Hash, L: Default> Drop for ArbitraryLockEntry<K, L> {
  fn drop(&mut self) {
    let key = self.inner.key.clone();
    // We must drop `inner` Arc before we attempt to drop the map's Arc.
    unsafe {
      ManuallyDrop::drop(&mut self.inner);
    };
    let Entry::Occupied(mut e) = self.map.map.entry(key) else {
      // 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.
      return;
    };
    let arc = e.get_mut().take().unwrap();
    // `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.
    match Arc::try_unwrap(arc) {
      Ok(_) => {
        // 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).
        e.remove();
      }
      Err(arc) => {
        // 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.
        *e.get_mut() = Some(arc);
      }
    };
  }
}

impl<K: Clone + Eq + Hash, L: Default> Deref for ArbitraryLockEntry<K, L> {
  type Target = L;

  fn deref(&self) -> &Self::Target {
    &self.inner.lock
  }
}

/// This can be cheaply cloned.
/// Provide your desired lock implementation for the `L` generic parameter. Options include `tokio::sync::RwLock`, `parking_lot::Mutex`, and `std::sync::Exclusive`.
pub struct ArbitraryLock<K: Clone + Eq + Hash, L: Default> {
  map: Arc<DashMap<K, Option<Arc<EntryInner<K, L>>>>>,
}

impl<K: Clone + Hash + Eq, L: Default> ArbitraryLock<K, L> {
  pub fn new() -> Self {
    Self {
      map: Default::default(),
    }
  }

  pub fn get(&self, k: K) -> ArbitraryLockEntry<K, L> {
    let inner = self
      .map
      .entry(k.clone())
      .or_insert_with(|| {
        Some(Arc::new(EntryInner {
          key: k,
          lock: L::default(),
        }))
      })
      .clone()
      // If the entry exists, it must be `Some`.
      .unwrap();
    ArbitraryLockEntry {
      inner: ManuallyDrop::new(inner),
      map: self.clone(),
    }
  }
}

impl<K: Clone + Hash + Eq, L: Default> Clone for ArbitraryLock<K, L> {
  fn clone(&self) -> Self {
    Self {
      map: self.map.clone(),
    }
  }
}