prometheus_client/encoding/
text.rs

1//! Open Metrics text format implementation.
2//!
3//! ```
4//! # use prometheus_client::encoding::text::{encode, encode_registry, encode_eof};
5//! # use prometheus_client::metrics::counter::Counter;
6//! # use prometheus_client::registry::Registry;
7//! #
8//! # // Create registry and counter and register the latter with the former.
9//! # let mut registry = Registry::default();
10//! # let counter: Counter = Counter::default();
11//! # registry.register(
12//! #   "my_counter",
13//! #   "This is my counter",
14//! #   counter.clone(),
15//! # );
16//! # counter.inc();
17//! let mut buffer = String::new();
18//!
19//! // Encode the complete OpenMetrics exposition into the message buffer
20//! encode(&mut buffer, &registry).unwrap();
21//! let expected_msg = "# HELP my_counter This is my counter.\n".to_owned() +
22//!                    "# TYPE my_counter counter\n" +
23//!                    "my_counter_total 1\n" +
24//!                    "# EOF\n";
25//! assert_eq!(expected_msg, buffer);
26//! buffer.clear();
27//!
28//! // Encode just the registry into the message buffer
29//! encode_registry(&mut buffer, &registry).unwrap();
30//! let expected_reg = "# HELP my_counter This is my counter.\n".to_owned() +
31//!                    "# TYPE my_counter counter\n" +
32//!                    "my_counter_total 1\n";
33//! assert_eq!(expected_reg, buffer);
34//!
35//! // Encode EOF marker into message buffer to complete the OpenMetrics exposition
36//! encode_eof(&mut buffer).unwrap();
37//! assert_eq!(expected_msg, buffer);
38//! ```
39
40use crate::encoding::{EncodeExemplarValue, EncodeLabelSet, NoLabelSet};
41use crate::metrics::exemplar::Exemplar;
42use crate::metrics::MetricType;
43use crate::registry::{Prefix, Registry, Unit};
44
45use std::borrow::Cow;
46use std::collections::HashMap;
47use std::fmt::Write;
48
49/// Encode both the metrics registered with the provided [`Registry`] and the
50/// EOF marker into the provided [`Write`]r using the OpenMetrics text format.
51///
52/// Note: This function encodes the **complete** OpenMetrics exposition.
53///
54/// Use [`encode_registry`] or [`encode_eof`] if partial encoding is needed.
55///
56/// See [OpenMetrics exposition format](https://github.com/prometheus/OpenMetrics/blob/v1.0.0/specification/OpenMetrics.md#text-format)
57/// for additional details.
58///
59/// # Examples
60///
61/// ```no_run
62/// # use prometheus_client::encoding::text::encode;
63/// # use prometheus_client::metrics::counter::Counter;
64/// # use prometheus_client::metrics::gauge::Gauge;
65/// # use prometheus_client::registry::Registry;
66/// #
67/// // Initialize registry with metric families
68/// let mut registry = Registry::default();
69/// let counter: Counter = Counter::default();
70/// registry.register(
71///     "my_counter",
72///     "This is my counter",
73///     counter.clone(),
74/// );
75/// let gauge: Gauge = Gauge::default();
76/// registry.register(
77///     "my_gauge",
78///     "This is my gauge",
79///     gauge.clone(),
80/// );
81///
82/// // Encode the complete OpenMetrics exposition into the buffer
83/// let mut buffer = String::new();
84/// encode(&mut buffer, &registry)?;
85/// # Ok::<(), std::fmt::Error>(())
86/// ```
87pub fn encode<W>(writer: &mut W, registry: &Registry) -> Result<(), std::fmt::Error>
88where
89    W: Write,
90{
91    encode_registry(writer, registry)?;
92    encode_eof(writer)
93}
94
95/// Encode the metrics registered with the provided [`Registry`] into the
96/// provided [`Write`]r using the OpenMetrics text format.
97///
98/// Note: The OpenMetrics exposition requires that a complete message must end
99/// with an EOF marker.
100///
101/// This function may be called repeatedly for the HTTP scrape response until
102/// [`encode_eof`] signals the end of the response.
103///
104/// This may also be used to compose a partial message with metrics assembled
105/// from multiple registries.
106///
107/// # Examples
108///
109/// ```no_run
110/// # use prometheus_client::encoding::text::encode_registry;
111/// # use prometheus_client::metrics::counter::Counter;
112/// # use prometheus_client::metrics::gauge::Gauge;
113/// # use prometheus_client::registry::Registry;
114/// #
115/// // Initialize registry with a counter
116/// let mut reg_counter = Registry::default();
117/// let counter: Counter = Counter::default();
118/// reg_counter.register(
119///     "my_counter",
120///     "This is my counter",
121///     counter.clone(),
122/// );
123///
124/// // Encode the counter registry into the buffer
125/// let mut buffer = String::new();
126/// encode_registry(&mut buffer, &reg_counter)?;
127///
128/// // Initialize another registry but with a gauge
129/// let mut reg_gauge = Registry::default();
130/// let gauge: Gauge = Gauge::default();
131/// reg_gauge.register(
132///   "my_gauge",
133///   "This is my gauge",
134///   gauge.clone(),
135/// );
136///
137/// // Encode the gauge registry into the buffer
138/// encode_registry(&mut buffer, &reg_gauge)?;
139/// # Ok::<(), std::fmt::Error>(())
140/// ```
141pub fn encode_registry<W>(writer: &mut W, registry: &Registry) -> Result<(), std::fmt::Error>
142where
143    W: Write,
144{
145    registry.encode(&mut DescriptorEncoder::new(writer).into())
146}
147
148/// Encode the EOF marker into the provided [`Write`]r using the OpenMetrics
149/// text format.
150///
151/// Note: This function is used to mark/signal the end of the exposition.
152///
153/// # Examples
154///
155/// ```no_run
156/// # use prometheus_client::encoding::text::{encode_registry, encode_eof};
157/// # use prometheus_client::metrics::counter::Counter;
158/// # use prometheus_client::metrics::gauge::Gauge;
159/// # use prometheus_client::registry::Registry;
160/// #
161/// // Initialize registry with a counter
162/// let mut registry = Registry::default();
163/// let counter: Counter = Counter::default();
164/// registry.register(
165///     "my_counter",
166///     "This is my counter",
167///     counter.clone(),
168/// );
169///
170/// // Encode registry into the buffer
171/// let mut buffer = String::new();
172/// encode_registry(&mut buffer, &registry)?;
173///
174/// // Encode EOF marker to complete the message
175/// encode_eof(&mut buffer)?;
176/// # Ok::<(), std::fmt::Error>(())
177/// ```
178pub fn encode_eof<W>(writer: &mut W) -> Result<(), std::fmt::Error>
179where
180    W: Write,
181{
182    writer.write_str("# EOF\n")
183}
184
185pub(crate) struct DescriptorEncoder<'a> {
186    writer: &'a mut dyn Write,
187    prefix: Option<&'a Prefix>,
188    labels: &'a [(Cow<'static, str>, Cow<'static, str>)],
189}
190
191impl std::fmt::Debug for DescriptorEncoder<'_> {
192    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
193        f.debug_struct("DescriptorEncoder").finish()
194    }
195}
196
197impl DescriptorEncoder<'_> {
198    pub(crate) fn new(writer: &mut dyn Write) -> DescriptorEncoder {
199        DescriptorEncoder {
200            writer,
201            prefix: Default::default(),
202            labels: Default::default(),
203        }
204    }
205
206    pub(crate) fn with_prefix_and_labels<'s>(
207        &'s mut self,
208        prefix: Option<&'s Prefix>,
209        labels: &'s [(Cow<'static, str>, Cow<'static, str>)],
210    ) -> DescriptorEncoder<'s> {
211        DescriptorEncoder {
212            prefix,
213            labels,
214            writer: self.writer,
215        }
216    }
217
218    pub fn encode_descriptor<'s>(
219        &'s mut self,
220        name: &'s str,
221        help: &str,
222        unit: Option<&'s Unit>,
223        metric_type: MetricType,
224    ) -> Result<MetricEncoder<'s>, std::fmt::Error> {
225        self.writer.write_str("# HELP ")?;
226        if let Some(prefix) = self.prefix {
227            self.writer.write_str(prefix.as_str())?;
228            self.writer.write_str("_")?;
229        }
230        self.writer.write_str(name)?;
231        if let Some(unit) = unit {
232            self.writer.write_str("_")?;
233            self.writer.write_str(unit.as_str())?;
234        }
235        self.writer.write_str(" ")?;
236        self.writer.write_str(help)?;
237        self.writer.write_str("\n")?;
238
239        self.writer.write_str("# TYPE ")?;
240        if let Some(prefix) = self.prefix {
241            self.writer.write_str(prefix.as_str())?;
242            self.writer.write_str("_")?;
243        }
244        self.writer.write_str(name)?;
245        if let Some(unit) = unit {
246            self.writer.write_str("_")?;
247            self.writer.write_str(unit.as_str())?;
248        }
249        self.writer.write_str(" ")?;
250        self.writer.write_str(metric_type.as_str())?;
251        self.writer.write_str("\n")?;
252
253        if let Some(unit) = unit {
254            self.writer.write_str("# UNIT ")?;
255            if let Some(prefix) = self.prefix {
256                self.writer.write_str(prefix.as_str())?;
257                self.writer.write_str("_")?;
258            }
259            self.writer.write_str(name)?;
260            self.writer.write_str("_")?;
261            self.writer.write_str(unit.as_str())?;
262            self.writer.write_str(" ")?;
263            self.writer.write_str(unit.as_str())?;
264            self.writer.write_str("\n")?;
265        }
266
267        Ok(MetricEncoder {
268            writer: self.writer,
269            prefix: self.prefix,
270            name,
271            unit,
272            const_labels: self.labels,
273            family_labels: None,
274        })
275    }
276}
277
278/// Helper type for [`EncodeMetric`](super::EncodeMetric), see
279/// [`EncodeMetric::encode`](super::EncodeMetric::encode).
280///
281// `MetricEncoder` does not take a trait parameter for `writer` and `labels`
282// because `EncodeMetric` which uses `MetricEncoder` needs to be usable as a
283// trait object in order to be able to register different metric types with a
284// `Registry`. Trait objects can not use type parameters.
285//
286// TODO: Alternative solutions to the above are very much appreciated.
287pub(crate) struct MetricEncoder<'a> {
288    writer: &'a mut dyn Write,
289    prefix: Option<&'a Prefix>,
290    name: &'a str,
291    unit: Option<&'a Unit>,
292    const_labels: &'a [(Cow<'static, str>, Cow<'static, str>)],
293    family_labels: Option<&'a dyn super::EncodeLabelSet>,
294}
295
296impl std::fmt::Debug for MetricEncoder<'_> {
297    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
298        let mut labels = String::new();
299        if let Some(l) = self.family_labels {
300            l.encode(LabelSetEncoder::new(&mut labels).into())?;
301        }
302
303        f.debug_struct("Encoder")
304            .field("name", &self.name)
305            .field("prefix", &self.prefix)
306            .field("unit", &self.unit)
307            .field("const_labels", &self.const_labels)
308            .field("labels", &labels.as_str())
309            .finish()
310    }
311}
312
313impl MetricEncoder<'_> {
314    pub fn encode_counter<
315        S: EncodeLabelSet,
316        CounterValue: super::EncodeCounterValue,
317        ExemplarValue: EncodeExemplarValue,
318    >(
319        &mut self,
320        v: &CounterValue,
321        exemplar: Option<&Exemplar<S, ExemplarValue>>,
322    ) -> Result<(), std::fmt::Error> {
323        self.write_prefix_name_unit()?;
324
325        self.write_suffix("total")?;
326
327        self.encode_labels::<NoLabelSet>(None)?;
328
329        v.encode(
330            &mut CounterValueEncoder {
331                writer: self.writer,
332            }
333            .into(),
334        )?;
335
336        if let Some(exemplar) = exemplar {
337            self.encode_exemplar(exemplar)?;
338        }
339
340        self.newline()?;
341
342        Ok(())
343    }
344
345    pub fn encode_gauge<GaugeValue: super::EncodeGaugeValue>(
346        &mut self,
347        v: &GaugeValue,
348    ) -> Result<(), std::fmt::Error> {
349        self.write_prefix_name_unit()?;
350
351        self.encode_labels::<NoLabelSet>(None)?;
352
353        v.encode(
354            &mut GaugeValueEncoder {
355                writer: self.writer,
356            }
357            .into(),
358        )?;
359
360        self.newline()?;
361
362        Ok(())
363    }
364
365    pub fn encode_info<S: EncodeLabelSet>(&mut self, label_set: &S) -> Result<(), std::fmt::Error> {
366        self.write_prefix_name_unit()?;
367
368        self.write_suffix("info")?;
369
370        self.encode_labels(Some(label_set))?;
371
372        self.writer.write_str(" ")?;
373        self.writer.write_str(itoa::Buffer::new().format(1))?;
374
375        self.newline()?;
376
377        Ok(())
378    }
379
380    /// Encode a set of labels. Used by wrapper metric types like
381    /// [`Family`](crate::metrics::family::Family).
382    pub fn encode_family<'s, S: EncodeLabelSet>(
383        &'s mut self,
384        label_set: &'s S,
385    ) -> Result<MetricEncoder<'s>, std::fmt::Error> {
386        debug_assert!(self.family_labels.is_none());
387
388        Ok(MetricEncoder {
389            writer: self.writer,
390            prefix: self.prefix,
391            name: self.name,
392            unit: self.unit,
393            const_labels: self.const_labels,
394            family_labels: Some(label_set),
395        })
396    }
397
398    pub fn encode_histogram<S: EncodeLabelSet>(
399        &mut self,
400        sum: f64,
401        count: u64,
402        buckets: &[(f64, u64)],
403        exemplars: Option<&HashMap<usize, Exemplar<S, f64>>>,
404    ) -> Result<(), std::fmt::Error> {
405        self.write_prefix_name_unit()?;
406        self.write_suffix("sum")?;
407        self.encode_labels::<NoLabelSet>(None)?;
408        self.writer.write_str(" ")?;
409        self.writer.write_str(dtoa::Buffer::new().format(sum))?;
410        self.newline()?;
411
412        self.write_prefix_name_unit()?;
413        self.write_suffix("count")?;
414        self.encode_labels::<NoLabelSet>(None)?;
415        self.writer.write_str(" ")?;
416        self.writer.write_str(itoa::Buffer::new().format(count))?;
417        self.newline()?;
418
419        let mut cummulative = 0;
420        for (i, (upper_bound, count)) in buckets.iter().enumerate() {
421            cummulative += count;
422
423            self.write_prefix_name_unit()?;
424            self.write_suffix("bucket")?;
425
426            if *upper_bound == f64::MAX {
427                self.encode_labels(Some(&[("le", "+Inf")]))?;
428            } else {
429                self.encode_labels(Some(&[("le", *upper_bound)]))?;
430            }
431
432            self.writer.write_str(" ")?;
433            self.writer
434                .write_str(itoa::Buffer::new().format(cummulative))?;
435
436            if let Some(exemplar) = exemplars.and_then(|e| e.get(&i)) {
437                self.encode_exemplar(exemplar)?
438            }
439
440            self.newline()?;
441        }
442
443        Ok(())
444    }
445
446    /// Encode an exemplar for the given metric.
447    fn encode_exemplar<S: EncodeLabelSet, V: EncodeExemplarValue>(
448        &mut self,
449        exemplar: &Exemplar<S, V>,
450    ) -> Result<(), std::fmt::Error> {
451        self.writer.write_str(" # {")?;
452        exemplar
453            .label_set
454            .encode(LabelSetEncoder::new(self.writer).into())?;
455        self.writer.write_str("} ")?;
456        exemplar.value.encode(
457            ExemplarValueEncoder {
458                writer: self.writer,
459            }
460            .into(),
461        )?;
462        Ok(())
463    }
464
465    fn newline(&mut self) -> Result<(), std::fmt::Error> {
466        self.writer.write_str("\n")
467    }
468    fn write_prefix_name_unit(&mut self) -> Result<(), std::fmt::Error> {
469        if let Some(prefix) = self.prefix {
470            self.writer.write_str(prefix.as_str())?;
471            self.writer.write_str("_")?;
472        }
473        self.writer.write_str(self.name)?;
474        if let Some(unit) = self.unit {
475            self.writer.write_str("_")?;
476            self.writer.write_str(unit.as_str())?;
477        }
478
479        Ok(())
480    }
481
482    fn write_suffix(&mut self, suffix: &'static str) -> Result<(), std::fmt::Error> {
483        self.writer.write_str("_")?;
484        self.writer.write_str(suffix)?;
485
486        Ok(())
487    }
488
489    // TODO: Consider caching the encoded labels for Histograms as they stay the
490    // same but are currently encoded multiple times.
491    fn encode_labels<S: EncodeLabelSet>(
492        &mut self,
493        additional_labels: Option<&S>,
494    ) -> Result<(), std::fmt::Error> {
495        if self.const_labels.is_empty()
496            && additional_labels.is_none()
497            && self.family_labels.is_none()
498        {
499            return Ok(());
500        }
501
502        self.writer.write_str("{")?;
503
504        self.const_labels
505            .encode(LabelSetEncoder::new(self.writer).into())?;
506
507        if let Some(additional_labels) = additional_labels {
508            if !self.const_labels.is_empty() {
509                self.writer.write_str(",")?;
510            }
511
512            additional_labels.encode(LabelSetEncoder::new(self.writer).into())?;
513        }
514
515        /// Writer impl which prepends a comma on the first call to write output to the wrapped writer
516        struct CommaPrependingWriter<'a> {
517            writer: &'a mut dyn Write,
518            should_prepend: bool,
519        }
520
521        impl Write for CommaPrependingWriter<'_> {
522            fn write_str(&mut self, s: &str) -> std::fmt::Result {
523                if self.should_prepend {
524                    self.writer.write_char(',')?;
525                    self.should_prepend = false;
526                }
527                self.writer.write_str(s)
528            }
529        }
530
531        if let Some(labels) = self.family_labels {
532            // if const labels or additional labels have been written, a comma must be prepended before writing the family labels.
533            // However, it could be the case that the family labels are `Some` and yet empty, so the comma should _only_
534            // be prepended if one of the `Write` methods are actually called when attempting to write the family labels.
535            // Therefore, wrap the writer on `Self` with a CommaPrependingWriter if other labels have been written and
536            // there may be a need to prepend an extra comma before writing additional labels.
537            if !self.const_labels.is_empty() || additional_labels.is_some() {
538                let mut writer = CommaPrependingWriter {
539                    writer: self.writer,
540                    should_prepend: true,
541                };
542                labels.encode(LabelSetEncoder::new(&mut writer).into())?;
543            } else {
544                labels.encode(LabelSetEncoder::new(self.writer).into())?;
545            };
546        }
547
548        self.writer.write_str("}")?;
549
550        Ok(())
551    }
552}
553
554pub(crate) struct CounterValueEncoder<'a> {
555    writer: &'a mut dyn Write,
556}
557
558impl std::fmt::Debug for CounterValueEncoder<'_> {
559    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
560        f.debug_struct("CounterValueEncoder").finish()
561    }
562}
563
564impl CounterValueEncoder<'_> {
565    pub fn encode_f64(&mut self, v: f64) -> Result<(), std::fmt::Error> {
566        self.writer.write_str(" ")?;
567        self.writer.write_str(dtoa::Buffer::new().format(v))?;
568        Ok(())
569    }
570
571    pub fn encode_u64(&mut self, v: u64) -> Result<(), std::fmt::Error> {
572        self.writer.write_str(" ")?;
573        self.writer.write_str(itoa::Buffer::new().format(v))?;
574        Ok(())
575    }
576}
577
578pub(crate) struct GaugeValueEncoder<'a> {
579    writer: &'a mut dyn Write,
580}
581
582impl std::fmt::Debug for GaugeValueEncoder<'_> {
583    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
584        f.debug_struct("GaugeValueEncoder").finish()
585    }
586}
587
588impl GaugeValueEncoder<'_> {
589    pub fn encode_u32(&mut self, v: u32) -> Result<(), std::fmt::Error> {
590        self.writer.write_str(" ")?;
591        self.writer.write_str(itoa::Buffer::new().format(v))?;
592        Ok(())
593    }
594
595    pub fn encode_i64(&mut self, v: i64) -> Result<(), std::fmt::Error> {
596        self.writer.write_str(" ")?;
597        self.writer.write_str(itoa::Buffer::new().format(v))?;
598        Ok(())
599    }
600
601    pub fn encode_f64(&mut self, v: f64) -> Result<(), std::fmt::Error> {
602        self.writer.write_str(" ")?;
603        self.writer.write_str(dtoa::Buffer::new().format(v))?;
604        Ok(())
605    }
606}
607
608pub(crate) struct ExemplarValueEncoder<'a> {
609    writer: &'a mut dyn Write,
610}
611
612impl std::fmt::Debug for ExemplarValueEncoder<'_> {
613    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
614        f.debug_struct("ExemplarValueEncoder").finish()
615    }
616}
617
618impl ExemplarValueEncoder<'_> {
619    pub fn encode(&mut self, v: f64) -> Result<(), std::fmt::Error> {
620        self.writer.write_str(dtoa::Buffer::new().format(v))
621    }
622}
623
624pub(crate) struct LabelSetEncoder<'a> {
625    writer: &'a mut dyn Write,
626    first: bool,
627}
628
629impl std::fmt::Debug for LabelSetEncoder<'_> {
630    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
631        f.debug_struct("LabelSetEncoder")
632            .field("first", &self.first)
633            .finish()
634    }
635}
636
637impl<'a> LabelSetEncoder<'a> {
638    fn new(writer: &'a mut dyn Write) -> Self {
639        Self {
640            writer,
641            first: true,
642        }
643    }
644
645    pub fn encode_label(&mut self) -> LabelEncoder {
646        let first = self.first;
647        self.first = false;
648        LabelEncoder {
649            writer: self.writer,
650            first,
651        }
652    }
653}
654
655pub(crate) struct LabelEncoder<'a> {
656    writer: &'a mut dyn Write,
657    first: bool,
658}
659
660impl std::fmt::Debug for LabelEncoder<'_> {
661    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
662        f.debug_struct("LabelEncoder")
663            .field("first", &self.first)
664            .finish()
665    }
666}
667
668impl LabelEncoder<'_> {
669    pub fn encode_label_key(&mut self) -> Result<LabelKeyEncoder, std::fmt::Error> {
670        if !self.first {
671            self.writer.write_str(",")?;
672        }
673        Ok(LabelKeyEncoder {
674            writer: self.writer,
675        })
676    }
677}
678
679pub(crate) struct LabelKeyEncoder<'a> {
680    writer: &'a mut dyn Write,
681}
682
683impl std::fmt::Debug for LabelKeyEncoder<'_> {
684    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
685        f.debug_struct("LabelKeyEncoder").finish()
686    }
687}
688
689impl<'a> LabelKeyEncoder<'a> {
690    pub fn encode_label_value(self) -> Result<LabelValueEncoder<'a>, std::fmt::Error> {
691        self.writer.write_str("=\"")?;
692        Ok(LabelValueEncoder {
693            writer: self.writer,
694        })
695    }
696}
697
698impl std::fmt::Write for LabelKeyEncoder<'_> {
699    fn write_str(&mut self, s: &str) -> std::fmt::Result {
700        self.writer.write_str(s)
701    }
702}
703
704pub(crate) struct LabelValueEncoder<'a> {
705    writer: &'a mut dyn Write,
706}
707
708impl std::fmt::Debug for LabelValueEncoder<'_> {
709    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
710        f.debug_struct("LabelValueEncoder").finish()
711    }
712}
713
714impl LabelValueEncoder<'_> {
715    pub fn finish(self) -> Result<(), std::fmt::Error> {
716        self.writer.write_str("\"")
717    }
718}
719
720impl std::fmt::Write for LabelValueEncoder<'_> {
721    fn write_str(&mut self, s: &str) -> std::fmt::Result {
722        self.writer.write_str(s)
723    }
724}
725
726#[cfg(test)]
727mod tests {
728    use super::*;
729    use crate::metrics::exemplar::HistogramWithExemplars;
730    use crate::metrics::family::Family;
731    use crate::metrics::gauge::Gauge;
732    use crate::metrics::histogram::{exponential_buckets, Histogram};
733    use crate::metrics::info::Info;
734    use crate::metrics::{counter::Counter, exemplar::CounterWithExemplar};
735    use pyo3::{prelude::*, types::PyModule};
736    use std::borrow::Cow;
737    use std::fmt::Error;
738    use std::sync::atomic::{AtomicI32, AtomicU32};
739
740    #[test]
741    fn encode_counter() {
742        let counter: Counter = Counter::default();
743        let mut registry = Registry::default();
744        registry.register("my_counter", "My counter", counter);
745
746        let counter_f32 = Counter::<f32, AtomicU32>::default();
747        registry.register("f32_counter", "Counter::<f32, AtomicU32>", counter_f32);
748        let counter_u32 = Counter::<u32, AtomicU32>::default();
749        registry.register("u32_counter", "Counter::<u32, AtomicU32>", counter_u32);
750        let mut encoded = String::new();
751
752        encode(&mut encoded, &registry).unwrap();
753
754        parse_with_python_client(encoded);
755    }
756
757    #[test]
758    fn encode_counter_with_unit() {
759        let mut registry = Registry::default();
760        let counter: Counter = Counter::default();
761        registry.register_with_unit("my_counter", "My counter", Unit::Seconds, counter);
762
763        let mut encoded = String::new();
764        encode(&mut encoded, &registry).unwrap();
765
766        let expected = "# HELP my_counter_seconds My counter.\n".to_owned()
767            + "# TYPE my_counter_seconds counter\n"
768            + "# UNIT my_counter_seconds seconds\n"
769            + "my_counter_seconds_total 0\n"
770            + "# EOF\n";
771        assert_eq!(expected, encoded);
772
773        parse_with_python_client(encoded);
774    }
775
776    #[test]
777    fn encode_counter_with_exemplar() {
778        let mut registry = Registry::default();
779
780        let counter_with_exemplar: CounterWithExemplar<Vec<(String, u64)>> =
781            CounterWithExemplar::default();
782        registry.register_with_unit(
783            "my_counter_with_exemplar",
784            "My counter with exemplar",
785            Unit::Seconds,
786            counter_with_exemplar.clone(),
787        );
788
789        counter_with_exemplar.inc_by(1, Some(vec![("user_id".to_string(), 42)]));
790
791        let mut encoded = String::new();
792        encode(&mut encoded, &registry).unwrap();
793
794        let expected = "# HELP my_counter_with_exemplar_seconds My counter with exemplar.\n"
795            .to_owned()
796            + "# TYPE my_counter_with_exemplar_seconds counter\n"
797            + "# UNIT my_counter_with_exemplar_seconds seconds\n"
798            + "my_counter_with_exemplar_seconds_total 1 # {user_id=\"42\"} 1.0\n"
799            + "# EOF\n";
800        assert_eq!(expected, encoded);
801
802        parse_with_python_client(encoded);
803    }
804
805    #[test]
806    fn encode_gauge() {
807        let mut registry = Registry::default();
808        let gauge: Gauge = Gauge::default();
809        registry.register("my_gauge", "My gauge", gauge);
810        let gauge = Gauge::<u32, AtomicU32>::default();
811        registry.register("u32_gauge", "Gauge::<u32, AtomicU32>", gauge);
812
813        let gauge_f32 = Gauge::<f32, AtomicU32>::default();
814        registry.register("f32_gauge", "Gauge::<f32, AtomicU32>", gauge_f32);
815
816        let gauge_i32 = Gauge::<i32, AtomicI32>::default();
817        registry.register("i32_gauge", "Gauge::<i32, AtomicU32>", gauge_i32);
818
819        let mut encoded = String::new();
820
821        encode(&mut encoded, &registry).unwrap();
822
823        parse_with_python_client(encoded);
824    }
825
826    #[test]
827    fn encode_counter_family() {
828        let mut registry = Registry::default();
829        let family = Family::<Vec<(String, String)>, Counter>::default();
830        registry.register("my_counter_family", "My counter family", family.clone());
831
832        family
833            .get_or_create(&vec![
834                ("method".to_string(), "GET".to_string()),
835                ("status".to_string(), "200".to_string()),
836            ])
837            .inc();
838
839        let mut encoded = String::new();
840
841        encode(&mut encoded, &registry).unwrap();
842
843        parse_with_python_client(encoded);
844    }
845
846    #[test]
847    fn encode_counter_family_with_prefix_with_label() {
848        let mut registry = Registry::default();
849        let sub_registry = registry.sub_registry_with_prefix("my_prefix");
850        let sub_sub_registry = sub_registry
851            .sub_registry_with_label((Cow::Borrowed("my_key"), Cow::Borrowed("my_value")));
852        let family = Family::<Vec<(String, String)>, Counter>::default();
853        sub_sub_registry.register("my_counter_family", "My counter family", family.clone());
854
855        family
856            .get_or_create(&vec![
857                ("method".to_string(), "GET".to_string()),
858                ("status".to_string(), "200".to_string()),
859            ])
860            .inc();
861
862        let mut encoded = String::new();
863
864        encode(&mut encoded, &registry).unwrap();
865
866        let expected = "# HELP my_prefix_my_counter_family My counter family.\n"
867            .to_owned()
868            + "# TYPE my_prefix_my_counter_family counter\n"
869            + "my_prefix_my_counter_family_total{my_key=\"my_value\",method=\"GET\",status=\"200\"} 1\n"
870            + "# EOF\n";
871        assert_eq!(expected, encoded);
872
873        parse_with_python_client(encoded);
874    }
875
876    #[test]
877    fn encode_info() {
878        let mut registry = Registry::default();
879        let info = Info::new(vec![("os".to_string(), "GNU/linux".to_string())]);
880        registry.register("my_info_metric", "My info metric", info);
881
882        let mut encoded = String::new();
883        encode(&mut encoded, &registry).unwrap();
884
885        let expected = "# HELP my_info_metric My info metric.\n".to_owned()
886            + "# TYPE my_info_metric info\n"
887            + "my_info_metric_info{os=\"GNU/linux\"} 1\n"
888            + "# EOF\n";
889        assert_eq!(expected, encoded);
890
891        parse_with_python_client(encoded);
892    }
893
894    #[test]
895    fn encode_histogram() {
896        let mut registry = Registry::default();
897        let histogram = Histogram::new(exponential_buckets(1.0, 2.0, 10));
898        registry.register("my_histogram", "My histogram", histogram.clone());
899        histogram.observe(1.0);
900
901        let mut encoded = String::new();
902
903        encode(&mut encoded, &registry).unwrap();
904
905        parse_with_python_client(encoded);
906    }
907
908    #[test]
909    fn encode_histogram_family() {
910        let mut registry = Registry::default();
911        let family =
912            Family::new_with_constructor(|| Histogram::new(exponential_buckets(1.0, 2.0, 10)));
913        registry.register("my_histogram", "My histogram", family.clone());
914        family
915            .get_or_create(&vec![
916                ("method".to_string(), "GET".to_string()),
917                ("status".to_string(), "200".to_string()),
918            ])
919            .observe(1.0);
920
921        let mut encoded = String::new();
922
923        encode(&mut encoded, &registry).unwrap();
924
925        parse_with_python_client(encoded);
926    }
927
928    #[test]
929    fn encode_histogram_family_with_empty_struct_family_labels() {
930        let mut registry = Registry::default();
931        let family =
932            Family::new_with_constructor(|| Histogram::new(exponential_buckets(1.0, 2.0, 10)));
933        registry.register("my_histogram", "My histogram", family.clone());
934
935        #[derive(Eq, PartialEq, Hash, Debug, Clone)]
936        struct EmptyLabels {}
937
938        impl EncodeLabelSet for EmptyLabels {
939            fn encode(&self, _encoder: crate::encoding::LabelSetEncoder) -> Result<(), Error> {
940                Ok(())
941            }
942        }
943
944        family.get_or_create(&EmptyLabels {}).observe(1.0);
945
946        let mut encoded = String::new();
947
948        encode(&mut encoded, &registry).unwrap();
949
950        parse_with_python_client(encoded);
951    }
952
953    #[test]
954    fn encode_histogram_with_exemplars() {
955        let mut registry = Registry::default();
956        let histogram = HistogramWithExemplars::new(exponential_buckets(1.0, 2.0, 10));
957        registry.register("my_histogram", "My histogram", histogram.clone());
958        histogram.observe(1.0, Some([("user_id".to_string(), 42u64)]));
959
960        let mut encoded = String::new();
961        encode(&mut encoded, &registry).unwrap();
962
963        let expected = "# HELP my_histogram My histogram.\n".to_owned()
964            + "# TYPE my_histogram histogram\n"
965            + "my_histogram_sum 1.0\n"
966            + "my_histogram_count 1\n"
967            + "my_histogram_bucket{le=\"1.0\"} 1 # {user_id=\"42\"} 1.0\n"
968            + "my_histogram_bucket{le=\"2.0\"} 1\n"
969            + "my_histogram_bucket{le=\"4.0\"} 1\n"
970            + "my_histogram_bucket{le=\"8.0\"} 1\n"
971            + "my_histogram_bucket{le=\"16.0\"} 1\n"
972            + "my_histogram_bucket{le=\"32.0\"} 1\n"
973            + "my_histogram_bucket{le=\"64.0\"} 1\n"
974            + "my_histogram_bucket{le=\"128.0\"} 1\n"
975            + "my_histogram_bucket{le=\"256.0\"} 1\n"
976            + "my_histogram_bucket{le=\"512.0\"} 1\n"
977            + "my_histogram_bucket{le=\"+Inf\"} 1\n"
978            + "# EOF\n";
979        assert_eq!(expected, encoded);
980
981        parse_with_python_client(encoded);
982    }
983
984    #[test]
985    fn sub_registry_with_prefix_and_label() {
986        let top_level_metric_name = "my_top_level_metric";
987        let mut registry = Registry::default();
988        let counter: Counter = Counter::default();
989        registry.register(top_level_metric_name, "some help", counter.clone());
990
991        let prefix_1 = "prefix_1";
992        let prefix_1_metric_name = "my_prefix_1_metric";
993        let sub_registry = registry.sub_registry_with_prefix(prefix_1);
994        sub_registry.register(prefix_1_metric_name, "some help", counter.clone());
995
996        let prefix_1_1 = "prefix_1_1";
997        let prefix_1_1_metric_name = "my_prefix_1_1_metric";
998        let sub_sub_registry = sub_registry.sub_registry_with_prefix(prefix_1_1);
999        sub_sub_registry.register(prefix_1_1_metric_name, "some help", counter.clone());
1000
1001        let label_1_2 = (Cow::Borrowed("registry"), Cow::Borrowed("1_2"));
1002        let prefix_1_2_metric_name = "my_prefix_1_2_metric";
1003        let sub_sub_registry = sub_registry.sub_registry_with_label(label_1_2.clone());
1004        sub_sub_registry.register(prefix_1_2_metric_name, "some help", counter.clone());
1005
1006        let labels_1_3 = vec![
1007            (Cow::Borrowed("label_1_3_1"), Cow::Borrowed("value_1_3_1")),
1008            (Cow::Borrowed("label_1_3_2"), Cow::Borrowed("value_1_3_2")),
1009        ];
1010        let prefix_1_3_metric_name = "my_prefix_1_3_metric";
1011        let sub_sub_registry =
1012            sub_registry.sub_registry_with_labels(labels_1_3.clone().into_iter());
1013        sub_sub_registry.register(prefix_1_3_metric_name, "some help", counter.clone());
1014
1015        let prefix_1_3_1 = "prefix_1_3_1";
1016        let prefix_1_3_1_metric_name = "my_prefix_1_3_1_metric";
1017        let sub_sub_sub_registry = sub_sub_registry.sub_registry_with_prefix(prefix_1_3_1);
1018        sub_sub_sub_registry.register(prefix_1_3_1_metric_name, "some help", counter.clone());
1019
1020        let prefix_2 = "prefix_2";
1021        let _ = registry.sub_registry_with_prefix(prefix_2);
1022
1023        let prefix_3 = "prefix_3";
1024        let prefix_3_metric_name = "my_prefix_3_metric";
1025        let sub_registry = registry.sub_registry_with_prefix(prefix_3);
1026        sub_registry.register(prefix_3_metric_name, "some help", counter);
1027
1028        let mut encoded = String::new();
1029        encode(&mut encoded, &registry).unwrap();
1030
1031        let expected = "# HELP my_top_level_metric some help.\n".to_owned()
1032            + "# TYPE my_top_level_metric counter\n"
1033            + "my_top_level_metric_total 0\n"
1034            + "# HELP prefix_1_my_prefix_1_metric some help.\n"
1035            + "# TYPE prefix_1_my_prefix_1_metric counter\n"
1036            + "prefix_1_my_prefix_1_metric_total 0\n"
1037            + "# HELP prefix_1_prefix_1_1_my_prefix_1_1_metric some help.\n"
1038            + "# TYPE prefix_1_prefix_1_1_my_prefix_1_1_metric counter\n"
1039            + "prefix_1_prefix_1_1_my_prefix_1_1_metric_total 0\n"
1040            + "# HELP prefix_1_my_prefix_1_2_metric some help.\n"
1041            + "# TYPE prefix_1_my_prefix_1_2_metric counter\n"
1042            + "prefix_1_my_prefix_1_2_metric_total{registry=\"1_2\"} 0\n"
1043            + "# HELP prefix_1_my_prefix_1_3_metric some help.\n"
1044            + "# TYPE prefix_1_my_prefix_1_3_metric counter\n"
1045            + "prefix_1_my_prefix_1_3_metric_total{label_1_3_1=\"value_1_3_1\",label_1_3_2=\"value_1_3_2\"} 0\n"
1046            + "# HELP prefix_1_prefix_1_3_1_my_prefix_1_3_1_metric some help.\n"
1047            + "# TYPE prefix_1_prefix_1_3_1_my_prefix_1_3_1_metric counter\n"
1048            + "prefix_1_prefix_1_3_1_my_prefix_1_3_1_metric_total{label_1_3_1=\"value_1_3_1\",label_1_3_2=\"value_1_3_2\"} 0\n"
1049            + "# HELP prefix_3_my_prefix_3_metric some help.\n"
1050            + "# TYPE prefix_3_my_prefix_3_metric counter\n"
1051            + "prefix_3_my_prefix_3_metric_total 0\n"
1052            + "# EOF\n";
1053        assert_eq!(expected, encoded);
1054
1055        parse_with_python_client(encoded);
1056    }
1057
1058    #[test]
1059    fn sub_registry_collector() {
1060        use crate::encoding::EncodeMetric;
1061
1062        #[derive(Debug)]
1063        struct Collector {
1064            name: String,
1065        }
1066
1067        impl Collector {
1068            fn new(name: impl Into<String>) -> Self {
1069                Self { name: name.into() }
1070            }
1071        }
1072
1073        impl crate::collector::Collector for Collector {
1074            fn encode(
1075                &self,
1076                mut encoder: crate::encoding::DescriptorEncoder,
1077            ) -> Result<(), std::fmt::Error> {
1078                let counter = crate::metrics::counter::ConstCounter::new(42u64);
1079                let metric_encoder = encoder.encode_descriptor(
1080                    &self.name,
1081                    "some help",
1082                    None,
1083                    counter.metric_type(),
1084                )?;
1085                counter.encode(metric_encoder)?;
1086                Ok(())
1087            }
1088        }
1089
1090        let mut registry = Registry::default();
1091        registry.register_collector(Box::new(Collector::new("top_level")));
1092
1093        let sub_registry = registry.sub_registry_with_prefix("prefix_1");
1094        sub_registry.register_collector(Box::new(Collector::new("sub_level")));
1095
1096        let sub_sub_registry = sub_registry.sub_registry_with_prefix("prefix_1_2");
1097        sub_sub_registry.register_collector(Box::new(Collector::new("sub_sub_level")));
1098
1099        let mut encoded = String::new();
1100        encode(&mut encoded, &registry).unwrap();
1101
1102        let expected = "# HELP top_level some help\n".to_owned()
1103            + "# TYPE top_level counter\n"
1104            + "top_level_total 42\n"
1105            + "# HELP prefix_1_sub_level some help\n"
1106            + "# TYPE prefix_1_sub_level counter\n"
1107            + "prefix_1_sub_level_total 42\n"
1108            + "# HELP prefix_1_prefix_1_2_sub_sub_level some help\n"
1109            + "# TYPE prefix_1_prefix_1_2_sub_sub_level counter\n"
1110            + "prefix_1_prefix_1_2_sub_sub_level_total 42\n"
1111            + "# EOF\n";
1112        assert_eq!(expected, encoded);
1113
1114        parse_with_python_client(encoded);
1115    }
1116
1117    #[test]
1118    fn encode_registry_eof() {
1119        let mut orders_registry = Registry::default();
1120
1121        let total_orders: Counter<u64> = Default::default();
1122        orders_registry.register("orders", "Total orders received", total_orders.clone());
1123        total_orders.inc();
1124
1125        let processing_times = Histogram::new(exponential_buckets(1.0, 2.0, 10));
1126        orders_registry.register_with_unit(
1127            "processing_times",
1128            "Order times",
1129            Unit::Seconds,
1130            processing_times.clone(),
1131        );
1132        processing_times.observe(2.4);
1133
1134        let mut user_auth_registry = Registry::default();
1135
1136        let successful_logins: Counter<u64> = Default::default();
1137        user_auth_registry.register(
1138            "successful_logins",
1139            "Total successful logins",
1140            successful_logins.clone(),
1141        );
1142        successful_logins.inc();
1143
1144        let failed_logins: Counter<u64> = Default::default();
1145        user_auth_registry.register(
1146            "failed_logins",
1147            "Total failed logins",
1148            failed_logins.clone(),
1149        );
1150
1151        let mut response = String::new();
1152
1153        encode_registry(&mut response, &orders_registry).unwrap();
1154        assert_eq!(&response[response.len() - 20..], "bucket{le=\"+Inf\"} 1\n");
1155
1156        encode_registry(&mut response, &user_auth_registry).unwrap();
1157        assert_eq!(&response[response.len() - 20..], "iled_logins_total 0\n");
1158
1159        encode_eof(&mut response).unwrap();
1160        assert_eq!(&response[response.len() - 20..], "ogins_total 0\n# EOF\n");
1161    }
1162
1163    fn parse_with_python_client(input: String) {
1164        pyo3::prepare_freethreaded_python();
1165
1166        println!("{:?}", input);
1167        Python::with_gil(|py| {
1168            let parser = PyModule::from_code_bound(
1169                py,
1170                r#"
1171from prometheus_client.openmetrics.parser import text_string_to_metric_families
1172
1173def parse(input):
1174    families = text_string_to_metric_families(input)
1175    list(families)
1176"#,
1177                "parser.py",
1178                "parser",
1179            )
1180            .map_err(|e| e.to_string())
1181            .unwrap();
1182
1183            parser
1184                .getattr("parse")
1185                .expect("`parse` to exist.")
1186                .call1((input.clone(),))
1187                .map_err(|e| e.to_string())
1188                .unwrap();
1189        })
1190    }
1191}