metrics_exporter_prometheus/exporter/
builder.rs

1use std::collections::HashMap;
2#[cfg(any(feature = "push-gateway", feature = "push-gateway-no-tls-provider"))]
3use std::convert::TryFrom;
4#[cfg(feature = "http-listener")]
5use std::net::{IpAddr, Ipv4Addr, SocketAddr};
6use std::num::NonZeroU32;
7use std::sync::RwLock;
8#[cfg(any(
9    feature = "http-listener",
10    feature = "push-gateway",
11    feature = "push-gateway-no-tls-provider"
12))]
13use std::thread;
14use std::time::Duration;
15
16#[cfg(any(feature = "push-gateway", feature = "push-gateway-no-tls-provider"))]
17use hyper::Uri;
18use indexmap::IndexMap;
19#[cfg(feature = "http-listener")]
20use ipnet::IpNet;
21use quanta::Clock;
22
23use metrics_util::{
24    parse_quantiles,
25    registry::{GenerationalStorage, Recency, Registry},
26    MetricKindMask, Quantile,
27};
28
29use crate::common::Matcher;
30use crate::distribution::DistributionBuilder;
31use crate::native_histogram::NativeHistogramConfig;
32use crate::recorder::{Inner, PrometheusRecorder};
33use crate::registry::AtomicStorage;
34use crate::{common::BuildError, PrometheusHandle};
35
36use super::ExporterConfig;
37#[cfg(any(
38    feature = "http-listener",
39    feature = "push-gateway",
40    feature = "push-gateway-no-tls-provider"
41))]
42use super::ExporterFuture;
43
44/// Builder for creating and installing a Prometheus recorder/exporter.
45#[derive(Debug)]
46pub struct PrometheusBuilder {
47    #[cfg_attr(
48        not(any(
49            feature = "http-listener",
50            feature = "push-gateway",
51            feature = "push-gateway-no-tls-provider"
52        )),
53        allow(dead_code)
54    )]
55    exporter_config: ExporterConfig,
56    #[cfg(feature = "http-listener")]
57    allowed_addresses: Option<Vec<IpNet>>,
58    quantiles: Vec<Quantile>,
59    bucket_duration: Option<Duration>,
60    bucket_count: Option<NonZeroU32>,
61    buckets: Option<Vec<f64>>,
62    bucket_overrides: Option<HashMap<Matcher, Vec<f64>>>,
63    native_histogram_overrides: Option<HashMap<Matcher, NativeHistogramConfig>>,
64    idle_timeout: Option<Duration>,
65    upkeep_timeout: Duration,
66    recency_mask: MetricKindMask,
67    global_labels: Option<IndexMap<String, String>>,
68    enable_recommended_naming: bool,
69    /// TODO Remove this field in next version and merge with `enable_recommended_naming`
70    enable_unit_suffix: bool,
71}
72
73impl PrometheusBuilder {
74    /// Creates a new [`PrometheusBuilder`].
75    pub fn new() -> Self {
76        let quantiles = parse_quantiles(&[0.0, 0.5, 0.9, 0.95, 0.99, 0.999, 1.0]);
77
78        #[cfg(feature = "http-listener")]
79        let exporter_config = ExporterConfig::HttpListener {
80            destination: super::ListenDestination::Tcp(SocketAddr::new(
81                IpAddr::V4(Ipv4Addr::UNSPECIFIED),
82                9000,
83            )),
84        };
85        #[cfg(not(feature = "http-listener"))]
86        let exporter_config = ExporterConfig::Unconfigured;
87
88        let upkeep_timeout = Duration::from_secs(5);
89
90        Self {
91            exporter_config,
92            #[cfg(feature = "http-listener")]
93            allowed_addresses: None,
94            quantiles,
95            bucket_duration: None,
96            bucket_count: None,
97            buckets: None,
98            bucket_overrides: None,
99            native_histogram_overrides: None,
100            idle_timeout: None,
101            upkeep_timeout,
102            recency_mask: MetricKindMask::NONE,
103            global_labels: None,
104            enable_recommended_naming: false,
105            enable_unit_suffix: false,
106        }
107    }
108
109    /// Configures the exporter to expose an HTTP listener that functions as a [scrape endpoint].
110    ///
111    /// The HTTP listener that is spawned will respond to GET requests on any request path.
112    ///
113    /// Running in HTTP listener mode is mutually exclusive with the push gateway i.e. enabling the HTTP listener will
114    /// disable the push gateway, and vise versa.
115    ///
116    /// Defaults to enabled, listening at `0.0.0.0:9000`.
117    ///
118    /// [scrape endpoint]: https://prometheus.io/docs/instrumenting/exposition_formats/#text-based-format
119    #[cfg(feature = "http-listener")]
120    #[cfg_attr(docsrs, doc(cfg(feature = "http-listener")))]
121    #[must_use]
122    pub fn with_http_listener(mut self, addr: impl Into<SocketAddr>) -> Self {
123        self.exporter_config = ExporterConfig::HttpListener {
124            destination: super::ListenDestination::Tcp(addr.into()),
125        };
126        self
127    }
128
129    /// Configures the exporter to push periodic requests to a Prometheus [push gateway].
130    ///
131    /// Running in push gateway mode is mutually exclusive with the HTTP listener i.e. enabling the push gateway will
132    /// disable the HTTP listener, and vise versa.
133    ///
134    /// Defaults to disabled.
135    ///
136    /// ## Errors
137    ///
138    /// If the given endpoint cannot be parsed into a valid URI, an error variant will be returned describing the error.
139    ///
140    /// [push gateway]: https://prometheus.io/docs/instrumenting/pushing/
141    #[cfg(any(feature = "push-gateway", feature = "push-gateway-no-tls-provider"))]
142    #[cfg_attr(
143        docsrs,
144        doc(cfg(any(feature = "push-gateway", feature = "push-gateway-no-tls-provider")))
145    )]
146    pub fn with_push_gateway<T>(
147        mut self,
148        endpoint: T,
149        interval: Duration,
150        username: Option<String>,
151        password: Option<String>,
152        use_http_post_method: bool,
153    ) -> Result<Self, BuildError>
154    where
155        T: AsRef<str>,
156    {
157        self.exporter_config = ExporterConfig::PushGateway {
158            endpoint: Uri::try_from(endpoint.as_ref())
159                .map_err(|e| BuildError::InvalidPushGatewayEndpoint(e.to_string()))?,
160            interval,
161            username,
162            password,
163            use_http_post_method,
164        };
165
166        Ok(self)
167    }
168
169    /// Configures the exporter to expose an HTTP listener that functions as a [scrape endpoint], listening on a Unix
170    /// Domain socket at the given path
171    ///
172    /// The HTTP listener that is spawned will respond to GET requests on any request path.
173    ///
174    /// Running in HTTP listener mode is mutually exclusive with the push gateway i.e. enabling the HTTP listener will
175    /// disable the push gateway, and vise versa.
176    ///
177    /// Defaults to disabled.
178    ///
179    /// [scrape endpoint]: https://prometheus.io/docs/instrumenting/exposition_formats/#text-based-format
180    #[cfg(feature = "uds-listener")]
181    #[cfg_attr(docsrs, doc(cfg(feature = "uds-listener")))]
182    #[must_use]
183    pub fn with_http_uds_listener(mut self, addr: impl Into<std::path::PathBuf>) -> Self {
184        self.exporter_config = ExporterConfig::HttpListener {
185            destination: super::ListenDestination::Uds(addr.into()),
186        };
187        self
188    }
189
190    /// Adds an IP address or subnet to the allowlist for the scrape endpoint.
191    ///
192    /// If a client makes a request to the scrape endpoint and their IP is not present in the allowlist, either directly
193    /// or within any of the allowed subnets, they will receive a 403 Forbidden response.
194    ///
195    /// Defaults to allowing all IPs.
196    ///
197    /// ## Security Considerations
198    ///
199    /// On its own, an IP allowlist is insufficient for access control, if the exporter is running in an environment
200    /// alongside applications (such as web browsers) that are susceptible to [DNS
201    /// rebinding](https://en.wikipedia.org/wiki/DNS_rebinding) attacks.
202    ///
203    /// ## Errors
204    ///
205    /// If the given address cannot be parsed into an IP address or subnet, an error variant will be returned describing
206    /// the error.
207    #[cfg(feature = "http-listener")]
208    #[cfg_attr(docsrs, doc(cfg(feature = "http-listener")))]
209    pub fn add_allowed_address<A>(mut self, address: A) -> Result<Self, BuildError>
210    where
211        A: AsRef<str>,
212    {
213        use std::str::FromStr;
214
215        let address = IpNet::from_str(address.as_ref())
216            .map_err(|e| BuildError::InvalidAllowlistAddress(e.to_string()))?;
217        self.allowed_addresses.get_or_insert(vec![]).push(address);
218
219        Ok(self)
220    }
221
222    /// Sets the quantiles to use when rendering histograms.
223    ///
224    /// Quantiles represent a scale of 0 to 1, where percentiles represent a scale of 1 to 100, so a quantile of 0.99 is
225    /// the 99th percentile, and a quantile of 0.99 is the 99.9th percentile.
226    ///
227    /// Defaults to a hard-coded set of quantiles: 0.0, 0.5, 0.9, 0.95, 0.99, 0.999, and 1.0. This means that all
228    /// histograms will be exposed as Prometheus summaries.
229    ///
230    /// If buckets are set (via [`set_buckets`][Self::set_buckets] or
231    /// [`set_buckets_for_metric`][Self::set_buckets_for_metric]) then all histograms will be exposed as summaries
232    /// instead.
233    ///
234    /// ## Errors
235    ///
236    /// If `quantiles` is empty, an error variant will be thrown.
237    pub fn set_quantiles(mut self, quantiles: &[f64]) -> Result<Self, BuildError> {
238        if quantiles.is_empty() {
239            return Err(BuildError::EmptyBucketsOrQuantiles);
240        }
241
242        self.quantiles = parse_quantiles(quantiles);
243        Ok(self)
244    }
245
246    /// Sets the bucket width when using summaries.
247    ///
248    /// Summaries are rolling, which means that they are divided into buckets of a fixed duration (width), and older
249    /// buckets are dropped as they age out. This means data from a period as large as the width will be dropped at a
250    /// time.
251    ///
252    /// The total amount of data kept for a summary is the number of buckets times the bucket width.  For example, a
253    /// bucket count of 3 and a bucket width of 20 seconds would mean that 60 seconds of data is kept at most, with the
254    /// oldest 20 second chunk of data being dropped as the summary rolls forward.
255    ///
256    /// Use more buckets with a smaller width to roll off smaller amounts of data at a time, or fewer buckets with a
257    /// larger width to roll it off in larger chunks.
258    ///
259    /// Defaults to 20 seconds.
260    ///
261    /// ## Errors
262    ///
263    /// If the duration given is zero, an error variant will be thrown.
264    pub fn set_bucket_duration(mut self, value: Duration) -> Result<Self, BuildError> {
265        if value.is_zero() {
266            return Err(BuildError::ZeroBucketDuration);
267        }
268
269        self.bucket_duration = Some(value);
270        Ok(self)
271    }
272
273    /// Sets the bucket count when using summaries.
274    ///
275    /// Summaries are rolling, which means that they are divided into buckets of a fixed duration (width), and older
276    /// buckets are dropped as they age out. This means data from a period as large as the width will be dropped at a
277    /// time.
278    ///
279    /// The total amount of data kept for a summary is the number of buckets times the bucket width.  For example, a
280    /// bucket count of 3 and a bucket width of 20 seconds would mean that 60 seconds of data is kept at most, with the
281    /// oldest 20 second chunk of data being dropped as the summary rolls forward.
282    ///
283    /// Use more buckets with a smaller width to roll off smaller amounts of data at a time, or fewer buckets with a
284    /// larger width to roll it off in larger chunks.
285    ///
286    /// Defaults to 3.
287    #[must_use]
288    pub fn set_bucket_count(mut self, count: NonZeroU32) -> Self {
289        self.bucket_count = Some(count);
290        self
291    }
292
293    /// Sets the buckets to use when rendering histograms.
294    ///
295    /// Buckets values represent the higher bound of each buckets.  If buckets are set, then all histograms will be
296    /// rendered as true Prometheus histograms, instead of summaries.
297    ///
298    /// ## Errors
299    ///
300    /// If `values` is empty, an error variant will be thrown.
301    pub fn set_buckets(mut self, values: &[f64]) -> Result<Self, BuildError> {
302        if values.is_empty() {
303            return Err(BuildError::EmptyBucketsOrQuantiles);
304        }
305
306        self.buckets = Some(values.to_vec());
307        Ok(self)
308    }
309
310    /// Sets whether a unit suffix is appended to metric names.
311    ///
312    /// When this is enabled and the [`Unit`][metrics::Unit] of metric is
313    /// given, then the exported metric name will be appended to according to
314    /// the [Prometheus Best Practices](https://prometheus.io/docs/practices/naming/).
315    ///
316    /// Defaults to false.
317    #[must_use]
318    #[deprecated(
319        since = "0.18.0",
320        note = "users should prefer `with_recommended_naming` which automatically enables unit suffixes"
321    )]
322    pub fn set_enable_unit_suffix(mut self, enabled: bool) -> Self {
323        self.enable_unit_suffix = enabled;
324        self
325    }
326
327    /// Enables Prometheus naming best practices for metrics.
328    ///
329    /// When set to `true`, counter names are suffixed with `_total` and unit suffixes are appended to metric names,
330    /// following [Prometheus Best Practices](https://prometheus.io/docs/practices/naming/).
331    ///
332    /// Defaults to `false`.
333    #[must_use]
334    pub fn with_recommended_naming(mut self, enabled: bool) -> Self {
335        self.enable_recommended_naming = enabled;
336        self
337    }
338
339    /// Sets the bucket for a specific pattern.
340    ///
341    /// The match pattern can be a full match (equality), prefix match, or suffix match.  The matchers are applied in
342    /// that order if two or more matchers would apply to a single metric.  That is to say, if a full match and a prefix
343    /// match applied to a metric, the full match would win, and if a prefix match and a suffix match applied to a
344    /// metric, the prefix match would win.
345    ///
346    /// Buckets values represent the higher bound of each buckets.  If buckets are set, then any histograms that match
347    /// will be rendered as true Prometheus histograms, instead of summaries.
348    ///
349    /// This option changes the observer's output of histogram-type metric into summaries.  It only affects matching
350    /// metrics if [`set_buckets`][Self::set_buckets] was not used.
351    ///
352    /// ## Errors
353    ///
354    /// If `values` is empty, an error variant will be thrown.
355    pub fn set_buckets_for_metric(
356        mut self,
357        matcher: Matcher,
358        values: &[f64],
359    ) -> Result<Self, BuildError> {
360        if values.is_empty() {
361            return Err(BuildError::EmptyBucketsOrQuantiles);
362        }
363
364        let buckets = self.bucket_overrides.get_or_insert_with(HashMap::new);
365        buckets.insert(matcher.sanitized(), values.to_vec());
366        Ok(self)
367    }
368
369    /// Sets native histogram configuration for a specific pattern.
370    ///
371    /// The match pattern can be a full match (equality), prefix match, or suffix match.  The matchers are applied in
372    /// that order if two or more matchers would apply to a single metric.  That is to say, if a full match and a prefix
373    /// match applied to a metric, the full match would win, and if a prefix match and a suffix match applied to a
374    /// metric, the prefix match would win.
375    ///
376    /// Native histograms use exponential buckets and take precedence over regular histograms and summaries.
377    /// They are only supported in the protobuf format.
378    #[must_use]
379    pub fn set_native_histogram_for_metric(
380        mut self,
381        matcher: Matcher,
382        config: NativeHistogramConfig,
383    ) -> Self {
384        let overrides = self.native_histogram_overrides.get_or_insert_with(HashMap::new);
385        overrides.insert(matcher.sanitized(), config);
386        self
387    }
388
389    /// Sets the idle timeout for metrics.
390    ///
391    /// If a metric hasn't been updated within this timeout, it will be removed from the registry and in turn removed
392    /// from the normal scrape output until the metric is emitted again.  This behavior is driven by requests to
393    /// generate rendered output, and so metrics will not be removed unless a request has been made recently enough to
394    /// prune the idle metrics.
395    ///
396    /// Further, the metric kind "mask" configures which metrics will be considered by the idle timeout.  If the kind of
397    /// a metric being considered for idle timeout is not of a kind represented by the mask, it will not be affected,
398    /// even if it would have otherwise been removed for exceeding the idle timeout.
399    ///
400    /// Refer to the documentation for [`MetricKindMask`](metrics_util::MetricKindMask) for more information on defining
401    /// a metric kind mask.
402    #[must_use]
403    pub fn idle_timeout(mut self, mask: MetricKindMask, timeout: Option<Duration>) -> Self {
404        self.idle_timeout = timeout;
405        self.recency_mask = if self.idle_timeout.is_none() { MetricKindMask::NONE } else { mask };
406        self
407    }
408
409    /// Sets the upkeep interval.
410    ///
411    /// The upkeep task handles periodic maintenance operations, such as draining histogram data, to ensure that all
412    /// recorded data is up-to-date and prevent unbounded memory growth.
413    #[must_use]
414    pub fn upkeep_timeout(mut self, timeout: Duration) -> Self {
415        self.upkeep_timeout = timeout;
416        self
417    }
418
419    /// Adds a global label to this exporter.
420    ///
421    /// Global labels are applied to all metrics.  Labels defined on the metric key itself have precedence over any
422    /// global labels.  If this method is called multiple times, the latest value for a given label key will be used.
423    #[must_use]
424    pub fn add_global_label<K, V>(mut self, key: K, value: V) -> Self
425    where
426        K: Into<String>,
427        V: Into<String>,
428    {
429        let labels = self.global_labels.get_or_insert_with(IndexMap::new);
430        labels.insert(key.into(), value.into());
431        self
432    }
433
434    /// Builds the recorder and exporter and installs them globally.
435    ///
436    /// When called from within a Tokio runtime, the exporter future is spawned directly into the runtime.  Otherwise, a
437    /// new single-threaded Tokio runtime is created on a background thread, and the exporter is spawned there.
438    ///
439    /// ## Errors
440    ///
441    /// If there is an error while either building the recorder and exporter, or installing the recorder and exporter,
442    /// an error variant will be returned describing the error.
443    #[cfg(any(
444        feature = "http-listener",
445        feature = "push-gateway",
446        feature = "push-gateway-no-tls-provider"
447    ))]
448    #[cfg_attr(
449        docsrs,
450        doc(cfg(any(
451            feature = "http-listener",
452            feature = "push-gateway",
453            feature = "push-gateway-no-tls-provider"
454        )))
455    )]
456    pub fn install(self) -> Result<(), BuildError> {
457        use tokio::runtime;
458
459        let recorder = if let Ok(handle) = runtime::Handle::try_current() {
460            let (recorder, exporter) = {
461                let _g = handle.enter();
462                self.build()?
463            };
464
465            handle.spawn(exporter);
466
467            recorder
468        } else {
469            let thread_name =
470                format!("metrics-exporter-prometheus-{}", self.exporter_config.as_type_str());
471
472            let runtime = runtime::Builder::new_current_thread()
473                .enable_all()
474                .build()
475                .map_err(|e| BuildError::FailedToCreateRuntime(e.to_string()))?;
476
477            let (recorder, exporter) = {
478                let _g = runtime.enter();
479                self.build()?
480            };
481
482            thread::Builder::new()
483                .name(thread_name)
484                .spawn(move || runtime.block_on(exporter))
485                .map_err(|e| BuildError::FailedToCreateRuntime(e.to_string()))?;
486
487            recorder
488        };
489
490        metrics::set_global_recorder(recorder)?;
491
492        Ok(())
493    }
494
495    /// Builds the recorder and installs it globally, returning a handle to it.
496    ///
497    /// The handle can be used to generate valid Prometheus scrape endpoint payloads directly.
498    ///
499    /// The caller is responsible for ensuring that upkeep is run periodically. See the **Upkeep and maintenance**
500    /// section in the top-level crate documentation for more information.
501    ///
502    /// ## Errors
503    ///
504    /// If there is an error while building the recorder, or installing the recorder, an error variant will be returned
505    /// describing the error.
506    pub fn install_recorder(self) -> Result<PrometheusHandle, BuildError> {
507        let recorder = self.build_recorder();
508        let handle = recorder.handle();
509
510        metrics::set_global_recorder(recorder)?;
511
512        Ok(handle)
513    }
514
515    /// Builds the recorder and exporter and returns them both.
516    ///
517    /// In most cases, users should prefer to use [`install`][PrometheusBuilder::install] to create and install the
518    /// recorder and exporter automatically for them.  If a caller is combining recorders, or needs to schedule the
519    /// exporter to run in a particular way, this method, or [`build_recorder`][PrometheusBuilder::build_recorder],
520    /// provide the flexibility to do so.
521    ///
522    /// ## Panics
523    ///
524    /// This method must be called from within an existing Tokio runtime or it will panic.
525    ///
526    /// ## Errors
527    ///
528    /// If there is an error while building the recorder and exporter, an error variant will be returned describing the
529    /// error.
530    #[warn(clippy::too_many_lines)]
531    #[cfg(any(
532        feature = "http-listener",
533        feature = "push-gateway",
534        feature = "push-gateway-no-tls-provider"
535    ))]
536    #[cfg_attr(
537        docsrs,
538        doc(cfg(any(
539            feature = "http-listener",
540            feature = "push-gateway",
541            feature = "push-gateway-no-tls-provider"
542        )))
543    )]
544    #[cfg_attr(not(feature = "http-listener"), allow(unused_mut))]
545    pub fn build(mut self) -> Result<(PrometheusRecorder, ExporterFuture), BuildError> {
546        #[cfg(feature = "http-listener")]
547        let allowed_addresses = self.allowed_addresses.take();
548        let exporter_config = self.exporter_config.clone();
549        let upkeep_timeout = self.upkeep_timeout;
550
551        let recorder = self.build_recorder();
552        let handle = recorder.handle();
553
554        let recorder_handle = handle.clone();
555        tokio::spawn(async move {
556            loop {
557                tokio::time::sleep(upkeep_timeout).await;
558                recorder_handle.run_upkeep();
559            }
560        });
561
562        Ok((
563            recorder,
564            match exporter_config {
565                ExporterConfig::Unconfigured => Err(BuildError::MissingExporterConfiguration)?,
566
567                #[cfg(feature = "http-listener")]
568                ExporterConfig::HttpListener { destination } => match destination {
569                    super::ListenDestination::Tcp(listen_address) => {
570                        super::http_listener::new_http_listener(
571                            handle,
572                            listen_address,
573                            allowed_addresses,
574                        )?
575                    }
576                    #[cfg(feature = "uds-listener")]
577                    super::ListenDestination::Uds(listen_path) => {
578                        super::http_listener::new_http_uds_listener(handle, listen_path)?
579                    }
580                },
581
582                #[cfg(any(feature = "push-gateway", feature = "push-gateway-no-tls-provider"))]
583                ExporterConfig::PushGateway {
584                    endpoint,
585                    interval,
586                    username,
587                    password,
588                    use_http_post_method,
589                } => super::push_gateway::new_push_gateway(
590                    endpoint,
591                    interval,
592                    username,
593                    password,
594                    use_http_post_method,
595                    handle,
596                ),
597            },
598        ))
599    }
600
601    /// Builds the recorder and returns it.
602    ///
603    /// The caller is responsible for ensuring that upkeep is run periodically. See the **Upkeep and maintenance**
604    /// section in the top-level crate documentation for more information.
605    pub fn build_recorder(self) -> PrometheusRecorder {
606        self.build_with_clock(Clock::new())
607    }
608
609    pub(crate) fn build_with_clock(self, clock: Clock) -> PrometheusRecorder {
610        let inner = Inner {
611            registry: Registry::new(GenerationalStorage::new(AtomicStorage)),
612            recency: Recency::new(clock, self.recency_mask, self.idle_timeout),
613            distributions: RwLock::new(HashMap::new()),
614            distribution_builder: DistributionBuilder::new(
615                self.quantiles,
616                self.bucket_duration,
617                self.buckets,
618                self.bucket_count,
619                self.bucket_overrides,
620                self.native_histogram_overrides,
621            ),
622            descriptions: RwLock::new(HashMap::new()),
623            global_labels: self.global_labels.unwrap_or_default(),
624            enable_unit_suffix: self.enable_recommended_naming || self.enable_unit_suffix,
625            counter_suffix: self.enable_recommended_naming.then_some("total"),
626        };
627
628        PrometheusRecorder::from(inner)
629    }
630}
631
632impl Default for PrometheusBuilder {
633    fn default() -> Self {
634        PrometheusBuilder::new()
635    }
636}
637
638#[cfg(test)]
639#[allow(clippy::approx_constant)]
640mod tests {
641    use std::time::Duration;
642
643    use quanta::Clock;
644
645    use metrics::{Key, KeyName, Label, Recorder, Unit};
646    use metrics_util::MetricKindMask;
647
648    use super::{Matcher, PrometheusBuilder};
649
650    static METADATA: metrics::Metadata =
651        metrics::Metadata::new(module_path!(), metrics::Level::INFO, Some(module_path!()));
652
653    #[test]
654    fn test_render() {
655        let recorder =
656            PrometheusBuilder::new().set_quantiles(&[0.0, 1.0]).unwrap().build_recorder();
657
658        let key = Key::from_name("basic_counter");
659        let counter1 = recorder.register_counter(&key, &METADATA);
660        counter1.increment(42);
661
662        let handle = recorder.handle();
663        let rendered = handle.render();
664        let expected_counter = "# TYPE basic_counter counter\nbasic_counter 42\n\n";
665
666        assert_eq!(rendered, expected_counter);
667
668        let labels = vec![Label::new("wutang", "forever")];
669        let key = Key::from_parts("basic_gauge", labels);
670        let gauge1 = recorder.register_gauge(&key, &METADATA);
671        gauge1.set(-3.14);
672        let rendered = handle.render();
673        let expected_gauge = format!(
674            "{expected_counter}# TYPE basic_gauge gauge\nbasic_gauge{{wutang=\"forever\"}} -3.14\n\n",
675        );
676
677        assert_eq!(rendered, expected_gauge);
678
679        let key = Key::from_name("basic_histogram");
680        let histogram1 = recorder.register_histogram(&key, &METADATA);
681        histogram1.record(12.0);
682        let rendered = handle.render();
683
684        let histogram_data = concat!(
685            "# TYPE basic_histogram summary\n",
686            "basic_histogram{quantile=\"0\"} 12\n",
687            "basic_histogram{quantile=\"1\"} 12\n",
688            "basic_histogram_sum 12\n",
689            "basic_histogram_count 1\n",
690            "\n"
691        );
692        let expected_histogram = format!("{expected_gauge}{histogram_data}");
693
694        assert_eq!(rendered, expected_histogram);
695    }
696
697    #[test]
698    fn test_render_recommended_naming_no_unit_or_description() {
699        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
700
701        let key = Key::from_name("basic_counter");
702        let counter = recorder.register_counter(&key, &METADATA);
703        counter.increment(42);
704
705        let handle = recorder.handle();
706        let rendered = handle.render();
707        let expected = "# TYPE basic_counter_total counter\nbasic_counter_total 42\n\n";
708
709        assert_eq!(rendered, expected);
710    }
711
712    #[test]
713    fn test_render_recommended_naming_with_unit_and_description() {
714        // Note: we need to create a new recorder, as the render order is not deterministic
715        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
716
717        let key_name = KeyName::from_const_str("counter_with_unit");
718        let key = Key::from_name(key_name.clone());
719        recorder.describe_counter(key_name, Some(Unit::Bytes), "A counter with a unit".into());
720        let counter = recorder.register_counter(&key, &METADATA);
721        counter.increment(42);
722
723        let handle = recorder.handle();
724        let rendered = handle.render();
725        let expected: &'static str = concat!(
726            "# HELP counter_with_unit_bytes_total A counter with a unit\n",
727            "# TYPE counter_with_unit_bytes_total counter\n",
728            "counter_with_unit_bytes_total 42\n",
729            "\n",
730        );
731        assert_eq!(rendered, expected);
732    }
733
734    #[test]
735    fn test_render_recommended_naming_manual_total_suffix_with_unit() {
736        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
737        let key_name = KeyName::from_const_str("foo_total");
738        let key = Key::from_name(key_name.clone());
739        recorder.describe_counter(key_name, Some(Unit::Bytes), "Some help".into());
740        let counter = recorder.register_counter(&key, &METADATA);
741        counter.increment(42);
742
743        let handle = recorder.handle();
744        let rendered = handle.render();
745        let expected = concat!(
746            "# HELP foo_bytes_total Some help\n",
747            "# TYPE foo_bytes_total counter\n",
748            "foo_bytes_total 42\n",
749            "\n",
750        );
751        assert_eq!(rendered, expected);
752    }
753
754    #[test]
755    fn test_render_recommended_naming_manual_counter_suffixes() {
756        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
757        let key_name = KeyName::from_const_str("foo_bytes_total");
758        let key = Key::from_name(key_name.clone());
759        recorder.describe_counter(key_name, Some(Unit::Bytes), "Some help".into());
760        let counter = recorder.register_counter(&key, &METADATA);
761        counter.increment(42);
762
763        let handle = recorder.handle();
764        let rendered = handle.render();
765        let expected = concat!(
766            "# HELP foo_bytes_total Some help\n",
767            "# TYPE foo_bytes_total counter\n",
768            "foo_bytes_total 42\n",
769            "\n",
770        );
771        assert_eq!(rendered, expected);
772    }
773
774    #[test]
775    fn test_render_recommended_naming_gauge_with_unit_in_name() {
776        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
777
778        let key_name = KeyName::from_const_str("gauge_with_unit_bytes");
779        let key = Key::from_name(key_name.clone());
780        recorder.describe_gauge(key_name, Some(Unit::Bytes), "A gauge with a unit".into());
781        let gauge = recorder.register_gauge(&key, &METADATA);
782        gauge.set(42.0);
783
784        let handle = recorder.handle();
785        let rendered = handle.render();
786        let expected = concat!(
787            "# HELP gauge_with_unit_bytes A gauge with a unit\n",
788            "# TYPE gauge_with_unit_bytes gauge\n",
789            "gauge_with_unit_bytes 42\n",
790            "\n",
791        );
792        assert_eq!(rendered, expected);
793    }
794
795    #[test]
796    fn test_buckets() {
797        const DEFAULT_VALUES: [f64; 3] = [10.0, 100.0, 1000.0];
798        const PREFIX_VALUES: [f64; 3] = [15.0, 105.0, 1005.0];
799        const SUFFIX_VALUES: [f64; 3] = [20.0, 110.0, 1010.0];
800        const FULL_VALUES: [f64; 3] = [25.0, 115.0, 1015.0];
801
802        let recorder = PrometheusBuilder::new()
803            .set_buckets_for_metric(
804                Matcher::Full("metrics.testing foo".to_owned()),
805                &FULL_VALUES[..],
806            )
807            .expect("bounds should not be empty")
808            .set_buckets_for_metric(
809                Matcher::Prefix("metrics.testing".to_owned()),
810                &PREFIX_VALUES[..],
811            )
812            .expect("bounds should not be empty")
813            .set_buckets_for_metric(Matcher::Suffix("foo".to_owned()), &SUFFIX_VALUES[..])
814            .expect("bounds should not be empty")
815            .set_buckets(&DEFAULT_VALUES[..])
816            .expect("bounds should not be empty")
817            .build_recorder();
818
819        let full_key = Key::from_name("metrics.testing_foo");
820        let full_key_histo = recorder.register_histogram(&full_key, &METADATA);
821        full_key_histo.record(FULL_VALUES[0]);
822
823        let prefix_key = Key::from_name("metrics.testing_bar");
824        let prefix_key_histo = recorder.register_histogram(&prefix_key, &METADATA);
825        prefix_key_histo.record(PREFIX_VALUES[1]);
826
827        let suffix_key = Key::from_name("metrics_testin_foo");
828        let suffix_key_histo = recorder.register_histogram(&suffix_key, &METADATA);
829        suffix_key_histo.record(SUFFIX_VALUES[2]);
830
831        let default_key = Key::from_name("metrics.wee");
832        let default_key_histo = recorder.register_histogram(&default_key, &METADATA);
833        default_key_histo.record(DEFAULT_VALUES[2] + 1.0);
834
835        let full_data = concat!(
836            "# TYPE metrics_testing_foo histogram\n",
837            "metrics_testing_foo_bucket{le=\"25\"} 1\n",
838            "metrics_testing_foo_bucket{le=\"115\"} 1\n",
839            "metrics_testing_foo_bucket{le=\"1015\"} 1\n",
840            "metrics_testing_foo_bucket{le=\"+Inf\"} 1\n",
841            "metrics_testing_foo_sum 25\n",
842            "metrics_testing_foo_count 1\n",
843        );
844
845        let prefix_data = concat!(
846            "# TYPE metrics_testing_bar histogram\n",
847            "metrics_testing_bar_bucket{le=\"15\"} 0\n",
848            "metrics_testing_bar_bucket{le=\"105\"} 1\n",
849            "metrics_testing_bar_bucket{le=\"1005\"} 1\n",
850            "metrics_testing_bar_bucket{le=\"+Inf\"} 1\n",
851            "metrics_testing_bar_sum 105\n",
852            "metrics_testing_bar_count 1\n",
853        );
854
855        let suffix_data = concat!(
856            "# TYPE metrics_testin_foo histogram\n",
857            "metrics_testin_foo_bucket{le=\"20\"} 0\n",
858            "metrics_testin_foo_bucket{le=\"110\"} 0\n",
859            "metrics_testin_foo_bucket{le=\"1010\"} 1\n",
860            "metrics_testin_foo_bucket{le=\"+Inf\"} 1\n",
861            "metrics_testin_foo_sum 1010\n",
862            "metrics_testin_foo_count 1\n",
863        );
864
865        let default_data = concat!(
866            "# TYPE metrics_wee histogram\n",
867            "metrics_wee_bucket{le=\"10\"} 0\n",
868            "metrics_wee_bucket{le=\"100\"} 0\n",
869            "metrics_wee_bucket{le=\"1000\"} 0\n",
870            "metrics_wee_bucket{le=\"+Inf\"} 1\n",
871            "metrics_wee_sum 1001\n",
872            "metrics_wee_count 1\n",
873        );
874
875        let handle = recorder.handle();
876        let rendered = handle.render();
877
878        assert!(rendered.contains(full_data));
879        assert!(rendered.contains(prefix_data));
880        assert!(rendered.contains(suffix_data));
881        assert!(rendered.contains(default_data));
882    }
883
884    #[test]
885    fn test_idle_timeout_all() {
886        let (clock, mock) = Clock::mock();
887
888        let recorder = PrometheusBuilder::new()
889            .idle_timeout(MetricKindMask::ALL, Some(Duration::from_secs(10)))
890            .set_quantiles(&[0.0, 1.0])
891            .unwrap()
892            .build_with_clock(clock);
893
894        let key = Key::from_name("basic_counter");
895        let counter1 = recorder.register_counter(&key, &METADATA);
896        counter1.increment(42);
897
898        let key = Key::from_name("basic_gauge");
899        let gauge1 = recorder.register_gauge(&key, &METADATA);
900        gauge1.set(-3.14);
901
902        let key = Key::from_name("basic_histogram");
903        let histo1 = recorder.register_histogram(&key, &METADATA);
904        histo1.record(1.0);
905
906        let handle = recorder.handle();
907        let rendered = handle.render();
908        let expected = concat!(
909            "# TYPE basic_counter counter\n",
910            "basic_counter 42\n\n",
911            "# TYPE basic_gauge gauge\n",
912            "basic_gauge -3.14\n\n",
913            "# TYPE basic_histogram summary\n",
914            "basic_histogram{quantile=\"0\"} 1\n",
915            "basic_histogram{quantile=\"1\"} 1\n",
916            "basic_histogram_sum 1\n",
917            "basic_histogram_count 1\n\n",
918        );
919
920        assert_eq!(rendered, expected);
921
922        mock.increment(Duration::from_secs(9));
923        let rendered = handle.render();
924        assert_eq!(rendered, expected);
925
926        mock.increment(Duration::from_secs(2));
927        let rendered = handle.render();
928        assert_eq!(rendered, "");
929    }
930
931    #[test]
932    fn test_idle_timeout_partial() {
933        let (clock, mock) = Clock::mock();
934
935        let recorder = PrometheusBuilder::new()
936            .idle_timeout(
937                MetricKindMask::COUNTER | MetricKindMask::HISTOGRAM,
938                Some(Duration::from_secs(10)),
939            )
940            .set_quantiles(&[0.0, 1.0])
941            .unwrap()
942            .build_with_clock(clock);
943
944        let key = Key::from_name("basic_counter");
945        let counter1 = recorder.register_counter(&key, &METADATA);
946        counter1.increment(42);
947
948        let key = Key::from_name("basic_gauge");
949        let gauge1 = recorder.register_gauge(&key, &METADATA);
950        gauge1.set(-3.14);
951
952        let key = Key::from_name("basic_histogram");
953        let histo1 = recorder.register_histogram(&key, &METADATA);
954        histo1.record(1.0);
955
956        let handle = recorder.handle();
957        let rendered = handle.render();
958        let expected = concat!(
959            "# TYPE basic_counter counter\n",
960            "basic_counter 42\n\n",
961            "# TYPE basic_gauge gauge\n",
962            "basic_gauge -3.14\n\n",
963            "# TYPE basic_histogram summary\n",
964            "basic_histogram{quantile=\"0\"} 1\n",
965            "basic_histogram{quantile=\"1\"} 1\n",
966            "basic_histogram_sum 1\n",
967            "basic_histogram_count 1\n\n",
968        );
969
970        assert_eq!(rendered, expected);
971
972        mock.increment(Duration::from_secs(9));
973        let rendered = handle.render();
974        assert_eq!(rendered, expected);
975
976        mock.increment(Duration::from_secs(2));
977        let rendered = handle.render();
978
979        let expected = "# TYPE basic_gauge gauge\nbasic_gauge -3.14\n\n";
980        assert_eq!(rendered, expected);
981    }
982
983    #[test]
984    fn test_idle_timeout_staggered_distributions() {
985        let (clock, mock) = Clock::mock();
986
987        let recorder = PrometheusBuilder::new()
988            .idle_timeout(MetricKindMask::ALL, Some(Duration::from_secs(10)))
989            .set_quantiles(&[0.0, 1.0])
990            .unwrap()
991            .build_with_clock(clock);
992
993        let key = Key::from_name("basic_counter");
994        let counter1 = recorder.register_counter(&key, &METADATA);
995        counter1.increment(42);
996
997        let key = Key::from_name("basic_gauge");
998        let gauge1 = recorder.register_gauge(&key, &METADATA);
999        gauge1.set(-3.14);
1000
1001        let key = Key::from_name("basic_histogram");
1002        let histo1 = recorder.register_histogram(&key, &METADATA);
1003        histo1.record(1.0);
1004
1005        let handle = recorder.handle();
1006        let rendered = handle.render();
1007        let expected = concat!(
1008            "# TYPE basic_counter counter\n",
1009            "basic_counter 42\n\n",
1010            "# TYPE basic_gauge gauge\n",
1011            "basic_gauge -3.14\n\n",
1012            "# TYPE basic_histogram summary\n",
1013            "basic_histogram{quantile=\"0\"} 1\n",
1014            "basic_histogram{quantile=\"1\"} 1\n",
1015            "basic_histogram_sum 1\n",
1016            "basic_histogram_count 1\n\n",
1017        );
1018
1019        assert_eq!(rendered, expected);
1020
1021        mock.increment(Duration::from_secs(9));
1022        let rendered = handle.render();
1023        assert_eq!(rendered, expected);
1024
1025        let key = Key::from_parts("basic_histogram", vec![Label::new("type", "special")]);
1026        let histo2 = recorder.register_histogram(&key, &METADATA);
1027        histo2.record(2.0);
1028
1029        let expected_second = concat!(
1030            "# TYPE basic_counter counter\n",
1031            "basic_counter 42\n\n",
1032            "# TYPE basic_gauge gauge\n",
1033            "basic_gauge -3.14\n\n",
1034            "# TYPE basic_histogram summary\n",
1035            "basic_histogram{quantile=\"0\"} 1\n",
1036            "basic_histogram{quantile=\"1\"} 1\n",
1037            "basic_histogram_sum 1\n",
1038            "basic_histogram_count 1\n",
1039            "basic_histogram{type=\"special\",quantile=\"0\"} 2\n",
1040            "basic_histogram{type=\"special\",quantile=\"1\"} 2\n",
1041            "basic_histogram_sum{type=\"special\"} 2\n",
1042            "basic_histogram_count{type=\"special\"} 1\n\n",
1043        );
1044        let rendered = handle.render();
1045        assert_eq!(rendered, expected_second);
1046
1047        let expected_after = concat!(
1048            "# TYPE basic_histogram summary\n",
1049            "basic_histogram{type=\"special\",quantile=\"0\"} 2\n",
1050            "basic_histogram{type=\"special\",quantile=\"1\"} 2\n",
1051            "basic_histogram_sum{type=\"special\"} 2\n",
1052            "basic_histogram_count{type=\"special\"} 1\n\n",
1053        );
1054
1055        mock.increment(Duration::from_secs(2));
1056        let rendered = handle.render();
1057        assert_eq!(rendered, expected_after);
1058    }
1059
1060    #[test]
1061    fn test_idle_timeout_doesnt_remove_recents() {
1062        let (clock, mock) = Clock::mock();
1063
1064        let recorder = PrometheusBuilder::new()
1065            .idle_timeout(MetricKindMask::ALL, Some(Duration::from_secs(10)))
1066            .build_with_clock(clock);
1067
1068        let key = Key::from_name("basic_counter");
1069        let counter1 = recorder.register_counter(&key, &METADATA);
1070        counter1.increment(42);
1071
1072        let key = Key::from_name("basic_gauge");
1073        let gauge1 = recorder.register_gauge(&key, &METADATA);
1074        gauge1.set(-3.14);
1075
1076        let handle = recorder.handle();
1077        let rendered = handle.render();
1078        let expected = concat!(
1079            "# TYPE basic_counter counter\n",
1080            "basic_counter 42\n\n",
1081            "# TYPE basic_gauge gauge\n",
1082            "basic_gauge -3.14\n\n",
1083        );
1084
1085        assert_eq!(rendered, expected);
1086
1087        mock.increment(Duration::from_secs(9));
1088        let rendered = handle.render();
1089        assert_eq!(rendered, expected);
1090
1091        let expected_second = concat!(
1092            "# TYPE basic_counter counter\n",
1093            "basic_counter 42\n\n",
1094            "# TYPE basic_gauge gauge\n",
1095            "basic_gauge -3.14\n\n",
1096        );
1097        let rendered = handle.render();
1098        assert_eq!(rendered, expected_second);
1099
1100        counter1.increment(1);
1101
1102        let expected_after = concat!("# TYPE basic_counter counter\n", "basic_counter 43\n\n",);
1103
1104        mock.increment(Duration::from_secs(2));
1105        let rendered = handle.render();
1106        assert_eq!(rendered, expected_after);
1107    }
1108
1109    #[test]
1110    fn test_idle_timeout_catches_delayed_idle() {
1111        let (clock, mock) = Clock::mock();
1112
1113        let recorder = PrometheusBuilder::new()
1114            .idle_timeout(MetricKindMask::ALL, Some(Duration::from_secs(10)))
1115            .build_with_clock(clock);
1116
1117        let key = Key::from_name("basic_counter");
1118        let counter1 = recorder.register_counter(&key, &METADATA);
1119        counter1.increment(42);
1120
1121        // First render, which starts tracking the counter in the recency state.
1122        let handle = recorder.handle();
1123        let rendered = handle.render();
1124        let expected = concat!("# TYPE basic_counter counter\n", "basic_counter 42\n\n",);
1125
1126        assert_eq!(rendered, expected);
1127
1128        // Now go forward by 9 seconds, which is close but still right unfer the idle timeout.
1129        mock.increment(Duration::from_secs(9));
1130        let rendered = handle.render();
1131        assert_eq!(rendered, expected);
1132
1133        // Now increment the counter and advance time by two seconds: this pushes it over the idle
1134        // timeout threshold, but it should not be removed since it has been updated.
1135        counter1.increment(1);
1136
1137        let expected_after = concat!("# TYPE basic_counter counter\n", "basic_counter 43\n\n",);
1138
1139        mock.increment(Duration::from_secs(2));
1140        let rendered = handle.render();
1141        assert_eq!(rendered, expected_after);
1142
1143        // Now advance by 11 seconds, right past the idle timeout threshold.  We've made no further
1144        // updates to the counter so it should be properly removed this time.
1145        mock.increment(Duration::from_secs(11));
1146        let rendered = handle.render();
1147        assert_eq!(rendered, "");
1148    }
1149
1150    #[test]
1151    pub fn test_global_labels() {
1152        let recorder = PrometheusBuilder::new()
1153            .add_global_label("foo", "foo")
1154            .add_global_label("foo", "bar")
1155            .build_recorder();
1156        let key = Key::from_name("basic_counter");
1157        let counter1 = recorder.register_counter(&key, &METADATA);
1158        counter1.increment(42);
1159
1160        let handle = recorder.handle();
1161        let rendered = handle.render();
1162        let expected_counter = "# TYPE basic_counter counter\nbasic_counter{foo=\"bar\"} 42\n\n";
1163
1164        assert_eq!(rendered, expected_counter);
1165    }
1166
1167    #[test]
1168    pub fn test_global_labels_overrides() {
1169        let recorder = PrometheusBuilder::new().add_global_label("foo", "foo").build_recorder();
1170
1171        let key =
1172            Key::from_name("overridden").with_extra_labels(vec![Label::new("foo", "overridden")]);
1173        let counter1 = recorder.register_counter(&key, &METADATA);
1174        counter1.increment(1);
1175
1176        let handle = recorder.handle();
1177        let rendered = handle.render();
1178        let expected_counter = "# TYPE overridden counter\noverridden{foo=\"overridden\"} 1\n\n";
1179
1180        assert_eq!(rendered, expected_counter);
1181    }
1182
1183    #[test]
1184    pub fn test_sanitized_render() {
1185        let recorder = PrometheusBuilder::new().add_global_label("foo:", "foo").build_recorder();
1186
1187        let key_name = KeyName::from("yee_haw:lets go");
1188        let key = Key::from_name(key_name.clone())
1189            .with_extra_labels(vec![Label::new("øhno", "\"yeet\nies\\\"")]);
1190        recorder.describe_counter(key_name, None, "\"Simplë stuff.\nRëally.\"".into());
1191        let counter1 = recorder.register_counter(&key, &METADATA);
1192        counter1.increment(1);
1193
1194        let handle = recorder.handle();
1195        let rendered = handle.render();
1196        let expected_counter = "# HELP yee_haw:lets_go \"Simplë stuff.\\nRëally.\"\n# TYPE yee_haw:lets_go counter\nyee_haw:lets_go{foo_=\"foo\",_hno=\"\\\"yeet\\nies\\\"\"} 1\n\n";
1197
1198        assert_eq!(rendered, expected_counter);
1199    }
1200}