src/collector: Introduce Collector abstraction

The `Collector` abstraction allows users to provide additional metrics
and their description on each scrape.

See also:

- https://pkg.go.dev/github.com/prometheus/client_golang/prometheus#hdr-Custom_Collectors_and_constant_Metrics
- #49
- #29

Signed-off-by: Max Inden <[email protected]>

src/collector.rs

@@ -0,0 +1,28 @@
+//! Metric collector implementation.
+//!
+//! See [`Collector`] for details.
+
+use std::borrow::Cow;
+
+use crate::{
+    registry::{Descriptor, Metric},
+    MaybeOwned,
+};
+
+/// The [`Collector`] abstraction allows users to provide additional metrics and
+/// their description on each scrape.
+///
+/// An example use-case is an exporter that retrieves a set of operating system metrics
+/// ad-hoc on each scrape.
+///
+/// Register a [`Collector`] with a [`Registry`](crate::registry::Registry) via
+/// [`Registry::register_collector`](crate::registry::Registry::register_collector).
+pub trait Collector: std::fmt::Debug + Send + Sync + 'static {
+    /// Once the [`Collector`] is registered, this method is called on each scrape.
+    ///
+    /// Note that the return type allows you to either return owned (convenient)
+    /// or borrowed (performant) descriptions and metrics.
+    fn collect<'a>(
+        &'a self,
+    ) -> Box<dyn Iterator<Item = (Cow<'a, Descriptor>, MaybeOwned<'a, Box<dyn Metric>>)> + 'a>;
+}

src/lib.rs

@@ -78,6 +78,32 @@
 //!
 //! [examples]: https://github.com/prometheus/client_rust/tree/master/examples
 
+pub mod collector;
 pub mod encoding;
 pub mod metrics;
 pub mod registry;
+
+/// Represents either borrowed or owned data.
+///
+/// In contrast to [`std::borrow::Cow`] does not require
+/// [`std::borrow::ToOwned`] or [`Clone`]respectively.
+///
+/// Needed for [`collector::Collector`].
+#[derive(Debug)]
+pub enum MaybeOwned<'a, T> {
+    /// Owned data
+    Owned(T),
+    /// Borrowed data
+    Borrowed(&'a T),
+}
+
+impl<'a, T> std::ops::Deref for MaybeOwned<'a, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Owned(t) => &t,
+            Self::Borrowed(t) => t,
+        }
+    }
+}

src/registry.rs

@@ -4,6 +4,9 @@
 
 use std::borrow::Cow;
 
+use crate::collector::Collector;
+use crate::MaybeOwned;
+
 /// A metric registry.
 ///
 /// First off one registers metrics with the registry via
@@ -59,6 +62,7 @@ pub struct Registry {
     prefix: Option<Prefix>,
     labels: Vec<(Cow<'static, str>, Cow<'static, str>)>,
     metrics: Vec<(Descriptor, Box<dyn Metric>)>,
+    collectors: Vec<Box<dyn Collector>>,
     sub_registries: Vec<Registry>,
 }
 
@@ -142,22 +146,49 @@ impl Registry {
         metric: impl Metric,
         unit: Option<Unit>,
     ) {
-        let name = name.into();
-        let help = help.into() + ".";
-        let descriptor = Descriptor {
-            name: self
-                .prefix
-                .as_ref()
-                .map(|p| (p.clone().0 + "_" + name.as_str()))
-                .unwrap_or(name),
-            help,
-            unit,
-            labels: self.labels.clone(),
-        };
+        let descriptor =
+            Descriptor::new(name, help, unit, self.prefix.as_ref(), self.labels.clone());
 
         self.metrics.push((descriptor, Box::new(metric)));
     }
 
+    /// Register a [`Collector`].
+    ///
+    /// ```
+    /// # use prometheus_client::metrics::counter::{Atomic as _, Counter};
+    /// # use prometheus_client::registry::{Descriptor, Metric, Registry};
+    /// # use prometheus_client::collector::Collector;
+    /// # use prometheus_client::MaybeOwned;
+    /// # use std::borrow::Cow;
+    /// #
+    /// #[derive(Debug)]
+    /// struct MyCollector {}
+    ///
+    /// impl Collector for MyCollector {
+    ///   fn collect<'a>(&'a self) -> Box<dyn Iterator<Item = (Cow<Descriptor>, MaybeOwned<'a, Box<dyn Metric>>)> + 'a> {
+    ///     let c: Counter = Counter::default();
+    ///     let c: Box<dyn Metric> = Box::new(c);
+    ///     let descriptor = Descriptor::new(
+    ///       "my_counter",
+    ///       "This is my counter",
+    ///       None,
+    ///       None,
+    ///       vec![],
+    ///     );
+    ///     Box::new(std::iter::once((Cow::Owned(descriptor), MaybeOwned::Owned(c))))
+    ///   }
+    /// }
+    ///
+    /// let my_collector = Box::new(MyCollector{});
+    ///
+    /// let mut registry = Registry::default();
+    ///
+    /// registry.register_collector(my_collector);
+    /// ```
+    pub fn register_collector(&mut self, collector: Box<dyn Collector>) {
+        self.collectors.push(collector);
+    }
+
     /// Create a sub-registry to register metrics with a common prefix.
     ///
     /// Say you would like to prefix one set of metrics with `subsystem_a` and
@@ -229,40 +260,62 @@ impl Registry {
     }
 
     /// [`Iterator`] over all metrics registered with the [`Registry`].
-    pub fn iter(&self) -> RegistryIterator {
+    pub fn iter(&self) -> std::iter::Chain<MetricIterator, CollectorIterator> {
+        return self.iter_metrics().chain(self.iter_collectors());
+    }
+
+    fn iter_metrics(&self) -> MetricIterator {
         let metrics = self.metrics.iter();
         let sub_registries = self.sub_registries.iter();
-        RegistryIterator {
+        MetricIterator {
             metrics,
             sub_registries,
             sub_registry: None,
         }
     }
+
+    fn iter_collectors(&self) -> CollectorIterator {
+        let collectors = self.collectors.iter();
+        let sub_registries = self.sub_registries.iter();
+        CollectorIterator {
+            prefix: self.prefix.as_ref(),
+            labels: &self.labels,
+
+            collector: None,
+            collectors,
+
+            sub_collector_iter: None,
+            sub_registries,
+        }
+    }
 }
 
-/// Iterator iterating both the metrics registered directly with the registry as
-/// well as all metrics registered with sub-registries.
+/// Iterator iterating both the metrics registered directly with the
+/// [`Registry`] as well as all metrics registered with sub [`Registry`]s.
 #[derive(Debug)]
-pub struct RegistryIterator<'a> {
+pub struct MetricIterator<'a> {
     metrics: std::slice::Iter<'a, (Descriptor, Box<dyn Metric>)>,
     sub_registries: std::slice::Iter<'a, Registry>,
-    sub_registry: Option<Box<RegistryIterator<'a>>>,
+    sub_registry: Option<Box<MetricIterator<'a>>>,
 }
 
-impl<'a> Iterator for RegistryIterator<'a> {
-    type Item = &'a (Descriptor, Box<dyn Metric>);
+impl<'a> Iterator for MetricIterator<'a> {
+    type Item = (Cow<'a, Descriptor>, MaybeOwned<'a, Box<dyn Metric>>);
 
     fn next(&mut self) -> Option<Self::Item> {
-        if let Some(metric) = self.metrics.next() {
-            return Some(metric);
-        }
-
         loop {
+            if let Some((descriptor, metric)) = self.metrics.next() {
+                return Some((Cow::Borrowed(descriptor), MaybeOwned::Borrowed(metric)));
+            }
+
             if let Some(metric) = self.sub_registry.as_mut().and_then(|i| i.next()) {
                 return Some(metric);
             }
 
-            self.sub_registry = self.sub_registries.next().map(|r| Box::new(r.iter()));
+            self.sub_registry = self
+                .sub_registries
+                .next()
+                .map(|r| Box::new(r.iter_metrics()));
 
             if self.sub_registry.is_none() {
                 break;
@@ -273,23 +326,89 @@ impl<'a> Iterator for RegistryIterator<'a> {
     }
 }
 
+/// Iterator iterating metrics retrieved from [`Collector`]s registered with the [`Registry`] or sub [`Registry`]s.
+pub struct CollectorIterator<'a> {
+    prefix: Option<&'a Prefix>,
+    labels: &'a [(Cow<'static, str>, Cow<'static, str>)],
+
+    collector: Option<
+        Box<dyn Iterator<Item = (Cow<'a, Descriptor>, MaybeOwned<'a, Box<dyn Metric>>)> + 'a>,
+    >,
+    collectors: std::slice::Iter<'a, Box<dyn Collector>>,
+
+    sub_collector_iter: Option<Box<CollectorIterator<'a>>>,
+    sub_registries: std::slice::Iter<'a, Registry>,
+}
+
+impl<'a> std::fmt::Debug for CollectorIterator<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("CollectorIterator")
+            .field("prefix", &self.prefix)
+            .field("labels", &self.labels)
+            .finish()
+    }
+}
+
+impl<'a> Iterator for CollectorIterator<'a> {
+    type Item = (Cow<'a, Descriptor>, MaybeOwned<'a, Box<dyn Metric>>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some((descriptor, metric)) = self
+                .collector
+                .as_mut()
+                .and_then(|c| c.next())
+                .or_else(|| self.sub_collector_iter.as_mut().and_then(|i| i.next()))
+            {
+                let Descriptor {
+                    name,
+                    help,
+                    unit,
+                    mut labels,
+                } = descriptor.into_owned();
+                labels.extend_from_slice(self.labels);
+                let enriched_descriptor = Descriptor::new(name, help, unit, self.prefix, labels);
+
+                return Some((Cow::Owned(enriched_descriptor), metric));
+            }
+
+            if let Some(collector) = self.collectors.next() {
+                self.collector = Some(collector.collect());
+                continue;
+            }
+
+            if let Some(collector_iter) = self
+                .sub_registries
+                .next()
+                .map(|r| Box::new(r.iter_collectors()))
+            {
+                self.sub_collector_iter = Some(collector_iter);
+                continue;
+            }
+
+            return None;
+        }
+    }
+}
+
+/// Metric prefix
 #[derive(Clone, Debug)]
-struct Prefix(String);
+pub struct Prefix(String);
 
-impl From<String> for Prefix {
-    fn from(s: String) -> Self {
-        Prefix(s)
+impl Prefix {
+    fn as_str(&self) -> &str {
+        self.0.as_str()
     }
 }
 
-impl From<Prefix> for String {
-    fn from(p: Prefix) -> Self {
-        p.0
+impl From<String> for Prefix {
+    fn from(s: String) -> Self {
+        Prefix(s)
     }
 }
 
 /// OpenMetrics metric descriptor.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Descriptor {
     name: String,
     help: String,
@@ -298,6 +417,30 @@ pub struct Descriptor {
 }
 
 impl Descriptor {
+    /// Create new [`Descriptor`].
+    pub fn new<N: Into<String>, H: Into<String>>(
+        name: N,
+        help: H,
+        unit: Option<Unit>,
+        prefix: Option<&Prefix>,
+        labels: Vec<(Cow<'static, str>, Cow<'static, str>)>,
+    ) -> Self {
+        let mut name = name.into();
+        if let Some(prefix) = prefix {
+            name.insert_str(0, "_");
+            name.insert_str(0, prefix.as_str());
+        }
+
+        let help = help.into() + ".";
+
+        Descriptor {
+            name,
+            help,
+            unit,
+            labels,
+        }
+    }
+
     /// Returns the name of the OpenMetrics metric [`Descriptor`].
     pub fn name(&self) -> &str {
         &self.name
@@ -322,7 +465,7 @@ impl Descriptor {
 /// Metric units recommended by Open Metrics.
 ///
 /// See [`Unit::Other`] to specify alternative units.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 #[allow(missing_docs)]
 pub enum Unit {
     Amperes,