refactor(transport): replace watch channel with rwlock

Author: mempirate

msg-socket/src/req/driver.rs

@@ -279,8 +279,7 @@ where
                     if let Ok(io) = result {
                         tracing::debug!(target = ?io.peer_addr(), "new connection");
 
-                        let tx = this.socket_state.transport.0.clone();
-                        let metered = MeteredIo::new(io, tx);
+                        let metered = MeteredIo::new(io, Arc::clone(&this.socket_state.transport));
 
                         let mut framed = Framed::new(metered, reqrep::Codec::new());
                         framed.set_backpressure_boundary(this.options.backpressure_boundary);

msg-socket/src/req/mod.rs

@@ -1,7 +1,10 @@
 use bytes::Bytes;
-use std::{sync::Arc, time::Duration};
+use std::{
+    sync::{Arc, RwLock},
+    time::Duration,
+};
 use thiserror::Error;
-use tokio::sync::{oneshot, watch};
+use tokio::sync::oneshot;
 
 use msg_wire::{
     compression::{CompressionType, Compressor},
@@ -185,8 +188,7 @@ impl ReqMessage {
 pub(crate) struct SocketState<S> {
     /// The socket stats.
     pub(crate) stats: Arc<SocketStats<ReqStats>>,
-    /// The transport stats. This is None until a connection is established.
-    pub(crate) transport: (watch::Sender<Arc<S>>, watch::Receiver<Arc<S>>),
+    pub(crate) transport: Arc<RwLock<Arc<S>>>,
 }
 
 // Manual clone implementation needed here because `S` is n`.

msg-socket/src/req/socket.rs

@@ -3,10 +3,7 @@ use rustc_hash::FxHashMap;
 use std::{marker::PhantomData, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
 use tokio::{
     net::{ToSocketAddrs, lookup_host},
-    sync::{
-        mpsc, oneshot,
-        watch::{self},
-    },
+    sync::{mpsc, oneshot},
 };
 
 use msg_transport::{Address, Transport};
@@ -76,7 +73,7 @@ where
             options: Arc::new(options),
             state: SocketState {
                 stats: Arc::new(SocketStats::default()),
-                transport: watch::channel(Arc::new(T::Stats::default())),
+                transport: Default::default(),
             },
             compressor: None,
             _marker: PhantomData,
@@ -96,9 +93,7 @@ where
 
     /// Borrow the latest transport-level stats snapshot.
     pub fn transport_stats(&self) -> Arc<T::Stats> {
-        // NOTE: We clone the Arc here because purely borrowing the inner stats
-        // would lock the channel.
-        self.state.transport.1.borrow().clone()
+        Arc::clone(&self.state.transport.read().unwrap())
     }
 
     pub async fn request(&self, message: Bytes) -> Result<Bytes, ReqError> {

msg-transport/src/lib.rs

@@ -10,17 +10,14 @@ use std::{
     net::SocketAddr,
     path::PathBuf,
     pin::Pin,
-    sync::Arc,
+    sync::{Arc, RwLock},
     task::{Context, Poll},
     time::{Duration, Instant},
 };
 
 use async_trait::async_trait;
 use futures::{Future, FutureExt};
-use tokio::{
-    io::{AsyncRead, AsyncWrite},
-    sync::watch,
-};
+use tokio::io::{AsyncRead, AsyncWrite};
 
 pub mod ipc;
 #[cfg(feature = "quic")]
@@ -47,7 +44,7 @@ where
     /// The inner IO object.
     inner: Io,
     /// The sender for the stats.
-    sender: watch::Sender<Arc<S>>,
+    stats: Arc<RwLock<Arc<S>>>,
     /// The next time the stats should be refreshed.
     next_refresh: Instant,
     /// The interval at which the stats should be refreshed.
@@ -130,10 +127,10 @@ where
     /// stats. The `sender` is used to send the latest stats to the caller.
     ///
     /// TODO: Specify configuration options.
-    pub fn new(inner: Io, sender: watch::Sender<Arc<S>>) -> Self {
+    pub fn new(inner: Io, stats: Arc<RwLock<Arc<S>>>) -> Self {
         Self {
             inner,
-            sender,
+            stats,
             _marker: PhantomData,
             next_refresh: Instant::now(),
             refresh_interval: Duration::from_secs(2),
@@ -146,9 +143,7 @@ where
         if self.next_refresh <= now {
             match S::try_from(&self.inner) {
                 Ok(stats) => {
-                    if let Err(e) = self.sender.send(Arc::new(stats)) {
-                        tracing::error!(err = ?e, "failed to update transport stats");
-                    }
+                    *self.stats.write().unwrap() = Arc::new(stats);
                 }
                 Err(e) => tracing::error!(errror = ?e, "failed to gather transport stats"),
             }