3 files changed, 7 insertions, 13 deletions
diff --git a/src/transport/firewall.rs b/src/transport/firewall.rs
index b8e05cd..c8592bd 100644
--- a/src/transport/firewall.rs
+++ b/src/transport/firewall.rs
@@ -6,6 +6,7 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use utoipa::ToSchema;
 
+use crate::locks::RwLockExt;
 use crate::transport::EncodedMessage;
 use crate::transport::{MessageHandler, ParsedMessage, Transport};
 
@@ -27,9 +28,7 @@ pub struct Firewall {
 
 impl Firewall {
     fn read(&self) -> RwLockReadGuard<'_, FirewallRules> {
-        self.rules
-            .read()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.rules.read_or_recover()
     }
 
     pub fn rules(&self) -> FirewallRules {
@@ -37,10 +36,7 @@ impl Firewall {
     }
 
     pub fn set_rules(&self, rules: FirewallRules) {
-        *self
-            .rules
-            .write()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = rules;
+        *self.rules.write_or_recover() = rules;
     }
 
     pub fn blocks_inbound(&self, peer: &str) -> bool {
diff --git a/src/transport/sim.rs b/src/transport/sim.rs
index ce5977f..675f382 100644
--- a/src/transport/sim.rs
+++ b/src/transport/sim.rs
@@ -26,6 +26,7 @@ use tracing::{debug, warn};
 
 use super::mycelium::{decode_message, encode_message};
 use super::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};
+use crate::locks::MutexExt;
 use crate::testrand::Rng;
 
 /// How one direction of a link behaves.
@@ -102,9 +103,7 @@ impl SimNet {
     }
 
     fn lock(&self) -> MutexGuard<'_, NetInner> {
-        self.inner
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.inner.lock_or_recover()
     }
 
     /// A transport for a new node `id` on this net.
diff --git a/src/transport/traffic.rs b/src/transport/traffic.rs
index 9dcd91a..3d17860 100644
--- a/src/transport/traffic.rs
+++ b/src/transport/traffic.rs
@@ -16,6 +16,7 @@ use async_trait::async_trait;
 use serde::Serialize;
 use utoipa::ToSchema;
 
+use crate::locks::MutexExt;
 use crate::transport::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};
 
 #[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
@@ -79,9 +80,7 @@ pub struct Traffic {
 
 impl Traffic {
     fn lock(&self) -> MutexGuard<'_, Table> {
-        self.table
-            .lock()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.table.lock_or_recover()
     }
 
     fn count(&self, peer: &str, bytes: usize, pick: impl Fn(&mut PeerTraffic) -> &mut Counter) {