Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- use dashmap::DashMap;
- use std::{
- collections::BTreeSet,
- hash::Hash,
- mem::ManuallyDrop,
- ops::{Deref, DerefMut},
- sync::Arc,
- };
- use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
- pub trait Empty {
- fn is_empty(&self) -> bool;
- }
- impl Empty for () {
- fn is_empty(&self) -> bool {
- true
- }
- }
- impl<V> Empty for BTreeSet<V> {
- fn is_empty(&self) -> bool {
- BTreeSet::is_empty(self)
- }
- }
- type KeyLockInner<K, V> = Arc<DashMap<K, Arc<RwLock<V>>>>;
- #[derive(Clone)]
- pub struct KeyLocks<K, V>(KeyLockInner<K, V>);
- struct KeyRef<K, V> {
- map: KeyLockInner<K, V>,
- key: K,
- }
- impl<K: Eq + Hash, V: Empty> KeyRef<K, V> {
- fn remove_if_needed<G: Deref<Target = V>>(&self, guard: &mut ManuallyDrop<G>) {
- if guard.is_empty() {
- self.map.remove_if(&self.key, |_, v| {
- unsafe { ManuallyDrop::drop(guard) };
- Arc::strong_count(v) == 1
- });
- } else {
- unsafe { ManuallyDrop::drop(guard) };
- }
- }
- }
- pub struct OwnedReadGuard<K: Eq + Hash, V: Empty> {
- key: KeyRef<K, V>,
- guard: ManuallyDrop<OwnedRwLockReadGuard<V>>,
- }
- impl<K: Eq + Hash, V: Empty> Deref for OwnedReadGuard<K, V> {
- type Target = V;
- fn deref(&self) -> &Self::Target {
- self.guard.deref()
- }
- }
- impl<K: Eq + Hash, V: Empty> Drop for OwnedReadGuard<K, V> {
- fn drop(&mut self) {
- self.key.remove_if_needed(&mut self.guard);
- }
- }
- pub struct OwnedWriteGuard<K: Eq + Hash, V: Empty> {
- key: KeyRef<K, V>,
- guard: ManuallyDrop<OwnedRwLockWriteGuard<V>>,
- }
- impl<K: Eq + Hash, V: Empty> Deref for OwnedWriteGuard<K, V> {
- type Target = V;
- fn deref(&self) -> &Self::Target {
- self.guard.deref()
- }
- }
- impl<K: Eq + Hash, V: Empty> DerefMut for OwnedWriteGuard<K, V> {
- fn deref_mut(&mut self) -> &mut Self::Target {
- self.guard.deref_mut()
- }
- }
- impl<K: Eq + Hash, V: Empty> Drop for OwnedWriteGuard<K, V> {
- fn drop(&mut self) {
- self.key.remove_if_needed(&mut self.guard);
- }
- }
- impl<K: Eq + Hash + Clone, V: Empty + Default> KeyLocks<K, V> {
- pub fn new() -> Self {
- KeyLocks(Arc::default())
- }
- fn key_ref(&self, key: K) -> KeyRef<K, V> {
- KeyRef {
- map: self.0.clone(),
- key,
- }
- }
- pub async fn read(&self, key: K) -> OwnedReadGuard<K, V> {
- OwnedReadGuard {
- key: self.key_ref(key.clone()),
- guard: ManuallyDrop::new(
- self.0
- .entry(key)
- .or_insert_with(Arc::default)
- .clone()
- .read_owned()
- .await,
- ),
- }
- }
- pub async fn write(&self, key: K) -> OwnedWriteGuard<K, V> {
- OwnedWriteGuard {
- key: self.key_ref(key.clone()),
- guard: ManuallyDrop::new(
- self.0
- .entry(key)
- .or_insert_with(Arc::default)
- .clone()
- .write_owned()
- .await,
- ),
- }
- }
- pub fn len(&self) -> usize {
- self.0.len()
- }
- }
- #[cfg(test)]
- mod test {
- use super::KeyLocks;
- use std::collections::BTreeSet;
- #[tokio::test]
- async fn drop_only_if_empty() {
- let locks = KeyLocks::<u32, BTreeSet<String>>::new();
- let mut lock = locks.write(1).await;
- lock.insert("Hello".to_owned());
- lock.insert("World".to_owned());
- drop(lock);
- // Value is not empty and thus is not dropped
- assert_eq!(locks.len(), 1);
- let mut lock = locks.write(1).await;
- assert_eq!(lock.len(), 2);
- lock.clear();
- drop(lock);
- // Should be dropped now
- assert_eq!(locks.len(), 0);
- }
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement