Refactored how metrics are handled - separate sender, receiver and reporter

mixnode/src/node/metrics.rs

@@ -8,58 +8,47 @@ use log::{debug, error};
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio::runtime::Handle;
+use tokio::task::JoinHandle;
 
-#[derive(Debug)]
-pub struct MetricsReporter {
+pub(crate) enum MetricEvent {
+    Sent(String),
+    Received,
+}
+
+#[derive(Debug, Clone)]
+struct MixMetrics {
+    inner: Arc<Mutex<MixMetricsInner>>,
+}
+
+struct MixMetricsInner {
     received: u64,
     sent: HashMap<String, u64>,
 }
 
-impl MetricsReporter {
+impl MixMetrics {
     pub(crate) fn new() -> Self {
-        MetricsReporter {
-            received: 0,
-            sent: HashMap::new(),
+        MixMetrics {
+            inner: Arc::new(Mutex::new(MixMetricsInner {
+                received: 0,
+                sent: HashMap::new(),
+            })),
         }
     }
 
-    pub(crate) fn add_arc_mutex(self) -> Arc<Mutex<Self>> {
-        Arc::new(Mutex::new(self))
-    }
-
-    async fn increment_received_metrics(metrics: Arc<Mutex<MetricsReporter>>) {
-        let mut unlocked = metrics.lock().await;
+    async fn increment_received_metrics(&mut self) {
+        let mut unlocked = self.inner.lock().await;
         unlocked.received += 1;
     }
 
-    pub(crate) async fn run_received_metrics_control(
-        metrics: Arc<Mutex<MetricsReporter>>,
-        mut rx: mpsc::Receiver<()>,
-    ) {
-        while let Some(_) = rx.next().await {
-            MetricsReporter::increment_received_metrics(metrics.clone()).await;
-        }
-    }
-
-    async fn increment_sent_metrics(metrics: Arc<Mutex<MetricsReporter>>, sent_to: String) {
-        let mut unlocked = metrics.lock().await;
-        let receiver_count = unlocked.sent.entry(sent_to).or_insert(0);
+    async fn increment_sent_metrics(&mut self, destination: String) {
+        let mut unlocked = self.inner.lock().await;
+        let receiver_count = unlocked.sent.entry(destination).or_insert(0);
         *receiver_count += 1;
     }
 
-    pub(crate) async fn run_sent_metrics_control(
-        metrics: Arc<Mutex<MetricsReporter>>,
-        mut rx: mpsc::Receiver<String>,
-    ) {
-        while let Some(sent_metric) = rx.next().await {
-            MetricsReporter::increment_sent_metrics(metrics.clone(), sent_metric).await;
-        }
-    }
-
-    async fn acquire_and_reset_metrics(
-        metrics: Arc<Mutex<MetricsReporter>>,
-    ) -> (u64, HashMap<String, u64>) {
-        let mut unlocked = metrics.lock().await;
+    async fn acquire_and_reset_metrics(&mut self) -> (u64, HashMap<String, u64>) {
+        let mut unlocked = self.inner.lock().await;
         let received = unlocked.received;
 
         let sent = std::mem::replace(&mut unlocked.sent, HashMap::new());
@@ -67,27 +56,138 @@ impl MetricsReporter {
 
         (received, sent)
     }
+}
 
-    pub(crate) async fn run_metrics_sender(
-        metrics: Arc<Mutex<MetricsReporter>>,
-        cfg: directory_client::Config,
-        pub_key_str: String,
-        sending_delay: Duration,
-    ) {
-        let directory_client = directory_client::Client::new(cfg);
-        loop {
-            tokio::time::delay_for(sending_delay).await;
-            let (received, sent) =
-                MetricsReporter::acquire_and_reset_metrics(metrics.clone()).await;
+struct MetricsReceiver {
+    metrics: MixMetrics,
+    metrics_rx: mpsc::UnboundedReceiver<MetricEvent>,
+}
 
-            match directory_client.metrics_post.post(&MixMetric {
-                pub_key: pub_key_str.clone(),
-                received,
-                sent,
-            }) {
-                Err(err) => error!("failed to send metrics - {:?}", err),
-                Ok(_) => debug!("sent metrics information"),
-            }
+impl MetricsReceiver {
+    fn new(metrics: MixMetrics, metrics_rx: mpsc::UnboundedReceiver<MetricEvent>) -> Self {
+        MetricsReceiver {
+            metrics,
+            metrics_rx,
         }
     }
+
+    fn start(mut self, handle: &Handle) -> JoinHandle<()> {
+        handle.spawn(async move {
+            while let Some(metrics_data) = self.metrics_rx.next().await {
+                match metrics_data {
+                    MetricEvent::Received => self.metrics.increment_received_metrics().await,
+                    MetricEvent::Sent(destination) => {
+                        self.metrics.increment_sent_metrics(destination).await
+                    }
+                }
+            }
+        })
+    }
+}
+
+struct MetricsSender {
+    metrics: MixMetrics,
+    directory_client: directory_client::Client,
+    pub_key_str: String,
+    sending_delay: Duration,
+}
+
+impl MetricsSender {
+    fn new(
+        metrics: MixMetrics,
+        directory_server: String,
+        pub_key_str: String,
+        sending_delay: Duration,
+    ) -> Self {
+        MetricsSender {
+            metrics,
+            directory_client: directory_client::Client::new(directory_client::Config::new(
+                directory_server,
+            )),
+            pub_key_str,
+            sending_delay,
+        }
+    }
+
+    fn start(mut self, handle: &Handle) -> JoinHandle<()> {
+        handle.spawn(async move {
+            loop {
+                tokio::time::delay_for(self.sending_delay).await;
+                let (received, sent) = self.metrics.acquire_and_reset_metrics().await;
+
+                match self.directory_client.metrics_post.post(&MixMetric {
+                    pub_key: self.pub_key_str.clone(),
+                    received,
+                    sent,
+                }) {
+                    Err(err) => error!("failed to send metrics - {:?}", err),
+                    Ok(_) => debug!("sent metrics information"),
+                }
+            }
+        })
+    }
+}
+
+#[derive(Clone)]
+pub struct MetricsReporter {
+    metrics_tx: mpsc::UnboundedSender<MetricEvent>,
+}
+
+impl MetricsReporter {
+    pub(crate) fn new(metrics_tx: mpsc::UnboundedSender<MetricEvent>) -> Self {
+        MetricsReporter { metrics_tx }
+    }
+
+    pub(crate) fn report_sent(&self, destination: String) {
+        // in unbounded_send() failed it means that the receiver channel was disconnected
+        // and hence something weird must have happened without a way of recovering
+        self.metrics_tx
+            .unbounded_send(MetricEvent::Sent(destination))
+            .unwrap()
+    }
+
+    pub(crate) fn report_received(&self) {
+        // in unbounded_send() failed it means that the receiver channel was disconnected
+        // and hence something weird must have happened without a way of recovering
+        self.metrics_tx
+            .unbounded_send(MetricEvent::Received)
+            .unwrap()
+    }
+}
+
+// basically an easy single entry point to start all metrics related tasks
+pub struct MetricsController {
+    receiver: MetricsReceiver,
+    reporter: MetricsReporter,
+    sender: MetricsSender,
+}
+
+impl MetricsController {
+    pub(crate) fn new(
+        directory_server: String,
+        pub_key_str: String,
+        sending_delay: Duration,
+    ) -> Self {
+        let (metrics_tx, metrics_rx) = mpsc::unbounded();
+        let shared_metrics = MixMetrics::new();
+
+        MetricsController {
+            sender: MetricsSender::new(
+                shared_metrics.clone(),
+                directory_server,
+                pub_key_str,
+                sending_delay,
+            ),
+            receiver: MetricsReceiver::new(shared_metrics, metrics_rx),
+            reporter: MetricsReporter::new(metrics_tx),
+        }
+    }
+
+    // reporter is how node is going to be accessing the metrics data
+    pub(crate) fn start(self, handle: &Handle) -> MetricsReporter {
+        // TODO: should we do anything with JoinHandle(s) returned by start methods?
+        self.receiver.start(handle);
+        self.sender.start(handle);
+        self.reporter
+    }
 }