diff --git a/Makefile b/Makefile index e2c9831f..b836120b 100644 --- a/Makefile +++ b/Makefile @@ -49,7 +49,7 @@ example: full: check-all test-all example udeps -fast: check test example +fast: check test example1 msrv: shellcheck ./scripts/* diff --git a/foyer-common/src/metrics/registry/prometheus.rs b/foyer-common/src/metrics/registry/prometheus.rs index 18a9bc35..fbb3fbe6 100644 --- a/foyer-common/src/metrics/registry/prometheus.rs +++ b/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>>> = + 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, +} + +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(&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,10 +228,14 @@ 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, + }, + ) } } @@ -118,23 +243,45 @@ impl RegistryOps for PrometheusMetricsRegistry { 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); + } }