unify lock logic

1 files changed, 2 insertions, 3 deletions

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) {