unify lock logic

10 files changed, 31 insertions, 70 deletions

diff --git a/src/swarm/auth.rs b/src/swarm/auth.rs
index 31230a5..1bca60d 100644
--- a/src/swarm/auth.rs
+++ b/src/swarm/auth.rs
@@ -19,6 +19,7 @@ use std::path::Path;
 use std::sync::Mutex;
 use std::time::{Duration, Instant};
 
+use crate::locks::MutexExt;
 use anyhow::{Context, Result, bail};
 use sha2::{Digest, Sha256};
 
@@ -67,10 +68,7 @@ impl SwarmKey {
     /// Records a rejected handshake. Returns how many went unlogged since the last log
     /// line if this one should be logged, None if it should be suppressed.
     pub(crate) fn note_rejection(&self) -> Option<u64> {
-        let mut rejections = self
-            .rejections
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let mut rejections = self.rejections.lock_or_recover();
         let due = rejections
             .last_logged
             .is_none_or(|last| last.elapsed() >= REJECTION_LOG_INTERVAL);
diff --git a/src/swarm/blob.rs b/src/swarm/blob.rs
index 42cfbf3..58522fa 100644
--- a/src/swarm/blob.rs
+++ b/src/swarm/blob.rs
@@ -29,6 +29,7 @@ use tokio::sync::{Semaphore, mpsc};
 use tokio::time::Instant;
 use tracing::{debug, info, warn};
 
+use crate::locks::MutexExt;
 use crate::state::AppState;
 use crate::transport::{BlobBody, BlobId, BlobMessage, ParsedMessage};
 
@@ -80,25 +81,19 @@ impl Blobs {
     fn fetches(
         &self,
     ) -> MutexGuard<'_, HashMap<BlobId, mpsc::UnboundedSender<(String, BlobBody)>>> {
-        self.fetches
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.fetches.lock_or_recover()
     }
 
     fn cached(&self, id: BlobId) -> Option<Arc<Prepared>> {
         self.cache
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+            .lock_or_recover()
             .as_ref()
             .filter(|(cached, _)| *cached == id)
             .map(|(_, prepared)| Arc::clone(prepared))
     }
 
     fn remember(&self, id: BlobId, prepared: Arc<Prepared>) {
-        *self
-            .cache
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some((id, prepared));
+        *self.cache.lock_or_recover() = Some((id, prepared));
     }
 }
 
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 651e067..ec6543b 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -14,6 +14,7 @@ use tracing::{debug, info, instrument, trace, warn};
 use crate::document::meta;
 use crate::document::replica::{Epoch, EpochRelation};
 use crate::document::store::{self, StoreError};
+use crate::locks::MutexExt;
 use crate::state::{AppState, HelloState, Liveness, PeerState, Standing};
 use crate::swarm::sync::{SyncEvent, handle_sync};
 use crate::swarm::topology;
@@ -543,8 +544,7 @@ pub(crate) async fn run_peer_worker(state: AppState) {
                 state.traffic.retain(|peer| known.contains_key(peer));
                 state
                     .throughput
-                    .lock()
-                    .unwrap_or_else(|poisoned| poisoned.into_inner())
+                    .lock_or_recover()
                     .retain(|peer| known.contains_key(peer));
             }
             // Fallback for what doesn't notify (`topology::run_worker`).
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index 703bc3f..5053be7 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -25,6 +25,7 @@ use std::time::Duration;
 use tokio::time::Instant;
 use tracing::debug;
 
+use crate::locks::MutexExt;
 use crate::state::{AppState, Standing};
 use crate::swarm::peer::send_ping;
 use crate::swarm::sync::SyncEvent;
@@ -45,10 +46,7 @@ pub(crate) async fn announce(state: &AppState, from: &str) {
         return;
     }
     let schedule = {
-        let mut pending = state
-            .announce_pending
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let mut pending = state.announce_pending.lock_or_recover();
         match pending.as_mut() {
             Some(senders) => {
                 senders.insert(from.to_owned());
@@ -68,8 +66,7 @@ pub(crate) async fn announce(state: &AppState, from: &str) {
         tokio::time::sleep(ANNOUNCE_DELAY).await;
         let senders = state
             .announce_pending
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+            .lock_or_recover()
             .take()
             .unwrap_or_default();
         announce_now(&state, &senders).await;
@@ -128,10 +125,7 @@ pub(crate) async fn pull_if_behind(state: &AppState, peer: &str, remote: &transp
 
     let now = Instant::now();
     {
-        let mut pulling = state
-            .pulling
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let mut pulling = state.pulling.lock_or_recover();
         let busy = matches!(
             &*pulling,
             Some((other, since)) if other != peer && now < *since + PULL_TIMEOUT
@@ -169,10 +163,7 @@ pub(crate) async fn pull_if_behind(state: &AppState, peer: &str, remote: &transp
 
 /// A sync message from `peer` was applied: a pull from it is answered.
 pub(crate) fn pulled(state: &AppState, peer: &str) {
-    let mut pulling = state
-        .pulling
-        .lock()
-        .unwrap_or_else(|poisoned| poisoned.into_inner());
+    let mut pulling = state.pulling.lock_or_recover();
     if pulling.as_ref().is_some_and(|(p, _)| p == peer) {
         *pulling = None;
     }
diff --git a/src/swarm/source.rs b/src/swarm/source.rs
index b1701b3..e7f06fb 100644
--- a/src/swarm/source.rs
+++ b/src/swarm/source.rs
@@ -35,6 +35,7 @@ use std::time::Duration;
 use tokio::time::Instant;
 use tracing::{debug, info, warn};
 
+use crate::locks::MutexExt;
 use crate::state::{AppState, CatchUpSource, PeerState};
 use crate::swarm::sync::SyncEvent;
 
@@ -55,17 +56,11 @@ const REPLY_TIMEOUTS: u32 = 3;
 const STALL_WARNING: Duration = Duration::from_secs(60);
 
 fn lock_preference(state: &AppState) -> std::sync::MutexGuard<'_, Option<Instant>> {
-    state
-        .catch_up_preference
-        .lock()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
+    state.catch_up_preference.lock_or_recover()
 }
 
 fn lock_avoided(state: &AppState) -> std::sync::MutexGuard<'_, Option<(String, Instant)>> {
-    state
-        .catch_up_avoided
-        .lock()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
+    state.catch_up_avoided.lock_or_recover()
 }
 
 fn catching_up(state: &AppState) -> bool {
@@ -73,10 +68,7 @@ fn catching_up(state: &AppState) -> bool {
 }
 
 fn lock(state: &AppState) -> std::sync::MutexGuard<'_, Option<CatchUpSource>> {
-    state
-        .catch_up_source
-        .lock()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
+    state.catch_up_source.lock_or_recover()
 }
 
 /// Returns how long to hold back sync with `peer`, if at all: while catching up, with
diff --git a/src/swarm/stability.rs b/src/swarm/stability.rs
index 9d79b17..1ee308f 100644
--- a/src/swarm/stability.rs
+++ b/src/swarm/stability.rs
@@ -22,6 +22,7 @@ use std::time::Duration;
 use tokio::time::Instant;
 use tracing::{info, warn};
 
+use crate::locks::MutexExt;
 use crate::state::{AppState, HelloState};
 use crate::swarm::members::{self, Tier};
 
@@ -129,10 +130,7 @@ impl Stability {
 }
 
 fn lock(state: &AppState) -> std::sync::MutexGuard<'_, Stability> {
-    state
-        .stability
-        .lock()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
+    state.stability.lock_or_recover()
 }
 
 /// The fraction of the observed time (up to a day) we had a live peer, and how much
diff --git a/src/swarm/state_file.rs b/src/swarm/state_file.rs
index 8b883fc..ba9192e 100644
--- a/src/swarm/state_file.rs
+++ b/src/swarm/state_file.rs
@@ -15,6 +15,7 @@ use tokio::time::Instant;
 use tracing::{error, info, warn};
 
 use crate::files::{FileKind, Files};
+use crate::locks::MutexExt;
 use crate::state::AppState;
 use crate::swarm::members;
 
@@ -90,9 +91,7 @@ impl StateFile {
     }
 
     fn lock(&self) -> MutexGuard<'_, Inner> {
-        self.inner
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.inner.lock_or_recover()
     }
 
     fn write(&self, remembered: &Remembered) -> std::io::Result<()> {
diff --git a/src/swarm/stats.rs b/src/swarm/stats.rs
index eb34276..f9a5b7b 100644
--- a/src/swarm/stats.rs
+++ b/src/swarm/stats.rs
@@ -9,6 +9,7 @@ use tokio::time::Instant;
 use tracing::{debug, warn};
 
 use crate::document::replica::Epoch;
+use crate::locks::MutexExt;
 use crate::state::AppState;
 
 /// Saving takes time proportional to the document; we spend at most about 1% of it on
@@ -120,9 +121,7 @@ impl Default for NodeStats {
 
 impl NodeStats {
     fn lock(&self) -> MutexGuard<'_, Inner> {
-        self.inner
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.inner.lock_or_recover()
     }
 
     pub(crate) fn caught_up(&self, sync_bytes: u64) {
diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs
index a05e346..fc6dc43 100644
--- a/src/swarm/sync.rs
+++ b/src/swarm/sync.rs
@@ -11,6 +11,7 @@ use tokio::sync::mpsc;
 use tokio::time::Instant;
 use tracing::{debug, info, trace, warn};
 
+use crate::locks::MutexExt;
 use crate::state::{AppState, Liveness, Sending, Standing};
 use crate::swarm::peer::{peer_swarm, send_hello_once};
 use crate::swarm::source;
@@ -355,8 +356,7 @@ pub(crate) async fn handle_sync(
             peer_state.last_sync_rx = Some(Instant::now());
             if let Some(wait) = &peer_state.sync.reply {
                 throughput
-                    .lock()
-                    .unwrap_or_else(|poisoned| poisoned.into_inner())
+                    .lock_or_recover()
                     .exchange(&peer, wait.started, arrived, rtt);
             }
             peer_state.sync.reply_settled();
diff --git a/src/swarm/topology.rs b/src/swarm/topology.rs
index 0dae7e9..73c2f2c 100644
--- a/src/swarm/topology.rs
+++ b/src/swarm/topology.rs
@@ -23,6 +23,7 @@ use sha2::{Digest, Sha256};
 use tokio::time::Instant;
 use tracing::info;
 
+use crate::locks::RwLockExt;
 use crate::state::AppState;
 use crate::swarm::members;
 
@@ -42,8 +43,7 @@ pub(crate) struct NeighbourView {
 pub(crate) fn is_full_mesh(state: &AppState) -> bool {
     state
         .neighbours
-        .read()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
+        .read_or_recover()
         .as_ref()
         .is_none_or(|view| view.full_mesh)
 }
@@ -65,8 +65,7 @@ pub(crate) fn summary(state: &AppState) -> Option<Summary> {
     let local = state.local_id();
     state
         .neighbours
-        .read()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
+        .read_or_recover()
         .as_ref()
         .map(|view| Summary {
             members: view.members.len(),
@@ -91,11 +90,7 @@ pub(crate) fn summary(state: &AppState) -> Option<Summary> {
 /// any member registered), so joining works as in a full mesh. We always *answer*
 /// anyone who contacts us.
 pub(crate) fn is_neighbour(state: &AppState, peer: &str) -> bool {
-    match &*state
-        .neighbours
-        .read()
-        .unwrap_or_else(|poisoned| poisoned.into_inner())
-    {
+    match &*state.neighbours.read_or_recover() {
         Some(view) => view.neighbours.contains(peer),
         None => true,
     }
@@ -154,10 +149,7 @@ pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
         }
     }
     let previous = {
-        let view = state
-            .neighbours
-            .read()
-            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let view = state.neighbours.read_or_recover();
         if view.as_ref().is_some_and(|view| {
             view.members == current
                 && view.k == k
@@ -207,10 +199,7 @@ pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
         .iter()
         .filter(|m| m.as_str() != state.local_id() && !unreachable.contains(*m))
         .count();
-    *state
-        .neighbours
-        .write()
-        .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(NeighbourView {
+    *state.neighbours.write_or_recover() = Some(NeighbourView {
         full_mesh: computed.len() >= reachable_others,
         members: current,
         k,