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 Oct 2, 2022
1 parent ef379f2 commit afe9a3f
Show file tree
Hide file tree
Showing 3 changed files with 223 additions and 33 deletions.
16 changes: 16 additions & 0 deletions src/collector.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
//! TODO

use std::borrow::Cow;

use crate::{
registry::{Descriptor, Metric},
MaybeOwned,
};

/// TODO
pub trait Collector: Send + Sync + 'static {
/// TODO
fn collect<'a>(
&'a self,
) -> Box<dyn Iterator<Item = (Cow<'a, Descriptor>, MaybeOwned<'a, Box<dyn Metric>>)> + 'a>;
}
21 changes: 21 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,27 @@
//!
//! [examples]: https://github.com/prometheus/client_rust/tree/master/examples

pub mod collector;
pub mod encoding;
pub mod metrics;
pub mod registry;

/// TODO
#[derive(Debug)]
pub enum MaybeOwned<'a, T> {
/// TODO
Owned(T),
/// TODO
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,
}
}
}
219 changes: 186 additions & 33 deletions src/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -54,20 +57,31 @@ use std::borrow::Cow;
/// # "# EOF\n";
/// # assert_eq!(expected, String::from_utf8(buffer).unwrap());
/// ```
#[derive(Debug)]
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>,
}

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 Default for Registry {
fn default() -> Self {
Self {
prefix: None,
labels: Default::default(),
metrics: Default::default(),
collectors: Default::default(),
sub_registries: vec![],
}
}
Expand Down Expand Up @@ -153,22 +167,47 @@ 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)));
}

// TODO Document
/// ```
/// # 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;
/// #
/// 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);
}

// TODO: Update doc.
/// Create a sub-registry to register metrics with a common prefix.
///
Expand Down Expand Up @@ -241,40 +280,63 @@ 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.
#[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> {
// TODO: Is the cow needed here?
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;
Expand All @@ -285,23 +347,90 @@ impl<'a> Iterator for RegistryIterator<'a> {
}
}

// TODO: Document that this one enriches the descriptors.
/// TODO
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;
}
}
}

/// TODO
#[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,
Expand All @@ -310,6 +439,30 @@ pub struct Descriptor {
}

impl Descriptor {
/// TODO
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
Expand All @@ -334,7 +487,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,
Expand Down

0 comments on commit afe9a3f

Please sign in to comment.