ref: add TestMetrics to abstract recorder reset and snapshots

CevizLee · Feb 25, 2022 · 177a924 · 177a924
1 parent 775d326
commit 177a924
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 66 deletions.
diff --git a/.integration/tests/metrics/mod.rs b/.integration/tests/metrics/mod.rs
@@ -23,35 +23,14 @@ use pea2pea::Pea2Pea;
 
 #[tokio::test]
 async fn metrics_initialization() {
-    // Start a test node.
-    let _test_node = TestNode::default().await;
-
     // Initialise the metrics, we need to call this manually in tests.
-    let snapshotter = metrics::initialize();
-
-    // Start a snarkOS node.
-    let _client_node = ClientNode::default().await;
+    let metrics = metrics::TestMetrics::new();
 
     // Verify the metrics have been properly initialised, expect the block height to be set.
-    assert_eq!(
-        metrics::get_metric(&snapshotter, metrics::blocks::HEIGHT),
-        metrics::MetricVal::Gauge(0.0)
-    );
-    assert_eq!(
-        metrics::get_metric(&snapshotter, metrics::peers::CONNECTED),
-        metrics::MetricVal::Gauge(0.0)
-    );
-    assert_eq!(
-        metrics::get_metric(&snapshotter, metrics::peers::CANDIDATE),
-        metrics::MetricVal::Gauge(0.0)
-    );
-    assert_eq!(
-        metrics::get_metric(&snapshotter, metrics::peers::RESTRICTED),
-        metrics::MetricVal::Gauge(0.0)
-    );
-
-    // Clear the recorder to avoid the global state bleeding into other tests.
-    metrics::clear_recorder();
+    assert_eq!(metrics.get_val_for(metrics::blocks::HEIGHT), metrics::MetricVal::Gauge(0.0));
+    assert_eq!(metrics.get_val_for(metrics::peers::CONNECTED), metrics::MetricVal::Gauge(0.0));
+    assert_eq!(metrics.get_val_for(metrics::peers::CANDIDATE), metrics::MetricVal::Gauge(0.0));
+    assert_eq!(metrics.get_val_for(metrics::peers::RESTRICTED), metrics::MetricVal::Gauge(0.0));
 }
 
 #[tokio::test]
@@ -60,7 +39,7 @@ async fn connect_disconnect() {
     let test_node = TestNode::default().await;
 
     // Initialise the metrics, we need to call this manually in tests.
-    let snapshotter = metrics::initialize();
+    let metrics = metrics::TestMetrics::new();
 
     // Start a snarkOS node.
     let client_node = ClientNode::default().await;
@@ -73,39 +52,18 @@ async fn connect_disconnect() {
         wait_until!(1, client_node.connected_peers().await.len() == 1);
 
         // Check the metrics.
-        assert_eq!(
-            metrics::get_metric(&snapshotter, metrics::peers::CONNECTED),
-            metrics::MetricVal::Gauge(1.0)
-        );
-        assert_eq!(
-            metrics::get_metric(&snapshotter, metrics::peers::CANDIDATE),
-            metrics::MetricVal::Gauge(0.0)
-        );
-        assert_eq!(
-            metrics::get_metric(&snapshotter, metrics::peers::RESTRICTED),
-            metrics::MetricVal::Gauge(0.0)
-        );
+        assert_eq!(metrics.get_val_for(metrics::peers::CONNECTED), metrics::MetricVal::Gauge(1.0));
+        assert_eq!(metrics.get_val_for(metrics::peers::CANDIDATE), metrics::MetricVal::Gauge(0.0));
+        assert_eq!(metrics.get_val_for(metrics::peers::RESTRICTED), metrics::MetricVal::Gauge(0.0));
 
         // Shut down the node, force a disconnect.
         test_node.node().disconnect(client_node.local_addr()).await;
 
         wait_until!(1, client_node.connected_peers().await.len() == 0);
 
         // Check the metrics.
-        assert_eq!(
-            metrics::get_metric(&snapshotter, metrics::peers::CONNECTED),
-            metrics::MetricVal::Gauge(0.0)
-        );
-        assert_eq!(
-            metrics::get_metric(&snapshotter, metrics::peers::CANDIDATE),
-            metrics::MetricVal::Gauge(1.0)
-        );
-        assert_eq!(
-            metrics::get_metric(&snapshotter, metrics::peers::RESTRICTED),
-            metrics::MetricVal::Gauge(0.0)
-        );
+        assert_eq!(metrics.get_val_for(metrics::peers::CONNECTED), metrics::MetricVal::Gauge(0.0));
+        assert_eq!(metrics.get_val_for(metrics::peers::CANDIDATE), metrics::MetricVal::Gauge(1.0));
+        assert_eq!(metrics.get_val_for(metrics::peers::RESTRICTED), metrics::MetricVal::Gauge(0.0));
     }
-
-    // Clear the recorder to avoid the global state bleeding into other tests.
-    metrics::clear_recorder();
 }
diff --git a/metrics/src/utils.rs b/metrics/src/utils.rs
@@ -17,19 +17,34 @@
 use metrics::Key;
 use metrics_util::{CompositeKey, DebugValue, MetricKind, Snapshotter};
 
+pub struct TestMetrics(Snapshotter);
+
+impl TestMetrics {
+    pub fn new() -> Self {
+        Self(crate::initialize())
+    }
+
+    pub fn get_val_for(&self, metric: &'static str) -> MetricVal {
+        let key = CompositeKey::new(MetricKind::Gauge, Key::from_name(metric));
+
+        match &self.0.snapshot().into_hashmap().get(&key).unwrap().2 {
+            DebugValue::Gauge(val) => MetricVal::Gauge(val.into_inner()),
+            DebugValue::Counter(val) => MetricVal::Counter(*val),
+            DebugValue::Histogram(vals) => MetricVal::Histogram(vals.iter().map(|val| val.into_inner()).collect()),
+        }
+    }
+}
+
+impl Drop for TestMetrics {
+    fn drop(&mut self) {
+        // Clear the recorder to avoid the global state bleeding into other tests.
+        metrics::clear_recorder();
+    }
+}
+
 #[derive(Debug, PartialEq)]
 pub enum MetricVal {
     Counter(u64),
     Gauge(f64),
     Histogram(Vec<f64>),
 }
-
-pub fn get_metric(snapshotter: &Snapshotter, metric: &'static str) -> MetricVal {
-    let key = CompositeKey::new(MetricKind::Gauge, Key::from_name(metric));
-
-    match &snapshotter.snapshot().into_hashmap().get(&key).unwrap().2 {
-        DebugValue::Gauge(val) => MetricVal::Gauge(val.into_inner()),
-        DebugValue::Counter(val) => MetricVal::Counter(*val),
-        DebugValue::Histogram(vals) => MetricVal::Histogram(vals.iter().map(|val| val.into_inner()).collect()),
-    }
-}
diff --git a/network/src/ledger.rs b/network/src/ledger.rs
@@ -513,7 +513,6 @@ impl<N: Network, E: Environment> Ledger<N, E> {
                 // Attempt to add the unconfirmed block as the next block in the canonical chain.
                 false => match self.canon.add_next_block(&unconfirmed_block) {
                     Ok(()) => {
-                        // Allocation necessary if prometheus is enabled.
                         let latest_block_height = self.canon.latest_block_height();
                         info!(
                             "Ledger successfully advanced to block {} ({})",
@@ -570,7 +569,6 @@ impl<N: Network, E: Environment> Ledger<N, E> {
 
         match self.canon.revert_to_block_height(block_height) {
             Ok(removed_blocks) => {
-                // Allocation necessary if prometheus is enabled.
                 let latest_block_height = self.canon.latest_block_height();
                 info!("Ledger successfully reverted to block {}", latest_block_height);