diff --git a/examples/tide.rs b/examples/tide.rs index dfb68125..3a09015a 100644 --- a/examples/tide.rs +++ b/examples/tide.rs @@ -31,7 +31,7 @@ async fn main() -> std::result::Result<(), std::io::Error> { app.at("/metrics") .get(|req: tide::Request| async move { let mut encoded = Vec::new(); - encode(&mut encoded, &req.state().registry).unwrap(); + encode(&mut encoded, req.state().registry.as_ref()).unwrap(); let response = tide::Response::builder(200) .body(encoded) .content_type("application/openmetrics-text; version=1.0.0; charset=utf-8") diff --git a/src/encoding/text.rs b/src/encoding/text.rs index 57914c5e..ff73739a 100644 --- a/src/encoding/text.rs +++ b/src/encoding/text.rs @@ -31,7 +31,7 @@ use crate::metrics::gauge::{self, Gauge}; use crate::metrics::histogram::Histogram; use crate::metrics::info::Info; use crate::metrics::{MetricType, TypedMetric}; -use crate::registry::{Registry, Unit}; +use crate::registry::{Collector, Unit}; use std::borrow::Cow; use std::collections::HashMap; @@ -40,12 +40,15 @@ use std::ops::Deref; pub use prometheus_client_derive_text_encode::*; -pub fn encode(writer: &mut W, registry: &Registry) -> Result<(), std::io::Error> +pub fn encode<'a, W, M>( + writer: &mut W, + registry: &'a dyn Collector<'a, M>, +) -> Result<(), std::io::Error> where W: Write, - M: EncodeMetric, + M: EncodeMetric + 'a, { - for (desc, metric) in registry.iter() { + for (desc, metric) in registry.collect() { writer.write_all(b"# HELP ")?; writer.write_all(desc.name().as_bytes())?; if let Some(unit) = desc.unit() { @@ -613,9 +616,40 @@ mod tests { use crate::metrics::counter::Counter; use crate::metrics::gauge::Gauge; use crate::metrics::histogram::exponential_buckets; + use crate::registry::Registry; use pyo3::{prelude::*, types::PyModule}; use std::borrow::Cow; + #[test] + fn collect_multiple_registries() { + let mut registry_single = Registry::default(); + let mut registry_1 = Registry::default(); + let mut registry_2 = Registry::default(); + + let counter_1_1: Counter = Counter::default(); + let counter_2_2: Counter = Counter::default(); + let counter_s_1: Counter = Counter::default(); + let counter_s_2: Counter = Counter::default(); + + counter_1_1.inc(); + counter_2_2.inc(); + counter_s_1.inc(); + counter_s_2.inc(); + + registry_single.register("counter_1", "Counter 1", counter_s_1); + registry_single.register("counter_2", "Counter 2", counter_s_2); + registry_1.register("counter_1", "Counter 1", counter_1_1); + registry_2.register("counter_2", "Counter 2", counter_2_2); + + let mut encoded_single = Vec::new(); + let mut encoded_combined = Vec::new(); + + encode(&mut encoded_single, ®istry_single).unwrap(); + encode(&mut encoded_combined, &vec![®istry_1, ®istry_2]).unwrap(); + + assert_eq!(encoded_single, encoded_combined); + } + #[test] fn encode_counter() { let counter: Counter = Counter::default(); diff --git a/src/registry.rs b/src/registry.rs index dcb01926..4704ab0f 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -2,6 +2,7 @@ //! //! See [`Registry`] for details. +use crate::encoding::text::{EncodeMetric, SendSyncEncodeMetric}; use std::borrow::Cow; /// A metric registry. @@ -58,7 +59,7 @@ use std::borrow::Cow; /// # assert_eq!(expected, String::from_utf8(buffer).unwrap()); /// ``` #[derive(Debug)] -pub struct Registry> { +pub struct Registry> { prefix: Option, labels: Vec<(Cow<'static, str>, Cow<'static, str>)>, metrics: Vec<(Descriptor, M)>, @@ -249,6 +250,35 @@ impl Registry { } } +/// Trait that allows manipulating Registries as a collection of Metrics. +/// +/// The main use is combining a bunch of registries for encoding. +pub trait Collector<'a, M> +where + M: EncodeMetric + 'a, +{ + fn collect(&'a self) -> Vec<&'a (Descriptor, M)>; +} + +impl<'a, M> Collector<'a, M> for Registry +where + M: EncodeMetric + 'a, +{ + fn collect(&'a self) -> Vec<&'a (Descriptor, M)> { + self.iter().collect() + } +} + +impl<'a, M, C> Collector<'a, M> for Vec<&'a C> +where + M: EncodeMetric + 'a, + C: Collector<'a, M>, +{ + fn collect(&'a self) -> Vec<&'a (Descriptor, M)> { + self.iter().flat_map(|r| r.collect()).collect() + } +} + /// Iterator iterating both the metrics registered directly with the registry as /// well as all metrics registered with sub-registries. #[derive(Debug)]