From 52479df9658d13c30b6165258ce6fb72ccf22aab Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Tue, 24 Oct 2023 17:16:42 +0100
Subject: [PATCH] less duplication, more sharing

---
 libs/metrics/src/tokio_metrics.rs | 92 +++++++++++++++----------------
 1 file changed, 46 insertions(+), 46 deletions(-)

diff --git a/libs/metrics/src/tokio_metrics.rs b/libs/metrics/src/tokio_metrics.rs
index 00de217778..0ef78a22c8 100644
--- a/libs/metrics/src/tokio_metrics.rs
+++ b/libs/metrics/src/tokio_metrics.rs
@@ -1,4 +1,7 @@
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
 
 use once_cell::sync::Lazy;
 use prometheus::{
@@ -6,14 +9,14 @@ use prometheus::{
     proto::{self, LabelPair},
 };
 
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct TokioCollector {
-    handles: Vec<tokio::runtime::Handle>,
+    handles: Arc<Mutex<Vec<tokio::runtime::Handle>>>,
 }
 
 impl TokioCollector {
     pub fn add_runtime(&mut self, handle: tokio::runtime::Handle) {
-        self.handles.push(handle);
+        self.handles.lock().unwrap().push(handle);
     }
 
     fn guage(
@@ -30,33 +33,6 @@ impl TokioCollector {
         })
     }
 
-    fn metric(
-        &self,
-        desc: &Desc,
-        typ: proto::MetricType,
-        f: impl Fn(&tokio::runtime::RuntimeMetrics) -> proto::Metric,
-    ) -> proto::MetricFamily {
-        let mut m = proto::MetricFamily::default();
-        m.set_name(desc.fq_name.clone());
-        m.set_help(desc.help.clone());
-        m.set_field_type(typ);
-        let mut metrics = vec![];
-        for rt in &self.handles {
-            let rt_metrics = rt.metrics();
-            let mut m = f(&rt_metrics);
-
-            let mut label_id = LabelPair::default();
-            label_id.set_name("id".to_owned());
-            label_id.set_value(rt.id().to_string());
-
-            m.set_label(vec![label_id]);
-            metrics.push(m);
-        }
-
-        m.set_metric(metrics);
-        m
-    }
-
     // fn guage_per_worker(
     //     &self,
     //     desc: &Desc,
@@ -85,35 +61,59 @@ impl TokioCollector {
         })
     }
 
+    fn metric(
+        &self,
+        desc: &Desc,
+        typ: proto::MetricType,
+        f: impl Fn(&tokio::runtime::RuntimeMetrics) -> proto::Metric,
+    ) -> proto::MetricFamily {
+        self.metrics(desc, typ, |rt, label_id, metrics| {
+            let mut m = f(rt);
+            m.set_label(vec![label_id.clone()]);
+            metrics.push(m);
+        })
+    }
+
     fn metric_per_worker(
         &self,
         desc: &Desc,
         typ: proto::MetricType,
         f: impl Fn(&tokio::runtime::RuntimeMetrics, usize) -> proto::Metric,
     ) -> proto::MetricFamily {
-        let mut m = proto::MetricFamily::default();
-        m.set_name(desc.fq_name.clone());
-        m.set_help(desc.help.clone());
-        m.set_field_type(typ);
-        let mut metrics = vec![];
-        for rt in &self.handles {
-            let rt_metrics = rt.metrics();
-            let workers = rt_metrics.num_workers();
-
+        self.metrics(desc, typ, |rt, label_id, metrics| {
+            let workers = rt.num_workers();
             for i in 0..workers {
-                let mut m = f(&rt_metrics, i);
-
-                let mut label_id = LabelPair::default();
-                label_id.set_name("id".to_owned());
-                label_id.set_value(rt.id().to_string());
+                let mut m = f(rt, i);
 
                 let mut label_worker = LabelPair::default();
                 label_worker.set_name("worker".to_owned());
                 label_worker.set_value(i.to_string());
 
-                m.set_label(vec![label_id, label_worker]);
+                m.set_label(vec![label_id.clone(), label_worker]);
                 metrics.push(m);
             }
+        })
+    }
+
+    fn metrics(
+        &self,
+        desc: &Desc,
+        typ: proto::MetricType,
+        f: impl Fn(&tokio::runtime::RuntimeMetrics, LabelPair, &mut Vec<proto::Metric>),
+    ) -> proto::MetricFamily {
+        let mut m = proto::MetricFamily::default();
+        m.set_name(desc.fq_name.clone());
+        m.set_help(desc.help.clone());
+        m.set_field_type(typ);
+        let mut metrics = vec![];
+        for rt in &*self.handles.lock().unwrap() {
+            let rt_metrics = rt.metrics();
+
+            let mut label_id = LabelPair::default();
+            label_id.set_name("id".to_owned());
+            label_id.set_value(rt.id().to_string());
+
+            f(&rt_metrics, label_id, &mut metrics);
         }
 
         m.set_metric(metrics);