rust_zero_core/
metric.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    fmt,
4    sync::{Arc, Mutex},
5};
6
7/// Prometheus-compatible metrics registry.
8#[derive(Default)]
9pub struct Metrics {
10    families: Mutex<BTreeMap<String, Arc<MetricFamily>>>,
11}
12
13impl Metrics {
14    pub fn new() -> Self {
15        Self::default()
16    }
17
18    /// Registers a monotonically increasing counter vector.
19    pub fn counter_vec(&self, options: VectorOptions) -> Result<CounterVec, MetricsError> {
20        let family = self.register(options, MetricKind::Counter)?;
21        Ok(CounterVec { family })
22    }
23
24    /// Registers a gauge vector.
25    pub fn gauge_vec(&self, options: VectorOptions) -> Result<GaugeVec, MetricsError> {
26        let family = self.register(options, MetricKind::Gauge)?;
27        Ok(GaugeVec { family })
28    }
29
30    /// Registers a histogram vector with the supplied upper-bound buckets.
31    pub fn histogram_vec(&self, options: HistogramOptions) -> Result<HistogramVec, MetricsError> {
32        let name = options.vector.name();
33        validate_options(&options.vector, &name)?;
34
35        let mut buckets = options.buckets;
36        if buckets.is_empty() {
37            buckets = DEFAULT_HISTOGRAM_BUCKETS.to_vec();
38        }
39        if buckets
40            .iter()
41            .any(|bucket| !bucket.is_finite() || *bucket <= 0.0)
42            || buckets.windows(2).any(|window| window[0] >= window[1])
43        {
44            return Err(MetricsError::InvalidHistogramBuckets(name));
45        }
46
47        let family = Arc::new(MetricFamily {
48            name: name.clone(),
49            help: options.vector.help,
50            labels: options.vector.labels,
51            kind: MetricKind::Histogram,
52            values: Mutex::new(MetricValues::Histogram {
53                buckets,
54                observations: BTreeMap::new(),
55            }),
56        });
57        self.insert(name, Arc::clone(&family))?;
58        Ok(HistogramVec { family })
59    }
60
61    /// Encodes all registered metrics using the Prometheus text exposition format.
62    pub fn render(&self) -> String {
63        let families = self
64            .families
65            .lock()
66            .expect("metrics registry mutex poisoned");
67        let mut output = String::new();
68
69        for family in families.values() {
70            output.push_str("# HELP ");
71            output.push_str(&family.name);
72            output.push(' ');
73            output.push_str(&escape_help(&family.help));
74            output.push('\n');
75            output.push_str("# TYPE ");
76            output.push_str(&family.name);
77            output.push(' ');
78            output.push_str(family.kind.prometheus_name());
79            output.push('\n');
80
81            let values = family.values.lock().expect("metric values mutex poisoned");
82            match &*values {
83                MetricValues::Counter(values) | MetricValues::Gauge(values) => {
84                    for (labels, value) in values {
85                        write_sample(&mut output, &family.name, &family.labels, labels, *value);
86                    }
87                }
88                MetricValues::Histogram {
89                    buckets,
90                    observations,
91                } => {
92                    for (labels, observation) in observations {
93                        let mut count = 0_u64;
94                        for bucket in buckets {
95                            count += observation
96                                .values
97                                .iter()
98                                .filter(|value| **value <= *bucket)
99                                .count() as u64;
100                            write_histogram_bucket(
101                                &mut output,
102                                &family.name,
103                                &family.labels,
104                                labels,
105                                *bucket,
106                                count,
107                            );
108                            count = 0;
109                        }
110                        write_histogram_bucket(
111                            &mut output,
112                            &family.name,
113                            &family.labels,
114                            labels,
115                            f64::INFINITY,
116                            observation.values.len() as u64,
117                        );
118                        write_sample(
119                            &mut output,
120                            &format!("{}_sum", family.name),
121                            &family.labels,
122                            labels,
123                            observation.values.iter().sum(),
124                        );
125                        write_sample(
126                            &mut output,
127                            &format!("{}_count", family.name),
128                            &family.labels,
129                            labels,
130                            observation.values.len() as f64,
131                        );
132                    }
133                }
134            }
135        }
136
137        output
138    }
139
140    fn register(
141        &self,
142        options: VectorOptions,
143        kind: MetricKind,
144    ) -> Result<Arc<MetricFamily>, MetricsError> {
145        let name = options.name();
146        validate_options(&options, &name)?;
147        let family = Arc::new(MetricFamily {
148            name: name.clone(),
149            help: options.help,
150            labels: options.labels,
151            kind,
152            values: Mutex::new(match kind {
153                MetricKind::Counter => MetricValues::Counter(BTreeMap::new()),
154                MetricKind::Gauge => MetricValues::Gauge(BTreeMap::new()),
155                MetricKind::Histogram => unreachable!("histograms use histogram_vec"),
156            }),
157        });
158        self.insert(name, Arc::clone(&family))?;
159        Ok(family)
160    }
161
162    fn insert(&self, name: String, family: Arc<MetricFamily>) -> Result<(), MetricsError> {
163        let mut families = self
164            .families
165            .lock()
166            .expect("metrics registry mutex poisoned");
167        if families.contains_key(&name) {
168            return Err(MetricsError::DuplicateMetric(name));
169        }
170        families.insert(name, family);
171        Ok(())
172    }
173}
174
175/// Shared options for counter and gauge vectors.
176#[derive(Debug, Clone, PartialEq, Eq)]
177pub struct VectorOptions {
178    pub namespace: String,
179    pub subsystem: String,
180    pub name: String,
181    pub help: String,
182    pub labels: Vec<String>,
183}
184
185impl VectorOptions {
186    pub fn new(name: impl Into<String>, help: impl Into<String>) -> Self {
187        Self {
188            namespace: String::new(),
189            subsystem: String::new(),
190            name: name.into(),
191            help: help.into(),
192            labels: Vec::new(),
193        }
194    }
195
196    pub fn with_namespace(mut self, namespace: impl Into<String>) -> Self {
197        self.namespace = namespace.into();
198        self
199    }
200
201    pub fn with_subsystem(mut self, subsystem: impl Into<String>) -> Self {
202        self.subsystem = subsystem.into();
203        self
204    }
205
206    pub fn with_labels<I, S>(mut self, labels: I) -> Self
207    where
208        I: IntoIterator<Item = S>,
209        S: Into<String>,
210    {
211        self.labels = labels.into_iter().map(Into::into).collect();
212        self
213    }
214
215    fn name(&self) -> String {
216        [&self.namespace, &self.subsystem, &self.name]
217            .into_iter()
218            .filter(|part| !part.is_empty())
219            .cloned()
220            .collect::<Vec<_>>()
221            .join("_")
222    }
223}
224
225/// Options for a histogram vector.
226#[derive(Debug, Clone, PartialEq)]
227pub struct HistogramOptions {
228    pub vector: VectorOptions,
229    pub buckets: Vec<f64>,
230}
231
232impl HistogramOptions {
233    pub fn new(name: impl Into<String>, help: impl Into<String>) -> Self {
234        Self {
235            vector: VectorOptions::new(name, help),
236            buckets: DEFAULT_HISTOGRAM_BUCKETS.to_vec(),
237        }
238    }
239
240    pub fn with_vector_options(mut self, options: VectorOptions) -> Self {
241        self.vector = options;
242        self
243    }
244
245    pub fn with_buckets(mut self, buckets: impl Into<Vec<f64>>) -> Self {
246        self.buckets = buckets.into();
247        self
248    }
249}
250
251/// A registered counter vector.
252#[derive(Clone)]
253pub struct CounterVec {
254    family: Arc<MetricFamily>,
255}
256
257impl CounterVec {
258    pub fn inc(&self, labels: &[&str]) -> Result<(), MetricsError> {
259        self.add(1.0, labels)
260    }
261
262    pub fn add(&self, value: f64, labels: &[&str]) -> Result<(), MetricsError> {
263        if !value.is_finite() || value < 0.0 {
264            return Err(MetricsError::InvalidCounterValue(value));
265        }
266        let labels = self.family.validate_labels(labels)?;
267        let mut values = self
268            .family
269            .values
270            .lock()
271            .expect("metric values mutex poisoned");
272        let MetricValues::Counter(values) = &mut *values else {
273            unreachable!("counter vector must contain counter values");
274        };
275        *values.entry(labels).or_default() += value;
276        Ok(())
277    }
278}
279
280/// A registered gauge vector.
281#[derive(Clone)]
282pub struct GaugeVec {
283    family: Arc<MetricFamily>,
284}
285
286impl GaugeVec {
287    pub fn set(&self, value: f64, labels: &[&str]) -> Result<(), MetricsError> {
288        self.update(value, labels, |current, value| *current = value)
289    }
290
291    pub fn inc(&self, labels: &[&str]) -> Result<(), MetricsError> {
292        self.add(1.0, labels)
293    }
294
295    pub fn dec(&self, labels: &[&str]) -> Result<(), MetricsError> {
296        self.add(-1.0, labels)
297    }
298
299    pub fn add(&self, value: f64, labels: &[&str]) -> Result<(), MetricsError> {
300        self.update(value, labels, |current, value| *current += value)
301    }
302
303    pub fn sub(&self, value: f64, labels: &[&str]) -> Result<(), MetricsError> {
304        self.add(-value, labels)
305    }
306
307    fn update(
308        &self,
309        value: f64,
310        labels: &[&str],
311        update: impl FnOnce(&mut f64, f64),
312    ) -> Result<(), MetricsError> {
313        if !value.is_finite() {
314            return Err(MetricsError::InvalidGaugeValue(value));
315        }
316        let labels = self.family.validate_labels(labels)?;
317        let mut values = self
318            .family
319            .values
320            .lock()
321            .expect("metric values mutex poisoned");
322        let MetricValues::Gauge(values) = &mut *values else {
323            unreachable!("gauge vector must contain gauge values");
324        };
325        update(values.entry(labels).or_default(), value);
326        Ok(())
327    }
328}
329
330/// A registered histogram vector.
331#[derive(Clone)]
332pub struct HistogramVec {
333    family: Arc<MetricFamily>,
334}
335
336impl HistogramVec {
337    pub fn observe(&self, value: f64, labels: &[&str]) -> Result<(), MetricsError> {
338        if !value.is_finite() {
339            return Err(MetricsError::InvalidObservation(value));
340        }
341        let labels = self.family.validate_labels(labels)?;
342        let mut values = self
343            .family
344            .values
345            .lock()
346            .expect("metric values mutex poisoned");
347        let MetricValues::Histogram { observations, .. } = &mut *values else {
348            unreachable!("histogram vector must contain histogram values");
349        };
350        observations.entry(labels).or_default().values.push(value);
351        Ok(())
352    }
353}
354
355#[derive(Debug, Clone, PartialEq)]
356pub enum MetricsError {
357    InvalidMetricName(String),
358    InvalidLabelName(String),
359    DuplicateLabel(String),
360    DuplicateMetric(String),
361    LabelCount { expected: usize, actual: usize },
362    InvalidCounterValue(f64),
363    InvalidGaugeValue(f64),
364    InvalidObservation(f64),
365    InvalidHistogramBuckets(String),
366}
367
368impl fmt::Display for MetricsError {
369    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
370        match self {
371            Self::InvalidMetricName(name) => {
372                write!(formatter, "invalid Prometheus metric name: {name}")
373            }
374            Self::InvalidLabelName(name) => {
375                write!(formatter, "invalid Prometheus label name: {name}")
376            }
377            Self::DuplicateLabel(name) => {
378                write!(formatter, "duplicate Prometheus label name: {name}")
379            }
380            Self::DuplicateMetric(name) => {
381                write!(formatter, "metric is already registered: {name}")
382            }
383            Self::LabelCount { expected, actual } => {
384                write!(
385                    formatter,
386                    "metric expects {expected} label values but received {actual}"
387                )
388            }
389            Self::InvalidCounterValue(value) => {
390                write!(
391                    formatter,
392                    "counter values must be finite and non-negative: {value}"
393                )
394            }
395            Self::InvalidGaugeValue(value) => {
396                write!(formatter, "gauge values must be finite: {value}")
397            }
398            Self::InvalidObservation(value) => {
399                write!(formatter, "histogram observations must be finite: {value}")
400            }
401            Self::InvalidHistogramBuckets(name) => {
402                write!(
403                    formatter,
404                    "histogram buckets must be finite, positive, and increasing: {name}"
405                )
406            }
407        }
408    }
409}
410
411impl std::error::Error for MetricsError {}
412
413const DEFAULT_HISTOGRAM_BUCKETS: &[f64] = &[
414    0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
415];
416
417struct MetricFamily {
418    name: String,
419    help: String,
420    labels: Vec<String>,
421    kind: MetricKind,
422    values: Mutex<MetricValues>,
423}
424
425impl MetricFamily {
426    fn validate_labels(&self, labels: &[&str]) -> Result<Vec<String>, MetricsError> {
427        if labels.len() != self.labels.len() {
428            return Err(MetricsError::LabelCount {
429                expected: self.labels.len(),
430                actual: labels.len(),
431            });
432        }
433        Ok(labels.iter().map(|label| (*label).to_owned()).collect())
434    }
435}
436
437#[derive(Clone, Copy)]
438enum MetricKind {
439    Counter,
440    Gauge,
441    Histogram,
442}
443
444impl MetricKind {
445    fn prometheus_name(self) -> &'static str {
446        match self {
447            Self::Counter => "counter",
448            Self::Gauge => "gauge",
449            Self::Histogram => "histogram",
450        }
451    }
452}
453
454enum MetricValues {
455    Counter(BTreeMap<Vec<String>, f64>),
456    Gauge(BTreeMap<Vec<String>, f64>),
457    Histogram {
458        buckets: Vec<f64>,
459        observations: BTreeMap<Vec<String>, HistogramObservation>,
460    },
461}
462
463#[derive(Default)]
464struct HistogramObservation {
465    values: Vec<f64>,
466}
467
468fn validate_options(options: &VectorOptions, name: &str) -> Result<(), MetricsError> {
469    if !is_valid_identifier(name) {
470        return Err(MetricsError::InvalidMetricName(name.to_owned()));
471    }
472
473    let mut labels = BTreeSet::new();
474    for label in &options.labels {
475        if !is_valid_identifier(label) {
476            return Err(MetricsError::InvalidLabelName(label.clone()));
477        }
478        if !labels.insert(label) {
479            return Err(MetricsError::DuplicateLabel(label.clone()));
480        }
481    }
482
483    Ok(())
484}
485
486fn is_valid_identifier(value: &str) -> bool {
487    let mut characters = value.chars();
488    matches!(characters.next(), Some(character) if character == '_' || character.is_ascii_alphabetic())
489        && characters.all(|character| character == '_' || character.is_ascii_alphanumeric())
490}
491
492fn write_sample(
493    output: &mut String,
494    name: &str,
495    label_names: &[String],
496    labels: &[String],
497    value: f64,
498) {
499    output.push_str(name);
500    write_labels(output, label_names, labels, None);
501    output.push(' ');
502    output.push_str(&format_float(value));
503    output.push('\n');
504}
505
506fn write_histogram_bucket(
507    output: &mut String,
508    name: &str,
509    label_names: &[String],
510    labels: &[String],
511    bucket: f64,
512    count: u64,
513) {
514    output.push_str(name);
515    output.push_str("_bucket");
516    write_labels(
517        output,
518        label_names,
519        labels,
520        Some(("le", format_float(bucket))),
521    );
522    output.push(' ');
523    output.push_str(&count.to_string());
524    output.push('\n');
525}
526
527fn write_labels(
528    output: &mut String,
529    label_names: &[String],
530    labels: &[String],
531    extra: Option<(&str, String)>,
532) {
533    if label_names.is_empty() && extra.is_none() {
534        return;
535    }
536
537    output.push('{');
538    for (index, (name, value)) in label_names.iter().zip(labels).enumerate() {
539        if index > 0 {
540            output.push(',');
541        }
542        write_label(output, name, value);
543    }
544    if let Some((name, value)) = extra {
545        if !label_names.is_empty() {
546            output.push(',');
547        }
548        write_label(output, name, &value);
549    }
550    output.push('}');
551}
552
553fn write_label(output: &mut String, name: &str, value: &str) {
554    output.push_str(name);
555    output.push_str("=\"");
556    output.push_str(&escape_label(value));
557    output.push('"');
558}
559
560fn format_float(value: f64) -> String {
561    if value == f64::INFINITY {
562        "+Inf".to_owned()
563    } else {
564        value.to_string()
565    }
566}
567
568fn escape_help(value: &str) -> String {
569    value.replace('\\', "\\\\").replace('\n', "\\n")
570}
571
572fn escape_label(value: &str) -> String {
573    value
574        .replace('\\', "\\\\")
575        .replace('\n', "\\n")
576        .replace('"', "\\\"")
577}
578
579#[cfg(test)]
580mod tests {
581    use super::*;
582
583    #[test]
584    fn renders_labeled_counters_and_gauges() {
585        let metrics = Metrics::new();
586        let requests = metrics
587            .counter_vec(
588                VectorOptions::new("requests_total", "Completed requests")
589                    .with_namespace("users")
590                    .with_labels(["method", "status"]),
591            )
592            .unwrap();
593        let connections = metrics
594            .gauge_vec(VectorOptions::new("connections", "Open connections"))
595            .unwrap();
596
597        requests.inc(&["GET", "200"]).unwrap();
598        requests.add(2.0, &["GET", "200"]).unwrap();
599        connections.set(4.0, &[]).unwrap();
600        connections.dec(&[]).unwrap();
601
602        assert_eq!(
603            metrics.render(),
604            concat!(
605                "# HELP connections Open connections\n",
606                "# TYPE connections gauge\n",
607                "connections 3\n",
608                "# HELP users_requests_total Completed requests\n",
609                "# TYPE users_requests_total counter\n",
610                "users_requests_total{method=\"GET\",status=\"200\"} 3\n",
611            )
612        );
613    }
614
615    #[test]
616    fn renders_cumulative_histogram_buckets() {
617        let metrics = Metrics::new();
618        let duration = metrics
619            .histogram_vec(
620                HistogramOptions::new("request_duration_seconds", "Request duration")
621                    .with_buckets(vec![0.1, 1.0]),
622            )
623            .unwrap();
624
625        duration.observe(0.05, &[]).unwrap();
626        duration.observe(0.5, &[]).unwrap();
627
628        assert_eq!(
629            metrics.render(),
630            concat!(
631                "# HELP request_duration_seconds Request duration\n",
632                "# TYPE request_duration_seconds histogram\n",
633                "request_duration_seconds_bucket{le=\"0.1\"} 1\n",
634                "request_duration_seconds_bucket{le=\"1\"} 2\n",
635                "request_duration_seconds_bucket{le=\"+Inf\"} 2\n",
636                "request_duration_seconds_sum 0.55\n",
637                "request_duration_seconds_count 2\n",
638            )
639        );
640    }
641
642    #[test]
643    fn rejects_invalid_metric_definitions_and_label_counts() {
644        let metrics = Metrics::new();
645
646        assert_eq!(
647            metrics
648                .counter_vec(VectorOptions::new("invalid-name", "Invalid"))
649                .err()
650                .expect("invalid metric name must be rejected"),
651            MetricsError::InvalidMetricName("invalid-name".to_owned())
652        );
653
654        let counter = metrics
655            .counter_vec(VectorOptions::new("events_total", "Events").with_labels(["kind"]))
656            .unwrap();
657        assert_eq!(
658            counter.inc(&[]).unwrap_err(),
659            MetricsError::LabelCount {
660                expected: 1,
661                actual: 0
662            }
663        );
664    }
665}
rust_zero_core/metric.rs

rust_zero_core/
metric.rs
