From 8b68d6306c33639076c6ee3d543fe2fc081e082f Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Fri, 14 Jun 2024 21:11:31 +0000
Subject: [PATCH] add MutexMap to utils

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/core/utils/mod.rs       |  2 ++
 src/core/utils/mutex_map.rs | 56 +++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 src/core/utils/mutex_map.rs

diff --git a/src/core/utils/mod.rs b/src/core/utils/mod.rs
index 9ffbbbd0..27d47fa5 100644
--- a/src/core/utils/mod.rs
+++ b/src/core/utils/mod.rs
@@ -4,6 +4,7 @@ pub mod defer;
 pub mod hash;
 pub mod html;
 pub mod json;
+pub mod mutex_map;
 pub mod sys;
 
 use std::{
@@ -14,6 +15,7 @@ use std::{
 pub use debug::slice_truncated as debug_slice_truncated;
 pub use html::Escape as HtmlEscape;
 pub use json::{deserialize_from_str, to_canonical_object};
+pub use mutex_map::MutexMap;
 use rand::prelude::*;
 use ring::digest;
 use ruma::OwnedUserId;
diff --git a/src/core/utils/mutex_map.rs b/src/core/utils/mutex_map.rs
new file mode 100644
index 00000000..f102487c
--- /dev/null
+++ b/src/core/utils/mutex_map.rs
@@ -0,0 +1,56 @@
+use std::{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>;
+
+/// Map of Mutexes
+pub struct MutexMap<Key, Val> {
+	map: Map<Key, Val>,
+}
+
+pub struct Guard<Val> {
+	_guard: tokio::sync::OwnedMutexGuard<Val>,
+}
+
+impl<Key, Val> MutexMap<Key, Val>
+where
+	Key: Send + Hash + Eq + Clone,
+	Val: Send + Default,
+{
+	#[must_use]
+	pub fn new() -> Self {
+		Self {
+			map: Map::<Key, Val>::new(HashMap::<Key, Val>::new()),
+		}
+	}
+
+	pub async fn lock<K>(&self, k: &K) -> Guard<Val>
+	where
+		K: ?Sized + Send + Sync,
+		Key: for<'a> From<&'a K>,
+	{
+		let val = self
+			.map
+			.lock()
+			.expect("map mutex locked")
+			.entry(k.into())
+			.or_default()
+			.clone();
+
+		let guard = val.lock_owned().await;
+		Guard::<Val> {
+			_guard: guard,
+		}
+	}
+}
+
+impl<Key, Val> Default for MutexMap<Key, Val>
+where
+	Key: Send + Hash + Eq + Clone,
+	Val: Send + Default,
+{
+	fn default() -> Self { Self::new() }
+}