fix: make sure TestMetrics doesn't remove state

src/tests/performance/connections.rs

@@ -200,10 +200,12 @@ async fn load_bearing() {
         let mut stats = Stats::new(MAX_PEERS, synth_count);
         stats.time = test_start.elapsed().as_secs_f64();
         {
-            stats.accepted = test_metrics.get_counter(METRIC_ACCEPTED) as u16;
-            stats.terminated = test_metrics.get_counter(METRIC_TERMINATED) as u16;
-            stats.rejected = test_metrics.get_counter(METRIC_REJECTED) as u16;
-            stats.conn_error = test_metrics.get_counter(METRIC_ERROR) as u16;
+            let snapshot = test_metrics.take_snapshot();
+
+            stats.accepted = snapshot.get_counter(METRIC_ACCEPTED) as u16;
+            stats.terminated = snapshot.get_counter(METRIC_TERMINATED) as u16;
+            stats.rejected = snapshot.get_counter(METRIC_REJECTED) as u16;
+            stats.conn_error = snapshot.get_counter(METRIC_ERROR) as u16;
 
             stats.timed_out = synth_count - stats.accepted - stats.rejected - stats.conn_error;
         }

src/tests/performance/getdata_blocks.rs

@@ -156,7 +156,8 @@ async fn throughput() {
 
         let time_taken_secs = test_start.elapsed().as_secs_f64();
 
-        if let Some(latencies) = test_metrics.construct_histogram(METRIC_LATENCY) {
+        let snapshot = test_metrics.take_snapshot();
+        if let Some(latencies) = snapshot.construct_histogram(METRIC_LATENCY) {
             if latencies.entries() >= 1 {
                 // add stats to table display
                 table.add_row(RequestStats::new(

src/tests/performance/ping_pong.rs

@@ -184,7 +184,8 @@ async fn throughput() {
 
         let time_taken_secs = test_start.elapsed().as_secs_f64();
 
-        if let Some(latencies) = test_metrics.construct_histogram(METRIC_LATENCY) {
+        let snapshot = test_metrics.take_snapshot();
+        if let Some(latencies) = snapshot.construct_histogram(METRIC_LATENCY) {
             if latencies.entries() >= 1 {
                 // add stats to table display
                 table.add_row(RequestStats::new(

src/tests/resistance/stress_test.rs

@@ -298,7 +298,8 @@ async fn throughput() {
         }
         let iteration_time = iteration_timer.elapsed().as_secs_f64();
 
-        if let Some(request_latencies) = test_metrics.construct_histogram(REQUEST_LATENCY) {
+        let snapshot = test_metrics.take_snapshot();
+        if let Some(request_latencies) = snapshot.construct_histogram(REQUEST_LATENCY) {
             if request_latencies.entries() >= 1 {
                 let row = RequestStats::new(
                     peers as u16,
@@ -310,34 +311,34 @@ async fn throughput() {
             }
         }
 
-        if let Some(handshake_latencies) = test_metrics.construct_histogram(HANDSHAKE_LATENCY) {
+        if let Some(handshake_latencies) = snapshot.construct_histogram(HANDSHAKE_LATENCY) {
             if handshake_latencies.entries() >= 1 {
                 let row = RequestStats::new(peers as u16, 1, handshake_latencies, iteration_time);
                 handshake_table.add_row(row);
+            }
 
-                // update stats table
-                let mut stat = Stats {
-                    peers,
-                    requests: MAX_VALID_MESSAGES,
-                    time: iteration_time,
-                    dangling: peers as u16 - completed,
-                    ..Default::default()
-                };
-                {
-                    stat.handshake_accepted = test_metrics.get_counter(HANDSHAKE_ACCEPTED) as u16;
-                    stat.handshake_rejected = test_metrics.get_counter(HANDSHAKE_REJECTED) as u16;
-
-                    stat.corrupt_terminated = test_metrics.get_counter(CORRUPT_TERMINATED) as u16;
-                    stat.corrupt_ignored = test_metrics.get_counter(CORRUPT_IGNORED) as u16;
-                    stat.corrupt_rejected = test_metrics.get_counter(CORRUPT_REJECTED) as u16;
-                    stat.corrupt_reply = test_metrics.get_counter(CORRUPT_REPLY) as u16;
-
-                    stat.peers_dropped = test_metrics.get_counter(CONNECTION_TERMINATED) as u16;
-                    stat.reply_errors = test_metrics.get_counter(BAD_REPLY) as u16;
-                }
-
-                stats.push(stat);
+            // update stats table
+            let mut stat = Stats {
+                peers,
+                requests: MAX_VALID_MESSAGES,
+                time: iteration_time,
+                dangling: peers as u16 - completed,
+                ..Default::default()
+            };
+            {
+                stat.handshake_accepted = snapshot.get_counter(HANDSHAKE_ACCEPTED) as u16;
+                stat.handshake_rejected = snapshot.get_counter(HANDSHAKE_REJECTED) as u16;
+
+                stat.corrupt_terminated = snapshot.get_counter(CORRUPT_TERMINATED) as u16;
+                stat.corrupt_ignored = snapshot.get_counter(CORRUPT_IGNORED) as u16;
+                stat.corrupt_rejected = snapshot.get_counter(CORRUPT_REJECTED) as u16;
+                stat.corrupt_reply = snapshot.get_counter(CORRUPT_REPLY) as u16;
+
+                stat.peers_dropped = snapshot.get_counter(CONNECTION_TERMINATED) as u16;
+                stat.reply_errors = snapshot.get_counter(BAD_REPLY) as u16;
             }
+
+            stats.push(stat);
         }
     }
 

src/tools/metrics/recorder.rs

@@ -1,60 +1,42 @@
 //! Metrics recording types and utilities.
 
+use std::collections::HashMap;
+
 use histogram::Histogram;
 use metrics::Key;
 use metrics_util::{
     debugging::{DebugValue, DebuggingRecorder, Snapshotter},
     CompositeKey, MetricKind,
 };
 
-pub fn initialize() -> Snapshotter {
-    let recorder = DebuggingRecorder::new();
-    let snapshotter = recorder.snapshotter();
-    let _ = recorder.install();
-
-    snapshotter
-}
-
-pub struct TestMetrics(Snapshotter);
-
-impl Default for TestMetrics {
-    fn default() -> Self {
-        Self(initialize())
-    }
-}
-
-impl TestMetrics {
-    pub fn get_val_for(&self, kind: MetricKind, metric: &'static str) -> MetricVal {
-        let key = CompositeKey::new(kind, Key::from_name(metric));
-
-        match &self.0.snapshot().into_hashmap().get(&key).unwrap().2 {
-            DebugValue::Counter(val) => MetricVal::Counter(*val),
-            DebugValue::Gauge(val) => MetricVal::Gauge(val.into_inner()),
-            DebugValue::Histogram(vals) => {
-                MetricVal::Histogram(vals.iter().map(|val| val.into_inner()).collect())
-            }
-        }
-    }
+// This is a false positive, `CompositeKey` does not depend on any interior mutable
+// types for its hashing implementation. Therefore, it is safe to use in our context.
+#[allow(clippy::mutable_key_type)]
+pub struct Snapshot(HashMap<CompositeKey, MetricVal>);
 
+impl Snapshot {
     pub fn get_counter(&self, metric: &'static str) -> u64 {
-        if let MetricVal::Counter(val) = self.get_val_for(MetricKind::Counter, metric) {
+        let key = CompositeKey::new(MetricKind::Counter, Key::from_name(metric));
+        if let MetricVal::Counter(val) = *self.0.get(&key).unwrap() {
             val
         } else {
             0
         }
     }
 
     pub fn get_gauge(&self, metric: &'static str) -> f64 {
-        if let MetricVal::Gauge(val) = self.get_val_for(MetricKind::Gauge, metric) {
+        let key = CompositeKey::new(MetricKind::Gauge, Key::from_name(metric));
+        if let MetricVal::Gauge(val) = *self.0.get(&key).unwrap() {
             val
         } else {
             0.0
         }
     }
 
     pub fn get_histogram(&self, metric: &'static str) -> Option<Vec<f64>> {
-        if let MetricVal::Histogram(vals) = self.get_val_for(MetricKind::Histogram, metric) {
-            Some(vals)
+        let key = CompositeKey::new(MetricKind::Histogram, Key::from_name(metric));
+        if let MetricVal::Histogram(vals) = self.0.get(&key).unwrap() {
+            Some(vals.to_vec())
         } else {
             None
         }
@@ -75,6 +57,44 @@ impl TestMetrics {
     }
 }
 
+pub struct TestMetrics(Snapshotter);
+
+impl TestMetrics {
+    fn new() -> TestMetrics {
+        let recorder = DebuggingRecorder::new();
+        let snapshotter = recorder.snapshotter();
+        let _ = recorder.install();
+
+        TestMetrics(snapshotter)
+    }
+}
+
+impl Default for TestMetrics {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TestMetrics {
+    #[allow(clippy::mutable_key_type)]
+    pub fn take_snapshot(&self) -> Snapshot {
+        let mut snapshot = HashMap::new();
+
+        for (key, val) in self.0.snapshot().into_hashmap().into_iter() {
+            match val.2 {
+                DebugValue::Counter(val) => snapshot.insert(key, MetricVal::Counter(val)),
+                DebugValue::Gauge(val) => snapshot.insert(key, MetricVal::Gauge(val.into_inner())),
+                DebugValue::Histogram(vals) => snapshot.insert(
+                    key,
+                    MetricVal::Histogram(vals.iter().map(|val| val.into_inner()).collect()),
+                ),
+            };
+        }
+
+        Snapshot(snapshot)
+    }
+}
+
 impl Drop for TestMetrics {
     fn drop(&mut self) {
         // Clear the recorder to avoid the global state bleeding into other tests.
@@ -100,6 +120,7 @@ mod tests {
     use super::*;
 
     const METRIC_NAME: &str = "test_metrics";
+    const METRIC_NAME_ALT: &str = "test_metrics_alt";
     const COUNTER_INC: u64 = 25;
     const HISTOGRAM_SIZE: usize = 50;
 
@@ -117,7 +138,9 @@ mod tests {
 
         counter.increment(COUNTER_INC);
 
-        assert_eq!(metrics.get_counter(METRIC_NAME), COUNTER_INC);
+        let snapshot = metrics.take_snapshot();
+
+        assert_eq!(snapshot.get_counter(METRIC_NAME), COUNTER_INC);
     }
 
     #[test]
@@ -130,7 +153,9 @@ mod tests {
         gauge.decrement(500.0);
         gauge.increment(25.0);
 
-        assert_eq!(metrics.get_gauge(METRIC_NAME), 525.0);
+        let snapshot = metrics.take_snapshot();
+
+        assert_eq!(snapshot.get_gauge(METRIC_NAME), 525.0);
     }
 
     #[test]
@@ -145,7 +170,9 @@ mod tests {
             values.push(i as f64);
         }
 
-        assert_eq!(metrics.get_histogram(METRIC_NAME), Some(values));
+        let snapshot = metrics.take_snapshot();
+
+        assert_eq!(snapshot.get_histogram(METRIC_NAME), Some(values));
     }
 
     #[test]
@@ -159,10 +186,29 @@ mod tests {
         histogram.record(5.0);
         histogram.record(9.0);
 
-        let constructed_histogram = metrics.construct_histogram(METRIC_NAME).unwrap();
+        let snapshot = metrics.take_snapshot();
+        let constructed_histogram = snapshot.construct_histogram(METRIC_NAME).unwrap();
 
         assert!(constructed_histogram.entries() == 4);
         assert_eq!(constructed_histogram.percentile(50.0).unwrap(), 5);
         assert_eq!(constructed_histogram.percentile(90.0).unwrap(), 9);
     }
+
+    #[test]
+    #[ignore]
+    fn can_construct_multiple_histograms() {
+        let metrics = TestMetrics::default();
+        let histogram = register_histogram!(METRIC_NAME);
+        histogram.record(1.0);
+
+        let histogram2 = register_histogram!(METRIC_NAME_ALT);
+        histogram2.record(1.0);
+
+        let snapshot = metrics.take_snapshot();
+        let constructed_histogram = snapshot.construct_histogram(METRIC_NAME).unwrap();
+        let constructed_histogram2 = snapshot.construct_histogram(METRIC_NAME_ALT).unwrap();
+
+        assert!(constructed_histogram.entries() == 1);
+        assert!(constructed_histogram2.entries() == 1);
+    }
 }