refactor channel metrics

common/channel/src/channel.rs

@@ -1,9 +1,10 @@
 use crate::error::Error;
-use metrics::{register_meter_with_group, Counter, CounterUsize, Histogram, Meter, Sample};
+use crate::metrics::unbounded_channel;
+use metrics::{Counter, CounterUsize};
 use std::sync::Arc;
-use std::time::{Duration, Instant};
-use tokio::sync::mpsc::error::{SendError, TryRecvError};
-use tokio::sync::{mpsc, oneshot};
+use std::time::Duration;
+use tokio::sync::mpsc::error::TryRecvError;
+use tokio::sync::oneshot;
 use tokio::time::timeout;
 
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(3);
@@ -22,63 +23,21 @@ pub struct Channel<N, Req, Res> {
 
 impl<N, Req, Res> Channel<N, Req, Res> {
     pub fn unbounded(name: &str) -> (Sender<N, Req, Res>, Receiver<N, Req, Res>) {
-        let (sender, receiver) = mpsc::unbounded_channel();
-
         let metrics_group = format!("common_channel_{}", name);
-        let metrics_queued = CounterUsize::register_with_group(metrics_group.as_str(), "size");
-
+        let (sender, receiver) = unbounded_channel(metrics_group.as_str());
+        let metrics_timeout = CounterUsize::register_with_group(metrics_group.as_str(), "timeout");
         (
             Sender {
                 chan: sender,
-                metrics_send_qps: register_meter_with_group(metrics_group.as_str(), "send"),
-                metrics_queued: metrics_queued.clone(),
-                metrics_timeout: CounterUsize::register_with_group(
-                    metrics_group.as_str(),
-                    "timeout",
-                ),
-            },
-            Receiver {
-                chan: receiver,
-                metrics_recv_qps: register_meter_with_group(metrics_group.as_str(), "recv"),
-                metrics_queued,
-                metrics_queue_latency: Sample::ExpDecay(0.015).register_with_group(
-                    metrics_group.as_str(),
-                    "latency",
-                    1024,
-                ),
+                metrics_timeout,
             },
+            Receiver { chan: receiver },
         )
     }
 }
 
-enum TimedMessage<N, Req, Res> {
-    Notification(Instant, N),
-    Request(Instant, Req, ResponseSender<Res>),
-}
-
-impl<N, Req, Res> From<Message<N, Req, Res>> for TimedMessage<N, Req, Res> {
-    fn from(value: Message<N, Req, Res>) -> Self {
-        match value {
-            Message::Notification(n) => TimedMessage::Notification(Instant::now(), n),
-            Message::Request(req, res) => TimedMessage::Request(Instant::now(), req, res),
-        }
-    }
-}
-
-impl<N, Req, Res> TimedMessage<N, Req, Res> {
-    fn into_message(self) -> (Instant, Message<N, Req, Res>) {
-        match self {
-            TimedMessage::Notification(since, n) => (since, Message::Notification(n)),
-            TimedMessage::Request(since, req, res) => (since, Message::Request(req, res)),
-        }
-    }
-}
-
 pub struct Sender<N, Req, Res> {
-    chan: mpsc::UnboundedSender<TimedMessage<N, Req, Res>>,
-
-    metrics_send_qps: Arc<dyn Meter>,
-    metrics_queued: Arc<dyn Counter<usize>>,
+    chan: crate::metrics::Sender<Message<N, Req, Res>>,
     metrics_timeout: Arc<dyn Counter<usize>>,
 }
 
@@ -86,8 +45,6 @@ impl<N, Req, Res> Clone for Sender<N, Req, Res> {
     fn clone(&self) -> Self {
         Sender {
             chan: self.chan.clone(),
-            metrics_send_qps: self.metrics_send_qps.clone(),
-            metrics_queued: self.metrics_queued.clone(),
             metrics_timeout: self.metrics_timeout.clone(),
         }
     }
@@ -95,13 +52,17 @@ impl<N, Req, Res> Clone for Sender<N, Req, Res> {
 
 impl<N, Req, Res> Sender<N, Req, Res> {
     pub fn notify(&self, msg: N) -> Result<(), Error<N, Req, Res>> {
-        self.send(Message::Notification(msg))
+        self.chan
+            .send(Message::Notification(msg))
+            .map_err(|e| Error::SendError(e))
     }
 
     pub async fn request(&self, request: Req) -> Result<Res, Error<N, Req, Res>> {
         let (sender, receiver) = oneshot::channel();
 
-        self.send(Message::Request(request, sender))?;
+        self.chan
+            .send(Message::Request(request, sender))
+            .map_err(|e| Error::SendError(e))?;
 
         timeout(DEFAULT_REQUEST_TIMEOUT, receiver)
             .await
@@ -111,49 +72,19 @@ impl<N, Req, Res> Sender<N, Req, Res> {
             })?
             .map_err(|e| Error::RecvError(e))
     }
-
-    fn send(&self, message: Message<N, Req, Res>) -> Result<(), Error<N, Req, Res>> {
-        match self.chan.send(message.into()) {
-            Ok(()) => {
-                self.metrics_send_qps.mark(1);
-                self.metrics_queued.inc(1);
-                Ok(())
-            }
-            Err(e) => {
-                let (_, msg) = e.0.into_message();
-                Err(Error::SendError(SendError(msg)))
-            }
-        }
-    }
 }
 
 pub struct Receiver<N, Req, Res> {
-    chan: mpsc::UnboundedReceiver<TimedMessage<N, Req, Res>>,
-
-    metrics_recv_qps: Arc<dyn Meter>,
-    metrics_queued: Arc<dyn Counter<usize>>,
-    metrics_queue_latency: Arc<dyn Histogram>,
+    chan: crate::metrics::Receiver<Message<N, Req, Res>>,
 }
 
 impl<N, Req, Res> Receiver<N, Req, Res> {
     pub async fn recv(&mut self) -> Option<Message<N, Req, Res>> {
-        let data = self.chan.recv().await?;
-        Some(self.on_recv_data(data))
+        self.chan.recv().await
     }
 
     pub fn try_recv(&mut self) -> Result<Message<N, Req, Res>, TryRecvError> {
-        let data = self.chan.try_recv()?;
-        Ok(self.on_recv_data(data))
-    }
-
-    fn on_recv_data(&self, data: TimedMessage<N, Req, Res>) -> Message<N, Req, Res> {
-        self.metrics_recv_qps.mark(1);
-        self.metrics_queued.dec(1);
-
-        let (since, msg) = data.into_message();
-        self.metrics_queue_latency.update_since(since);
-
-        msg
+        self.chan.try_recv()
     }
 }
 

common/channel/src/lib.rs

@@ -1,5 +1,6 @@
 mod channel;
 pub mod error;
+pub mod metrics;
 pub mod test_util;
 
 pub use crate::channel::{Channel, Message, Receiver, ResponseSender, Sender};

common/channel/src/metrics.rs

@@ -0,0 +1,112 @@
+use std::{fmt::Debug, sync::Arc, time::Instant};
+
+use metrics::{register_meter_with_group, Counter, CounterUsize, Histogram, Meter, Sample};
+use tokio::sync::mpsc::{
+    error::{SendError, TryRecvError},
+    unbounded_channel as new_unbounded_channel, UnboundedReceiver, UnboundedSender,
+};
+
+pub fn unbounded_channel<T>(metric_name: &str) -> (Sender<T>, Receiver<T>) {
+    let (sender, receiver) = new_unbounded_channel();
+    let metrics_queued = CounterUsize::register_with_group(metric_name, "size");
+    (
+        Sender::new(sender, metric_name, metrics_queued.clone()),
+        Receiver::new(receiver, metric_name, metrics_queued),
+    )
+}
+
+pub struct Sender<T> {
+    sender: UnboundedSender<(Instant, T)>,
+    metrics_send_qps: Arc<dyn Meter>,
+    metrics_queued: Arc<dyn Counter<usize>>,
+}
+
+impl<T> Clone for Sender<T> {
+    fn clone(&self) -> Self {
+        Self {
+            sender: self.sender.clone(),
+            metrics_send_qps: self.metrics_send_qps.clone(),
+            metrics_queued: self.metrics_queued.clone(),
+        }
+    }
+}
+
+impl<T> Debug for Sender<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?}", self.sender)
+    }
+}
+
+impl<T> Sender<T> {
+    pub(crate) fn new(
+        sender: UnboundedSender<(Instant, T)>,
+        metrics_group: &str,
+        metrics_queued: Arc<dyn Counter<usize>>,
+    ) -> Self {
+        Self {
+            sender,
+            metrics_send_qps: register_meter_with_group(metrics_group, "send"),
+            metrics_queued,
+        }
+    }
+
+    pub fn send(&self, value: T) -> Result<(), SendError<T>> {
+        match self.sender.send((Instant::now(), value)) {
+            Ok(()) => {
+                self.metrics_send_qps.mark(1);
+                self.metrics_queued.inc(1);
+                Ok(())
+            }
+            Err(e) => Err(SendError(e.0 .1)),
+        }
+    }
+}
+
+pub struct Receiver<T> {
+    receiver: UnboundedReceiver<(Instant, T)>,
+    metrics_recv_qps: Arc<dyn Meter>,
+    metrics_queued: Arc<dyn Counter<usize>>,
+    metrics_queue_latency: Arc<dyn Histogram>,
+}
+
+impl<T> Debug for Receiver<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?}", self.receiver)
+    }
+}
+
+impl<T> Receiver<T> {
+    pub(crate) fn new(
+        receiver: UnboundedReceiver<(Instant, T)>,
+        metrics_group: &str,
+        metrics_queued: Arc<dyn Counter<usize>>,
+    ) -> Self {
+        Self {
+            receiver,
+            metrics_recv_qps: register_meter_with_group(metrics_group, "recv"),
+            metrics_queued,
+            metrics_queue_latency: Sample::ExpDecay(0.015).register_with_group(
+                metrics_group,
+                "latency",
+                1024,
+            ),
+        }
+    }
+
+    fn on_recv(&self, value: (Instant, T)) -> T {
+        self.metrics_recv_qps.mark(1);
+        self.metrics_queued.dec(1);
+        self.metrics_queue_latency.update_since(value.0);
+        value.1
+    }
+
+    pub async fn recv(&mut self) -> Option<T> {
+        let value = self.receiver.recv().await?;
+        Some(self.on_recv(value))
+    }
+
+    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
+        let value = self.receiver.try_recv()?;
+        Ok(self.on_recv(value))
+    }
+}