unify lock logic

2 files changed, 5 insertions, 11 deletions

diff --git a/src/remake/local.rs b/src/remake/local.rs
index 517367b..8e9d8ba 100644
--- a/src/remake/local.rs
+++ b/src/remake/local.rs
@@ -16,6 +16,7 @@ use tokio::time::Instant;
 use super::records::Retention;
 use super::snapshot::Snapshot;
 use crate::document::replica::Epoch;
+use crate::locks::MutexExt;
 
 pub(crate) struct RemakeState {
     /// When this node started (Unix milliseconds): votes it cast before may be lost
@@ -185,9 +186,7 @@ impl RemakeState {
     }
 
     fn lock(&self) -> MutexGuard<'_, Inner> {
-        self.inner
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.inner.lock_or_recover()
     }
 
     pub(crate) fn store_prebuilt(&self, id: &str, snapshot: Snapshot) {
diff --git a/src/remake/switch.rs b/src/remake/switch.rs
index 65f296f..903baea 100644
--- a/src/remake/switch.rs
+++ b/src/remake/switch.rs
@@ -38,6 +38,7 @@ use super::vote;
 use crate::document::meta::{self, REMAKES};
 use crate::document::replica::Epoch;
 use crate::document::store;
+use crate::locks::MutexExt;
 use crate::swarm::sync::paths_related;
 
 const TICK: Duration = Duration::from_secs(1);
@@ -311,16 +312,10 @@ pub(super) async fn swap_in<H: Host>(
                     "Replayed late changes"
                 );
             }
-            *late_out
-                .lock()
-                .unwrap_or_else(|poisoned| poisoned.into_inner()) = late;
+            *late_out.lock_or_recover() = late;
             Ok(new_doc)
         })
         .await?;
-    let late = std::mem::take(
-        &mut *late_stats
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()),
-    );
+    let late = std::mem::take(&mut *late_stats.lock_or_recover());
     Ok((old, late))
 }