Skip to content

Commit

Permalink
src/collector: Introduce Collector abstraction
Browse files Browse the repository at this point in the history
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
- prometheus#49
- prometheus#29
  • Loading branch information
mxinden committed Aug 29, 2022
1 parent f5018cf commit 5011475
Show file tree
Hide file tree
Showing 4 changed files with 180 additions and 36 deletions.
9 changes: 9 additions & 0 deletions src/collector.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
// TODO: Document

use std::borrow::Cow;

use crate::registry::Descriptor;

pub trait Collector<M: Clone> {
fn collect<'a>(&'a self) -> Box<dyn Iterator<Item = (Cow<Descriptor>, Cow<M>)> + 'a>;
}
2 changes: 1 addition & 1 deletion src/encoding/text.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ 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<W, M: Clone>(writer: &mut W, registry: &Registry<M>) -> Result<(), std::io::Error>
where
W: Write,
M: EncodeMetric,
Expand Down
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@
//!
//! [examples]: https://github.com/prometheus/client_rust/tree/master/examples

pub mod collector;
pub mod encoding;
pub mod metrics;
pub mod registry;
204 changes: 169 additions & 35 deletions src/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

use std::borrow::Cow;

use crate::collector::Collector;

/// A metric registry.
///
/// First off one registers metrics with the registry via
Expand Down Expand Up @@ -57,26 +59,37 @@ use std::borrow::Cow;
/// # "# EOF\n";
/// # assert_eq!(expected, String::from_utf8(buffer).unwrap());
/// ```
#[derive(Debug)]
pub struct Registry<M = Box<dyn crate::encoding::text::SendSyncEncodeMetric>> {
prefix: Option<Prefix>,
labels: Vec<(Cow<'static, str>, Cow<'static, str>)>,
metrics: Vec<(Descriptor, M)>,
collectors: Vec<Box<dyn Collector<M>>>,
sub_registries: Vec<Registry<M>>,
}

impl std::fmt::Debug for Registry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Registry")
.field("prefix", &self.prefix)
.field("labels", &self.labels)
.field("sub_registries", &self.sub_registries)
.finish()
}
}

impl<M> Default for Registry<M> {
fn default() -> Self {
Self {
prefix: None,
labels: Default::default(),
metrics: Default::default(),
collectors: Default::default(),
sub_registries: vec![],
}
}
}

impl<M> Registry<M> {
impl<M: Clone> Registry<M> {
/// Creates a new default [`Registry`] with the given prefix.
pub fn with_prefix(prefix: impl Into<String>) -> Self {
Self {
Expand Down Expand Up @@ -151,22 +164,45 @@ impl<M> Registry<M> {
metric: M,
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, metric));
}

// TODO Document
/// ```
/// # use prometheus_client::metrics::counter::{Atomic as _, Counter};
/// # use prometheus_client::registry::{Descriptor, Registry};
/// # use prometheus_client::collector::Collector;
/// # use std::borrow::Cow;
/// #
/// struct MyCollector {}
///
/// impl Collector<Counter> for MyCollector {
/// fn collect<'a>(&'a self) -> Box<dyn Iterator<Item = (Cow<Descriptor>, Cow<Counter>)> + 'a> {
/// let c = Counter::default();
/// let descriptor = Descriptor::new(
/// "my_counter",
/// "This is my counter",
/// None,
/// None,
/// vec![],
/// );
/// Box::new(std::iter::once((Cow::Owned(descriptor), Cow::Owned(c))))
/// }
/// }
///
/// let my_collector = Box::new(MyCollector{});
///
/// let mut registry: Registry<Counter> = Registry::default();
///
/// registry.register_collector(my_collector);
/// ```
pub fn register_collector(&mut self, collector: Box<dyn Collector<M>>) {
self.collectors.push(collector);
}

// TODO: Update doc.
/// Create a sub-registry to register metrics with a common prefix.
///
Expand Down Expand Up @@ -238,40 +274,62 @@ impl<M> Registry<M> {
.expect("sub_registries not to be empty.")
}

pub fn iter(&self) -> RegistryIterator<M> {
pub fn iter(&self) -> std::iter::Chain<MetricIterator<M>, CollectorIterator<M>> {
return self.iter_metrics().chain(self.iter_collectors());
}

fn iter_metrics(&self) -> MetricIterator<M> {
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<M> {
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.
#[derive(Debug)]
pub struct RegistryIterator<'a, M> {
pub struct MetricIterator<'a, M> {
metrics: std::slice::Iter<'a, (Descriptor, M)>,
sub_registries: std::slice::Iter<'a, Registry<M>>,
sub_registry: Option<Box<RegistryIterator<'a, M>>>,
sub_registry: Option<Box<MetricIterator<'a, M>>>,
}

impl<'a, M> Iterator for RegistryIterator<'a, M> {
type Item = &'a (Descriptor, M);
impl<'a, M: Clone> Iterator for MetricIterator<'a, M> {
// TODO: Is the cow needed here?
type Item = (Cow<'a, Descriptor>, Cow<'a, M>);

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), Cow::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;
Expand All @@ -282,22 +340,76 @@ impl<'a, M> Iterator for RegistryIterator<'a, M> {
}
}

// TODO: Document that this one enriches the descriptors.
pub struct CollectorIterator<'a, M: Clone> {
prefix: Option<&'a Prefix>,
labels: &'a [(Cow<'static, str>, Cow<'static, str>)],

collector: Option<Box<dyn Iterator<Item = (Cow<'a, Descriptor>, Cow<'a, M>)> + 'a>>,
collectors: std::slice::Iter<'a, Box<dyn Collector<M>>>,

sub_collector_iter: Option<Box<CollectorIterator<'a, M>>>,
sub_registries: std::slice::Iter<'a, Registry<M>>,
}

impl<'a, M: Clone> Iterator for CollectorIterator<'a, M> {
type Item = (Cow<'a, Descriptor>, Cow<'a, M>);

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

#[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)
}
}

#[derive(Debug)]
#[derive(Debug, Clone)]
pub struct Descriptor {
name: String,
help: String,
Expand All @@ -306,6 +418,28 @@ pub struct Descriptor {
}

impl 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, prefix.as_str());
}

let help = help.into() + ".";

Descriptor {
name,
help,
unit,
labels,
}
}

pub fn name(&self) -> &str {
&self.name
}
Expand All @@ -326,7 +460,7 @@ impl Descriptor {
/// Metric units recommended by Open Metrics.
///
/// See [`Unit::Other`] to specify alternative units.
#[derive(Debug)]
#[derive(Debug, Clone)]
pub enum Unit {
Amperes,
Bytes,
Expand Down

0 comments on commit 5011475

Please sign in to comment.