Collector pattern

examples/tide.rs

@@ -31,7 +31,7 @@ async fn main() -> std::result::Result<(), std::io::Error> {
     app.at("/metrics")
         .get(|req: tide::Request<State>| 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")

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<W, M>(writer: &mut W, registry: &Registry<M>) -> 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, &registry_single).unwrap();
+        encode(&mut encoded_combined, &vec![&registry_1, &registry_2]).unwrap();
+
+        assert_eq!(encoded_single, encoded_combined);
+    }
+
     #[test]
     fn encode_counter() {
         let counter: Counter = Counter::default();

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<M = Box<dyn crate::encoding::text::SendSyncEncodeMetric>> {
+pub struct Registry<M = Box<dyn SendSyncEncodeMetric>> {
     prefix: Option<Prefix>,
     labels: Vec<(Cow<'static, str>, Cow<'static, str>)>,
     metrics: Vec<(Descriptor, M)>,
@@ -249,6 +250,35 @@ impl<M> Registry<M> {
     }
 }
 
+/// 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<M>
+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)]