Add callback lock

Author: seanaye

src/util.rs

@@ -22,6 +22,7 @@ pub mod progress;
 pub use mem_or_file::{FileAndSize, MemOrFile};
 mod sparse_mem_file;
 pub use sparse_mem_file::SparseMemFile;
+pub mod callback_lock;
 pub mod local_pool;
 
 #[cfg(test)]

src/util/callback_lock.rs

@@ -0,0 +1,67 @@
+//! This module defines a wrapper around a [`tokio::sync::RwLock`] that runs a callback
+//! After any write operation occurs
+
+use std::future::Future;
+
+#[derive(derive_more::Debug)]
+pub struct CallbackLock<T, F> {
+    inner: tokio::sync::RwLock<T>,
+    #[debug(skip)]
+    callback: F,
+}
+
+pub struct CallbackLockWriteGuard<'a, T, F: Fn(&T)> {
+    inner: tokio::sync::RwLockWriteGuard<'a, T>,
+    callback: &'a F,
+}
+
+impl<'a, T, F: Fn(&T)> std::ops::Deref for CallbackLockWriteGuard<'a, T, F> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &*self.inner
+    }
+}
+
+impl<'a, T, F: Fn(&T)> std::ops::DerefMut for CallbackLockWriteGuard<'a, T, F> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut *self.inner
+    }
+}
+
+impl<'a, T, F: Fn(&T)> Drop for CallbackLockWriteGuard<'a, T, F> {
+    fn drop(&mut self) {
+        (self.callback)(&*self.inner);
+    }
+}
+
+impl<T, F> CallbackLock<T, F>
+where
+    F: Fn(&T),
+{
+    pub fn new(val: T, callback: F) -> Self {
+        CallbackLock {
+            inner: tokio::sync::RwLock::new(val),
+            callback,
+        }
+    }
+
+    pub async fn write<'a>(&'a self) -> CallbackLockWriteGuard<'a, T, F> {
+        let guard = self.inner.write().await;
+
+        CallbackLockWriteGuard {
+            inner: guard,
+            callback: &self.callback,
+        }
+    }
+
+    pub fn read(&self) -> impl Future<Output = tokio::sync::RwLockReadGuard<'_, T>> {
+        self.inner.read()
+    }
+
+    pub fn try_read(
+        &self,
+    ) -> Result<tokio::sync::RwLockReadGuard<'_, T>, tokio::sync::TryLockError> {
+        self.inner.try_read()
+    }
+}