Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Nov 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ example:

full: check-all test-all example udeps

fast: check test example
fast: check test example1

msrv:
shellcheck ./scripts/*
Expand Down
195 changes: 171 additions & 24 deletions foyer-common/src/metrics/registry/prometheus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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!(),

Check warning on line 53 in foyer-common/src/metrics/registry/prometheus.rs

View check run for this annotation

Codecov / codecov/patch

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

Added line #L53 was not covered by tests
}
}

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!(),

Check warning on line 75 in foyer-common/src/metrics/registry/prometheus.rs

View check run for this annotation

Codecov / codecov/patch

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

Added line #L75 was not covered by tests
}
}

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!(),

Check warning on line 97 in foyer-common/src/metrics/registry/prometheus.rs

View check run for this annotation

Codecov / codecov/patch

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

Added line #L97 was not covered by tests
}
}

#[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) {
Expand Down Expand Up @@ -64,15 +157,35 @@
}

/// 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
}
}

Expand All @@ -83,10 +196,14 @@
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(
Expand All @@ -95,10 +212,14 @@
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(
Expand All @@ -107,34 +228,60 @@
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);
}
}
Loading