cleanup on drop for utils::mutex_map.

Signed-off-by: Jason Volk <jason@zemos.net>

src/core/utils/mutex_map.rs

@@ -1,20 +1,22 @@
-use std::{hash::Hash, sync::Arc};
+use std::{fmt::Debug, hash::Hash, sync::Arc};
 
-type Value<Val> = tokio::sync::Mutex<Val>;
-type ArcMutex<Val> = Arc<Value<Val>>;
-type HashMap<Key, Val> = std::collections::HashMap<Key, ArcMutex<Val>>;
-type MapMutex<Key, Val> = std::sync::Mutex<HashMap<Key, Val>>;
-type Map<Key, Val> = MapMutex<Key, Val>;
+use tokio::sync::OwnedMutexGuard as Omg;
 
 /// Map of Mutexes
 pub struct MutexMap<Key, Val> {
 	map: Map<Key, Val>,
 }
 
-pub struct Guard<Val> {
-	_guard: tokio::sync::OwnedMutexGuard<Val>,
+pub struct Guard<Key, Val> {
+	map: Map<Key, Val>,
+	val: Omg<Val>,
 }
 
+type Map<Key, Val> = Arc<MapMutex<Key, Val>>;
+type MapMutex<Key, Val> = std::sync::Mutex<HashMap<Key, Val>>;
+type HashMap<Key, Val> = std::collections::HashMap<Key, Value<Val>>;
+type Value<Val> = Arc<tokio::sync::Mutex<Val>>;
+
 impl<Key, Val> MutexMap<Key, Val>
 where
 	Key: Send + Hash + Eq + Clone,
@@ -23,28 +25,38 @@ where
 	#[must_use]
 	pub fn new() -> Self {
 		Self {
-			map: Map::<Key, Val>::new(HashMap::<Key, Val>::new()),
+			map: Map::new(MapMutex::new(HashMap::new())),
 		}
 	}
 
-	pub async fn lock<K>(&self, k: &K) -> Guard<Val>
+	#[tracing::instrument(skip(self), level = "debug")]
+	pub async fn lock<K>(&self, k: &K) -> Guard<Key, Val>
 	where
-		K: ?Sized + Send + Sync,
+		K: ?Sized + Send + Sync + Debug,
 		Key: for<'a> From<&'a K>,
 	{
 		let val = self
 			.map
 			.lock()
-			.expect("map mutex locked")
+			.expect("locked")
 			.entry(k.into())
 			.or_default()
 			.clone();
 
-		let guard = val.lock_owned().await;
-		Guard::<Val> {
-			_guard: guard,
+		Guard::<Key, Val> {
+			map: Arc::clone(&self.map),
+			val: val.lock_owned().await,
 		}
 	}
+
+	#[must_use]
+	pub fn contains(&self, k: &Key) -> bool { self.map.lock().expect("locked").contains_key(k) }
+
+	#[must_use]
+	pub fn is_empty(&self) -> bool { self.map.lock().expect("locked").is_empty() }
+
+	#[must_use]
+	pub fn len(&self) -> usize { self.map.lock().expect("locked").len() }
 }
 
 impl<Key, Val> Default for MutexMap<Key, Val>
@@ -54,3 +66,14 @@ where
 {
 	fn default() -> Self { Self::new() }
 }
+
+impl<Key, Val> Drop for Guard<Key, Val> {
+	fn drop(&mut self) {
+		if Arc::strong_count(Omg::mutex(&self.val)) <= 2 {
+			self.map
+				.lock()
+				.expect("locked")
+				.retain(|_, val| !Arc::ptr_eq(val, Omg::mutex(&self.val)) || Arc::strong_count(val) > 2);
+		}
+	}
+}