refactor: Allow with_lock to return data and take an async closure

src/core/utils/with_lock.rs

@@ -1,89 +1,212 @@
 //! Traits for explicitly scoping the lifetime of locks.
 
-use std::sync::{Arc, Mutex};
+use std::{
+	future::Future,
+	sync::{Arc, Mutex},
+};
 
 pub trait WithLock<T: ?Sized> {
-	/// Acquires a lock and executes the given closure with the locked data.
-	fn with_lock<F>(&self, f: F)
+	/// Acquires a lock and executes the given closure with the locked data,
+	/// returning the result.
+	fn with_lock<R, F>(&self, f: F) -> R
 	where
-		F: FnMut(&mut T);
+		F: FnMut(&mut T) -> R;
 }
 
 impl<T> WithLock<T> for Mutex<T> {
-	fn with_lock<F>(&self, mut f: F)
+	fn with_lock<R, F>(&self, mut f: F) -> R
 	where
-		F: FnMut(&mut T),
+		F: FnMut(&mut T) -> R,
 	{
 		// The locking and unlocking logic is hidden inside this function.
 		let mut data_guard = self.lock().unwrap();
-		f(&mut data_guard);
+		f(&mut data_guard)
 		// Lock is released here when `data_guard` goes out of scope.
 	}
 }
 
 impl<T> WithLock<T> for Arc<Mutex<T>> {
-	fn with_lock<F>(&self, mut f: F)
+	fn with_lock<R, F>(&self, mut f: F) -> R
 	where
-		F: FnMut(&mut T),
+		F: FnMut(&mut T) -> R,
 	{
 		// The locking and unlocking logic is hidden inside this function.
 		let mut data_guard = self.lock().unwrap();
-		f(&mut data_guard);
+		f(&mut data_guard)
 		// Lock is released here when `data_guard` goes out of scope.
 	}
 }
 
 impl<R: lock_api::RawMutex, T: ?Sized> WithLock<T> for lock_api::Mutex<R, T> {
-	fn with_lock<F>(&self, mut f: F)
+	fn with_lock<Ret, F>(&self, mut f: F) -> Ret
 	where
-		F: FnMut(&mut T),
+		F: FnMut(&mut T) -> Ret,
 	{
 		// The locking and unlocking logic is hidden inside this function.
 		let mut data_guard = self.lock();
-		f(&mut data_guard);
+		f(&mut data_guard)
 		// Lock is released here when `data_guard` goes out of scope.
 	}
 }
 
 impl<R: lock_api::RawMutex, T: ?Sized> WithLock<T> for Arc<lock_api::Mutex<R, T>> {
-	fn with_lock<F>(&self, mut f: F)
+	fn with_lock<Ret, F>(&self, mut f: F) -> Ret
 	where
-		F: FnMut(&mut T),
+		F: FnMut(&mut T) -> Ret,
 	{
 		// The locking and unlocking logic is hidden inside this function.
 		let mut data_guard = self.lock();
-		f(&mut data_guard);
+		f(&mut data_guard)
 		// Lock is released here when `data_guard` goes out of scope.
 	}
 }
 
 pub trait WithLockAsync<T> {
-	/// Acquires a lock and executes the given closure with the locked data.
-	fn with_lock<F>(&self, f: F) -> impl Future<Output = ()>
+	/// Acquires a lock and executes the given closure with the locked data,
+	/// returning the result.
+	fn with_lock<R, F>(&self, f: F) -> impl Future<Output = R>
 	where
-		F: FnMut(&mut T);
+		F: FnMut(&mut T) -> R;
+
+	/// Acquires a lock and executes the given async closure with the locked
+	/// data.
+	fn with_lock_async<R, F>(&self, f: F) -> impl std::future::Future<Output = R>
+	where
+		F: AsyncFnMut(&mut T) -> R;
 }
 
 impl<T> WithLockAsync<T> for futures::lock::Mutex<T> {
-	async fn with_lock<F>(&self, mut f: F)
+	async fn with_lock<R, F>(&self, mut f: F) -> R
 	where
-		F: FnMut(&mut T),
+		F: FnMut(&mut T) -> R,
 	{
 		// The locking and unlocking logic is hidden inside this function.
 		let mut data_guard = self.lock().await;
-		f(&mut data_guard);
+		f(&mut data_guard)
+		// Lock is released here when `data_guard` goes out of scope.
+	}
+
+	async fn with_lock_async<R, F>(&self, mut f: F) -> R
+	where
+		F: AsyncFnMut(&mut T) -> R,
+	{
+		// The locking and unlocking logic is hidden inside this function.
+		let mut data_guard = self.lock().await;
+		f(&mut data_guard).await
 		// Lock is released here when `data_guard` goes out of scope.
 	}
 }
 
 impl<T> WithLockAsync<T> for Arc<futures::lock::Mutex<T>> {
-	async fn with_lock<F>(&self, mut f: F)
+	async fn with_lock<R, F>(&self, mut f: F) -> R
 	where
-		F: FnMut(&mut T),
+		F: FnMut(&mut T) -> R,
 	{
 		// The locking and unlocking logic is hidden inside this function.
 		let mut data_guard = self.lock().await;
-		f(&mut data_guard);
+		f(&mut data_guard)
+		// Lock is released here when `data_guard` goes out of scope.
+	}
+
+	async fn with_lock_async<R, F>(&self, mut f: F) -> R
+	where
+		F: AsyncFnMut(&mut T) -> R,
+	{
+		// The locking and unlocking logic is hidden inside this function.
+		let mut data_guard = self.lock().await;
+		f(&mut data_guard).await
 		// Lock is released here when `data_guard` goes out of scope.
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn test_with_lock_return_value() {
+		let mutex = Mutex::new(5);
+		let result = mutex.with_lock(|v| {
+			*v += 1;
+			*v * 2
+		});
+		assert_eq!(result, 12);
+		let value = mutex.lock().unwrap();
+		assert_eq!(*value, 6);
+	}
+
+	#[test]
+	fn test_with_lock_unit_return() {
+		let mutex = Mutex::new(10);
+		mutex.with_lock(|v| {
+			*v += 2;
+		});
+		let value = mutex.lock().unwrap();
+		assert_eq!(*value, 12);
+	}
+
+	#[test]
+	fn test_with_lock_arc_mutex() {
+		let mutex = Arc::new(Mutex::new(1));
+		let result = mutex.with_lock(|v| {
+			*v *= 10;
+			*v
+		});
+		assert_eq!(result, 10);
+		assert_eq!(*mutex.lock().unwrap(), 10);
+	}
+
+	#[tokio::test]
+	async fn test_with_lock_async_return_value() {
+		use futures::lock::Mutex as AsyncMutex;
+		let mutex = AsyncMutex::new(7);
+		let result = mutex
+			.with_lock(|v| {
+				*v += 3;
+				*v * 2
+			})
+			.await;
+		assert_eq!(result, 20);
+		let value = mutex.lock().await;
+		assert_eq!(*value, 10);
+	}
+
+	#[tokio::test]
+	async fn test_with_lock_async_unit_return() {
+		use futures::lock::Mutex as AsyncMutex;
+		let mutex = AsyncMutex::new(100);
+		mutex
+			.with_lock(|v| {
+				*v -= 50;
+			})
+			.await;
+		let value = mutex.lock().await;
+		assert_eq!(*value, 50);
+	}
+
+	#[tokio::test]
+	async fn test_with_lock_async_closure() {
+		use futures::lock::Mutex as AsyncMutex;
+		let mutex = AsyncMutex::new(1);
+		mutex
+			.with_lock_async(async |v| {
+				*v += 9;
+			})
+			.await;
+		let value = mutex.lock().await;
+		assert_eq!(*value, 10);
+	}
+
+	#[tokio::test]
+	async fn test_with_lock_async_arc_mutex() {
+		use futures::lock::Mutex as AsyncMutex;
+		let mutex = Arc::new(AsyncMutex::new(2));
+		mutex
+			.with_lock_async(async |v: &mut i32| {
+				*v *= 5;
+			})
+			.await;
+		let value = mutex.lock().await;
+		assert_eq!(*value, 10);
+	}
+}