prometheus_closure_metric/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5// Copyright 2014 The Prometheus Authors
6// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
7
8//! This library implements a `ClosureMetric` for crate `prometheus` whose value
9//! is computed at the time of collection by a provided closure.
10
11// TODO: add example usage once constructor macros are implemented.
12// (For now, look at tests for an example.)
13
14use anyhow::{Result, anyhow};
15use prometheus::{core, proto};
16
17/// A Prometheus metric whose value is computed at collection time by the
18/// provided closure.
19///
20/// WARNING: The provided closure must be fast (~milliseconds or faster), since
21/// it blocks metric collection.
22#[derive(Debug)]
23pub struct ClosureMetric<F> {
24    desc: core::Desc,
25    f: F,
26    value_type: ValueType,
27    label_pairs: Vec<proto::LabelPair>,
28}
29
30impl<F, T> ClosureMetric<F>
31where
32    F: Fn() -> T + Sync + Send,
33    T: core::Number,
34{
35    pub fn new<D: core::Describer>(
36        describer: D,
37        value_type: ValueType,
38        f: F,
39        label_values: &[&str],
40    ) -> Result<Self> {
41        let desc = describer.describe()?;
42        let label_pairs = make_label_pairs(&desc, label_values)?;
43
44        Ok(Self {
45            desc,
46            f,
47            value_type,
48            label_pairs,
49        })
50    }
51
52    pub fn metric(&self) -> proto::Metric {
53        let mut m = proto::Metric::default();
54        m.set_label(self.label_pairs.clone());
55
56        let val = (self.f)().into_f64();
57        match self.value_type {
58            ValueType::Counter => {
59                let mut counter = proto::Counter::default();
60                counter.set_value(val);
61                m.set_counter(counter);
62            }
63            ValueType::Gauge => {
64                let mut gauge = proto::Gauge::default();
65                gauge.set_value(val);
66                m.set_gauge(gauge);
67            }
68        }
69
70        m
71    }
72}
73
74impl<F, T> prometheus::core::Collector for ClosureMetric<F>
75where
76    F: Fn() -> T + Sync + Send,
77    T: core::Number,
78{
79    fn desc(&self) -> Vec<&prometheus::core::Desc> {
80        vec![&self.desc]
81    }
82
83    fn collect(&self) -> Vec<prometheus::proto::MetricFamily> {
84        let mut m = proto::MetricFamily::default();
85        m.set_name(self.desc.fq_name.clone());
86        m.set_help(self.desc.help.clone());
87        m.set_field_type(self.value_type.metric_type());
88        m.set_metric(vec![self.metric()]);
89        vec![m]
90    }
91}
92
93#[derive(Debug, Clone, Copy)]
94pub enum ValueType {
95    Counter,
96    Gauge,
97}
98
99impl ValueType {
100    /// `metric_type` returns the corresponding proto metric type.
101    pub fn metric_type(self) -> proto::MetricType {
102        match self {
103            ValueType::Counter => proto::MetricType::COUNTER,
104            ValueType::Gauge => proto::MetricType::GAUGE,
105        }
106    }
107}
108
109pub fn make_label_pairs(desc: &core::Desc, label_values: &[&str]) -> Result<Vec<proto::LabelPair>> {
110    if desc.variable_labels.len() != label_values.len() {
111        return Err(anyhow!("inconsistent cardinality"));
112    }
113
114    let total_len = desc.variable_labels.len() + desc.const_label_pairs.len();
115    if total_len == 0 {
116        return Ok(vec![]);
117    }
118
119    if desc.variable_labels.is_empty() {
120        return Ok(desc.const_label_pairs.clone());
121    }
122
123    let mut label_pairs = Vec::with_capacity(total_len);
124    for (i, n) in desc.variable_labels.iter().enumerate() {
125        let mut label_pair = proto::LabelPair::default();
126        label_pair.set_name(n.clone());
127        label_pair.set_value(label_values[i].to_owned());
128        label_pairs.push(label_pair);
129    }
130
131    for label_pair in &desc.const_label_pairs {
132        label_pairs.push(label_pair.clone());
133    }
134    label_pairs.sort();
135    Ok(label_pairs)
136}
137
138// TODO: add and test macros for easier ClosureMetric construction.