feat: support multiple foyer instance share the same prometheus (#801)

Signed-off-by: MrCroxx <[email protected]>

Makefile

@@ -49,7 +49,7 @@ example:
 
 full: check-all test-all example udeps
 
-fast: check test example
+fast: check test example1
 
 msrv:
 	shellcheck ./scripts/*

foyer-common/src/metrics/registry/prometheus.rs

@@ -12,12 +12,105 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
+use std::{
+    collections::HashMap,
+    hash::{Hash, Hasher},
+    sync::{Arc, LazyLock},
+};
+
+use parking_lot::Mutex;
 use prometheus::{
     register_histogram_vec_with_registry, register_int_counter_vec_with_registry, register_int_gauge_vec_with_registry,
     Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Registry,
 };
 
-use crate::metrics::{CounterOps, CounterVecOps, GaugeOps, GaugeVecOps, HistogramOps, HistogramVecOps, RegistryOps};
+use crate::{
+    metrics::{CounterOps, CounterVecOps, GaugeOps, GaugeVecOps, HistogramOps, HistogramVecOps, RegistryOps},
+    scope::Scope,
+};
+
+static METRICS: LazyLock<Mutex<HashMap<PrometheusMetricsRegistry, HashMap<Metadata, MetricVec>>>> =
+    LazyLock::new(|| Mutex::new(HashMap::new()));
+
+fn get_or_register_counter_vec(registry: &PrometheusMetricsRegistry, metadata: Metadata) -> IntCounterVec {
+    let vec = METRICS.lock().with(|mut metrics| {
+        metrics
+            .get_mut(registry)
+            .expect("registry must be registered when creating")
+            .entry(metadata.clone())
+            .or_insert_with(|| {
+                MetricVec::Counter(
+                    register_int_counter_vec_with_registry! {
+                        metadata.name, metadata.desc, metadata.label_names, registry.registry
+                    }
+                    .unwrap(),
+                )
+            })
+            .clone()
+    });
+    match vec {
+        MetricVec::Counter(v) => v,
+        _ => unreachable!(),
+    }
+}
+
+fn get_or_register_gauge_vec(registry: &PrometheusMetricsRegistry, metadata: Metadata) -> IntGaugeVec {
+    let vec = METRICS.lock().with(|mut metrics| {
+        metrics
+            .get_mut(registry)
+            .expect("registry must be registered when creating")
+            .entry(metadata.clone())
+            .or_insert_with(|| {
+                MetricVec::Gauge(
+                    register_int_gauge_vec_with_registry! {
+                        metadata.name, metadata.desc, metadata.label_names, registry.registry
+                    }
+                    .unwrap(),
+                )
+            })
+            .clone()
+    });
+    match vec {
+        MetricVec::Gauge(v) => v,
+        _ => unreachable!(),
+    }
+}
+
+fn get_or_register_histogram_vec(registry: &PrometheusMetricsRegistry, metadata: Metadata) -> HistogramVec {
+    let vec = METRICS.lock().with(|mut metrics| {
+        metrics
+            .get_mut(registry)
+            .expect("registry must be registered when creating")
+            .entry(metadata.clone())
+            .or_insert_with(|| {
+                MetricVec::Histogram(
+                    register_histogram_vec_with_registry! {
+                        metadata.name, metadata.desc, metadata.label_names, registry.registry
+                    }
+                    .unwrap(),
+                )
+            })
+            .clone()
+    });
+    match vec {
+        MetricVec::Histogram(v) => v,
+        _ => unreachable!(),
+    }
+}
+
+#[derive(Debug, Clone)]
+enum MetricVec {
+    Counter(IntCounterVec),
+    Gauge(IntGaugeVec),
+    Histogram(HistogramVec),
+}
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+struct Metadata {
+    name: &'static str,
+    desc: &'static str,
+    label_names: &'static [&'static str],
+}
 
 impl CounterOps for IntCounter {
     fn increase(&self, val: u64) {
@@ -64,15 +157,35 @@ impl HistogramVecOps for HistogramVec {
 }
 
 /// Prometheus metrics registry.
-#[derive(Debug)]
+///
+/// The [`PrometheusMetricsRegistry`] can be cloned and used by multiple foyer instances, without worrying about
+/// duplicately registering.
+#[derive(Debug, Clone)]
 pub struct PrometheusMetricsRegistry {
-    registry: Registry,
+    registry: Arc<Registry>,
+}
+
+impl PartialEq for PrometheusMetricsRegistry {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.registry, &other.registry)
+    }
+}
+
+impl Eq for PrometheusMetricsRegistry {}
+
+impl Hash for PrometheusMetricsRegistry {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        Arc::as_ptr(&self.registry).hash(state);
+    }
 }
 
 impl PrometheusMetricsRegistry {
     /// Create an Prometheus metrics registry.
     pub fn new(registry: Registry) -> Self {
-        Self { registry }
+        let registry = Arc::new(registry);
+        let this = Self { registry };
+        METRICS.lock().insert(this.clone(), HashMap::new());
+        this
     }
 }
 
@@ -83,10 +196,14 @@ impl RegistryOps for PrometheusMetricsRegistry {
         desc: &'static str,
         label_names: &'static [&'static str],
     ) -> impl CounterVecOps {
-        register_int_counter_vec_with_registry! {
-            name, desc, label_names, self.registry
-        }
-        .unwrap()
+        get_or_register_counter_vec(
+            self,
+            Metadata {
+                name,
+                desc,
+                label_names,
+            },
+        )
     }
 
     fn register_gauge_vec(
@@ -95,10 +212,14 @@ impl RegistryOps for PrometheusMetricsRegistry {
         desc: &'static str,
         label_names: &'static [&'static str],
     ) -> impl GaugeVecOps {
-        register_int_gauge_vec_with_registry! {
-            name, desc, label_names, self.registry
-        }
-        .unwrap()
+        get_or_register_gauge_vec(
+            self,
+            Metadata {
+                name,
+                desc,
+                label_names,
+            },
+        )
     }
 
     fn register_histogram_vec(
@@ -107,34 +228,60 @@ impl RegistryOps for PrometheusMetricsRegistry {
         desc: &'static str,
         label_names: &'static [&'static str],
     ) -> impl HistogramVecOps {
-        register_histogram_vec_with_registry! {
-            name, desc, label_names, self.registry
-        }
-        .unwrap()
+        get_or_register_histogram_vec(
+            self,
+            Metadata {
+                name,
+                desc,
+                label_names,
+            },
+        )
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    #[test]
-    fn test() {
-        let registry = Registry::new();
-        let p8s = PrometheusMetricsRegistry::new(registry);
-
-        let cv = p8s.register_counter_vec("test_counter_1", "test counter 1", &["label1", "label2"]);
+    fn case(registry: &PrometheusMetricsRegistry) {
+        let cv = registry.register_counter_vec("test_counter_1", "test counter 1", &["label1", "label2"]);
         let c = cv.counter(&["l1", "l2"]);
         c.increase(42);
 
-        let gv = p8s.register_gauge_vec("test_gauge_1", "test gauge 1", &["label1", "label2"]);
+        let gv = registry.register_gauge_vec("test_gauge_1", "test gauge 1", &["label1", "label2"]);
         let g = gv.gauge(&["l1", "l2"]);
         g.increase(514);
         g.decrease(114);
         g.absolute(114514);
 
-        let hv = p8s.register_histogram_vec("test_histogram_1", "test histogram 1", &["label1", "label2"]);
+        let hv = registry.register_histogram_vec("test_histogram_1", "test histogram 1", &["label1", "label2"]);
         let h = hv.histogram(&["l1", "l2"]);
         h.record(114.514);
     }
+
+    #[test]
+    fn test_prometheus_metrics_registry() {
+        let registry = Registry::new();
+        let p8s = PrometheusMetricsRegistry::new(registry);
+        case(&p8s);
+    }
+
+    #[should_panic]
+    #[test]
+    fn test_duplicated_prometheus_metrics_registry_wrongly() {
+        let registry = Registry::new();
+        let p8s1 = PrometheusMetricsRegistry::new(registry.clone());
+        let p8s2 = PrometheusMetricsRegistry::new(registry);
+        case(&p8s1);
+        case(&p8s2);
+    }
+
+    #[test]
+    fn test_duplicated_prometheus_metrics_registry() {
+        let registry = Registry::new();
+        let p8s1 = PrometheusMetricsRegistry::new(registry);
+        let p8s2 = p8s1.clone();
+        case(&p8s1);
+        case(&p8s2);
+    }
 }