Merge 6fb21d1424 into c1f465e009

common/channel/src/channel.rs

@@ -1,10 +1,9 @@
 use crate::error::Error;
-use crate::metrics::unbounded_channel;
+use metrics::{register_meter_with_group, Counter, CounterUsize, Histogram, Meter, Sample};
-use metrics::{Counter, CounterUsize};
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
-use tokio::sync::mpsc::error::TryRecvError;
+use tokio::sync::mpsc::error::{SendError, TryRecvError};
-use tokio::sync::oneshot;
+use tokio::sync::{mpsc, oneshot};
 use tokio::time::timeout;
 
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(3);
@@ -23,21 +22,63 @@ 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 (sender, receiver) = unbounded_channel(metrics_group.as_str());
+        let metrics_queued = CounterUsize::register_with_group(metrics_group.as_str(), "size");
-        let metrics_timeout = CounterUsize::register_with_group(metrics_group.as_str(), "timeout");
+
         (
             Sender {
                 chan: sender,
-                metrics_timeout,
+                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,
+                ),
             },
-            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: crate::metrics::Sender<Message<N, Req, Res>>,
+    chan: mpsc::UnboundedSender<TimedMessage<N, Req, Res>>,
+
+    metrics_send_qps: Arc<dyn Meter>,
+    metrics_queued: Arc<dyn Counter<usize>>,
     metrics_timeout: Arc<dyn Counter<usize>>,
 }
 
@@ -45,6 +86,8 @@ 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(),
         }
     }
@@ -52,17 +95,13 @@ 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.chan
+        self.send(Message::Notification(msg))
-            .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.chan
+        self.send(Message::Request(request, sender))?;
-            .send(Message::Request(request, sender))
-            .map_err(|e| Error::SendError(e))?;
 
         timeout(DEFAULT_REQUEST_TIMEOUT, receiver)
             .await
@@ -72,19 +111,49 @@ 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: crate::metrics::Receiver<Message<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>,
 }
 
 impl<N, Req, Res> Receiver<N, Req, Res> {
     pub async fn recv(&mut self) -> Option<Message<N, Req, Res>> {
-        self.chan.recv().await
+        let data = self.chan.recv().await?;
+        Some(self.on_recv_data(data))
     }
 
     pub fn try_recv(&mut self) -> Result<Message<N, Req, Res>, TryRecvError> {
-        self.chan.try_recv()
+        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
     }
 }
 

common/channel/src/lib.rs

@@ -1,6 +1,5 @@
 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

@@ -1,112 +0,0 @@
-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))
-    }
-}