unify lock logic

3 files changed, 8 insertions, 17 deletions

diff --git a/src/state/host.rs b/src/state/host.rs
index 863cc70..9b149b8 100644
--- a/src/state/host.rs
+++ b/src/state/host.rs
@@ -13,6 +13,7 @@ use tracing::debug;
 
 use super::{AppState, HelloState, Liveness, NodeLifecycle, PeerState, StaleView, Standing};
 use crate::document::replica::Epoch;
+use crate::locks::MutexExt;
 use crate::remake::{Host, PeerSummary, RemakeState};
 use crate::swarm::sync::{ChangeSet, SyncEvent};
 
@@ -195,14 +196,8 @@ impl Host for AppState {
     fn reset_initial_sync(&self) {
         self.initial_sync_done
             .store(false, std::sync::atomic::Ordering::Release);
-        *self
-            .catch_up_source
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
-        *self
-            .catch_up_preference
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
+        *self.catch_up_source.lock_or_recover() = None;
+        *self.catch_up_preference.lock_or_recover() = None;
     }
 
     fn subscribe_changes(&self) -> broadcast::Receiver<Arc<ChangeSet>> {
diff --git a/src/state/mod.rs b/src/state/mod.rs
index c5a6885..3336b00 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -30,6 +30,7 @@ use tokio::time::Instant;
 use tracing::{debug, error, warn};
 
 use crate::document::replica::{Epoch, Replica};
+use crate::locks::MutexExt;
 use crate::remake::{Host, RemakeState};
 use crate::swarm::sync::{ChangeSet, SyncEvent};
 use crate::swarm::throughput::ThroughputTable;
@@ -174,9 +175,7 @@ impl AppState {
     }
 
     pub(crate) fn lock_rtt(&self) -> std::sync::MutexGuard<'_, crate::swarm::rtt::RttTable> {
-        self.rtt
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.rtt.lock_or_recover()
     }
 
     /// The local epoch. Code that labels or applies sync messages must use the epoch
diff --git a/src/state/read.rs b/src/state/read.rs
index b76582c..f8323a2 100644
--- a/src/state/read.rs
+++ b/src/state/read.rs
@@ -6,6 +6,7 @@ use automerge::{AutoCommit, ChangeHash};
 
 use super::AppState;
 use crate::document::replica::Epoch;
+use crate::locks::RwLockExt;
 
 /// A read-only view of the document the node served before it started rejoining.
 pub(crate) struct StaleView {
@@ -46,8 +47,7 @@ impl AppState {
 
         let stale = self
             .stale
-            .read()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+            .read_or_recover()
             .as_ref()
             .map(|view| (view.epoch, Arc::clone(&view.doc)));
         let Some((epoch, doc)) = stale else {
@@ -92,9 +92,6 @@ impl AppState {
     }
 
     pub(super) fn set_stale(&self, view: Option<StaleView>) {
-        *self
-            .stale
-            .write()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = view;
+        *self.stale.write_or_recover() = view;
     }
 }