Deduplicate a lot of code

Author: lithp

trin-core/src/lib.rs

@@ -10,38 +10,28 @@ pub mod utils;
 pub const ACQUIRE_TIMEOUT_MS: u64 = 100;
 pub const HOLD_TIMEOUT_MS: u64 = 100;
 
-pub struct TimedRwReadGuard<'a, T: ?Sized> {
-    pub inner: tokio::sync::RwLockReadGuard<'a, T>,
+pub struct TimedGuard<T> {
+    pub inner: T,
     pub acquisition_line: u32,
     pub acquisition_file: &'static str,
     pub acquisition_time: std::time::Instant,
     pub sleep_task: tokio::task::JoinHandle<()>,
 }
 
-async fn sleep_then_log(file: &'static str, line: u32) {
-    tokio::time::sleep(std::time::Duration::from_millis(HOLD_TIMEOUT_MS)).await;
-    log::warn!(
-        "[{}:{}] lock held for over {}ms, not yet released",
-        file,
-        line,
-        HOLD_TIMEOUT_MS.to_string()
-    );
-}
-
-impl<'a, T: ?Sized> TimedRwReadGuard<'a, T> {
+impl<T> TimedGuard<T> {
     pub async fn new(
-        inner: tokio::sync::RwLockReadGuard<'a, T>,
+        inner: T,
         acquisition_line: u32,
         acquisition_file: &'static str,
-    ) -> TimedRwReadGuard<'a, T> {
+    ) -> TimedGuard<T> {
         let now = std::time::Instant::now();
         let move_line = acquisition_line;
         let move_file = acquisition_file;
         let handle = tokio::spawn(async move {
             sleep_then_log(move_file, move_line).await;
         });
 
-        TimedRwReadGuard {
+        TimedGuard {
             inner,
             acquisition_line,
             acquisition_file,
@@ -51,15 +41,21 @@ impl<'a, T: ?Sized> TimedRwReadGuard<'a, T> {
     }
 }
 
-impl<'a, T> std::ops::Deref for TimedRwReadGuard<'a, T> {
-    type Target = tokio::sync::RwLockReadGuard<'a, T>;
+impl<T> std::ops::Deref for TimedGuard<T> {
+    type Target = T;
 
     fn deref(&self) -> &Self::Target {
         &self.inner
     }
 }
 
-impl<'a, T: ?Sized> Drop for TimedRwReadGuard<'a, T> {
+impl<T> std::ops::DerefMut for TimedGuard<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl<T> Drop for TimedGuard<T> {
     fn drop(&mut self) {
         self.sleep_task.abort();
         let held_for = self.acquisition_time.elapsed().as_millis();
@@ -74,9 +70,19 @@ impl<'a, T: ?Sized> Drop for TimedRwReadGuard<'a, T> {
     }
 }
 
+async fn sleep_then_log(file: &'static str, line: u32) {
+    tokio::time::sleep(std::time::Duration::from_millis(HOLD_TIMEOUT_MS)).await;
+    log::warn!(
+        "[{}:{}] lock held for over {}ms, not yet released",
+        file,
+        line,
+        HOLD_TIMEOUT_MS.to_string()
+    );
+}
+
 #[macro_export]
-macro_rules! rw_read {
-    ($lock_name:expr) => {{
+macro_rules! inner_rw {
+    ($meth: ident, $lock_name:expr) => {{
         let acquire_timeout = std::time::Duration::from_millis($crate::ACQUIRE_TIMEOUT_MS);
         let sleep = tokio::time::sleep(acquire_timeout);
         tokio::pin!(sleep);
@@ -95,7 +101,7 @@ macro_rules! rw_read {
                     );
                     did_log = true;
                 }
-                guard = $lock_name.read() => {
+                guard = $lock_name.$meth() => {
                     if did_log {
                         let wait_time = now.elapsed().as_millis();
                         log::warn!(
@@ -106,7 +112,7 @@ macro_rules! rw_read {
                         );
                     }
 
-                    let wrapped = $crate::TimedRwReadGuard::new(
+                    let wrapped = $crate::TimedGuard::new(
                         guard, std::line!(), std::file!(),
                     ).await;
                     break wrapped;
@@ -119,24 +125,13 @@ macro_rules! rw_read {
 #[macro_export]
 macro_rules! rw_write {
     ($lock_name:expr) => {{
-        let acquire_timeout = std::time::Duration::from_millis($crate::ACQUIRE_TIMEOUT_MS);
-        let sleep = tokio::time::sleep(acquire_timeout);
-        tokio::pin!(sleep);
+        $crate::inner_rw!(write, $lock_name)
+    }}
+}
 
-        loop {
-            tokio::select! {
-                _ = &mut sleep, if !sleep.is_elapsed() => {
-                    log::warn!(
-                        "[{}:{}] took more than {} ms to acquire lock",
-                        std::file!(),
-                        std::line!(),
-                        $crate::ACQUIRE_TIMEOUT_MS,
-                    );
-                }
-                guard = $lock_name.write() => {
-                    break guard;
-                }
-            }
-        }
-    }};
+#[macro_export]
+macro_rules! rw_read {
+    ($lock_name:expr) => {{
+        $crate::inner_rw!(read, $lock_name)
+    }}
 }