vise/
wrappers.rs

1//! Wrappers for metric types defined in `prometheus-client`.
2
3use std::{
4    borrow::Borrow,
5    collections::HashMap,
6    fmt,
7    hash::Hash,
8    marker::PhantomData,
9    ops,
10    sync::Arc,
11    time::{Duration, Instant},
12};
13
14use elsa::sync::FrozenMap;
15use once_cell::sync::OnceCell;
16use prometheus_client::{
17    encoding::{
18        EncodeLabelKey, EncodeLabelValue, EncodeMetric, LabelKeyEncoder, LabelValueEncoder,
19        MetricEncoder,
20    },
21    metrics::{
22        counter::Counter, gauge::Gauge as GaugeInner, histogram::Histogram as HistogramInner,
23        MetricType, TypedMetric,
24    },
25    registry::Unit,
26};
27
28use crate::{
29    buckets::Buckets,
30    builder::BuildMetric,
31    encoding::{EncodeGroupedMetric, FullLabelSet, LabelSetWrapper},
32    traits::{EncodeLabelSet, EncodedGaugeValue, GaugeValue, HistogramValue, MapLabels},
33};
34
35/// Label with a unit suffix implementing [`EncodeLabelKey`].
36#[doc(hidden)] // used in proc macros only
37#[derive(Debug)]
38pub struct LabelWithUnit {
39    name: &'static str,
40    unit: Unit,
41}
42
43impl LabelWithUnit {
44    pub const fn new(name: &'static str, unit: Unit) -> Self {
45        Self { name, unit }
46    }
47}
48
49impl EncodeLabelKey for LabelWithUnit {
50    fn encode(&self, encoder: &mut LabelKeyEncoder<'_>) -> fmt::Result {
51        use std::fmt::Write as _;
52
53        write!(encoder, "{}_{}", self.name, self.unit.as_str())
54    }
55}
56
57/// Wraps a [`Duration`] so that it can be used as a label value, which will be set to the fractional
58/// number of seconds in the duration, i.e. [`Duration::as_secs_f64()`]. Mostly useful for [`Info`] metrics.
59#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
60pub struct DurationAsSecs(pub Duration);
61
62impl From<Duration> for DurationAsSecs {
63    fn from(duration: Duration) -> Self {
64        Self(duration)
65    }
66}
67
68impl EncodeLabelValue for DurationAsSecs {
69    fn encode(&self, encoder: &mut LabelValueEncoder) -> fmt::Result {
70        EncodeLabelValue::encode(&self.0.as_secs_f64(), encoder)
71    }
72}
73
74impl<N, A> EncodeGroupedMetric for Counter<N, A> where Self: EncodeMetric + TypedMetric {}
75
76/// Gauge metric.
77///
78/// Gauges are integer or floating-point values that can go up or down. Logically, a reported gauge value
79/// can be treated as valid until the next value is reported.
80///
81/// Gauge values must implement the [`GaugeValue`] trait.
82pub struct Gauge<V: GaugeValue = i64>(GaugeInner<V, V::Atomic>);
83
84impl<V: GaugeValue> fmt::Debug for Gauge<V> {
85    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
86        fmt::Debug::fmt(&self.0, formatter)
87    }
88}
89
90impl<V: GaugeValue> Clone for Gauge<V> {
91    fn clone(&self) -> Self {
92        Self(self.0.clone())
93    }
94}
95
96impl<V: GaugeValue> Default for Gauge<V> {
97    fn default() -> Self {
98        Self(GaugeInner::default())
99    }
100}
101
102impl<V: GaugeValue> Gauge<V> {
103    /// Increases this [`Gauge`] by `v`, returning the previous value.
104    pub fn inc_by(&self, v: V) -> V {
105        self.0.inc_by(v)
106    }
107
108    /// Increases this [`Gauge`] by `v` and returns a guard that will decrement this value back
109    /// when dropped. This can be useful for gauges that measure consumption of a certain resource.
110    pub fn inc_guard(&self, v: V) -> GaugeGuard<V> {
111        let guard = GaugeGuard {
112            gauge: self.clone(),
113            increment: v,
114        };
115        self.0.inc_by(v);
116        guard
117    }
118
119    /// Decreases this [`Gauge`] by `v`, returning the previous value.
120    ///
121    /// # Panics
122    ///
123    /// Depending on the value type, this method may panic on underflow; use with care.
124    pub fn dec_by(&self, v: V) -> V {
125        self.0.dec_by(v)
126    }
127
128    /// Sets the value of this [`Gauge`] returning the previous value.
129    pub fn set(&self, value: V) -> V {
130        self.0.set(value)
131    }
132
133    /// Gets the current value of the gauge.
134    pub fn get(&self) -> V {
135        self.0.get()
136    }
137}
138
139impl<V: GaugeValue> EncodeMetric for Gauge<V> {
140    fn encode(&self, mut encoder: MetricEncoder<'_>) -> fmt::Result {
141        match self.get().encode() {
142            EncodedGaugeValue::I64(value) => encoder.encode_gauge(&value),
143            EncodedGaugeValue::F64(value) => encoder.encode_gauge(&value),
144        }
145    }
146
147    fn metric_type(&self) -> MetricType {
148        <Self as TypedMetric>::TYPE
149    }
150}
151
152impl<V: GaugeValue> TypedMetric for Gauge<V> {
153    const TYPE: MetricType = MetricType::Gauge;
154}
155
156impl<V: GaugeValue> EncodeGroupedMetric for Gauge<V> {}
157
158/// Guard for a [`Gauge`] returned by [`Gauge::inc_guard()`]. When dropped, a guard decrements
159/// the gauge by the same value that it was increased by when creating the guard.
160#[derive(Debug)]
161pub struct GaugeGuard<V: GaugeValue = i64> {
162    gauge: Gauge<V>,
163    increment: V,
164}
165
166impl<V: GaugeValue> Drop for GaugeGuard<V> {
167    fn drop(&mut self) {
168        self.gauge.dec_by(self.increment);
169    }
170}
171
172/// Histogram metric.
173///
174/// Histograms are floating-point values counted in configurable buckets. Logically, a histogram observes
175/// a certain probability distribution, and observations are transient (unlike gauge values).
176///
177/// Histogram values must implement the [`HistogramValue`] trait.
178#[derive(Debug)]
179pub struct Histogram<V: HistogramValue = f64> {
180    inner: HistogramInner,
181    _value: PhantomData<V>,
182}
183
184impl<V: HistogramValue> Clone for Histogram<V> {
185    fn clone(&self) -> Self {
186        Self {
187            inner: self.inner.clone(),
188            _value: PhantomData,
189        }
190    }
191}
192
193impl<V: HistogramValue> Histogram<V> {
194    pub(crate) fn new(buckets: Buckets) -> Self {
195        Self {
196            inner: HistogramInner::new(buckets.iter()),
197            _value: PhantomData,
198        }
199    }
200
201    /// Observes the specified `value` of the metric.
202    pub fn observe(&self, value: V) {
203        self.inner.observe(value.encode());
204    }
205}
206
207impl Histogram<Duration> {
208    /// Starts latency observation for the metric. When the observation is finished,
209    /// call [`LatencyObserver::observe()`].
210    pub fn start(&self) -> LatencyObserver<'_> {
211        LatencyObserver {
212            start: Instant::now(),
213            histogram: self,
214        }
215    }
216}
217
218impl<V: HistogramValue> EncodeMetric for Histogram<V> {
219    fn encode(&self, encoder: MetricEncoder<'_>) -> fmt::Result {
220        self.inner.encode(encoder)
221    }
222
223    fn metric_type(&self) -> MetricType {
224        <Self as TypedMetric>::TYPE
225    }
226}
227
228impl<V: HistogramValue> TypedMetric for Histogram<V> {
229    const TYPE: MetricType = MetricType::Histogram;
230}
231
232impl<V: HistogramValue> EncodeGroupedMetric for Histogram<V> {}
233
234/// Observer of latency for a [`Histogram`].
235#[must_use = "`LatencyObserver` should be `observe()`d"]
236#[derive(Debug)]
237pub struct LatencyObserver<'a> {
238    start: Instant,
239    histogram: &'a Histogram<Duration>,
240}
241
242impl LatencyObserver<'_> {
243    /// Observes and returns the latency passed since this observer was created.
244    pub fn observe(self) -> Duration {
245        let elapsed = self.start.elapsed();
246        self.histogram.observe(elapsed);
247        elapsed
248    }
249}
250
251/// Information metric.
252///
253/// Information metrics represent pieces of information that are not changed during program lifetime
254/// (e.g., config parameters of a certain component).
255#[derive(Debug)]
256pub struct Info<S>(Arc<OnceCell<S>>);
257
258impl<S> Default for Info<S> {
259    fn default() -> Self {
260        Self(Arc::default())
261    }
262}
263
264impl<S> Clone for Info<S> {
265    fn clone(&self) -> Self {
266        Self(self.0.clone())
267    }
268}
269
270impl<S: EncodeLabelSet> Info<S> {
271    /// Gets the current value of the metric.
272    pub fn get(&self) -> Option<&S> {
273        self.0.get()
274    }
275
276    /// Sets the value of this metric.
277    ///
278    /// # Errors
279    ///
280    /// Returns an error if the value is already set.
281    pub fn set(&self, value: S) -> Result<(), SetInfoError<S>> {
282        self.0.set(value).map_err(SetInfoError)
283    }
284}
285
286impl<S: EncodeLabelSet> EncodeMetric for Info<S> {
287    fn encode(&self, mut encoder: MetricEncoder<'_>) -> fmt::Result {
288        if let Some(value) = self.0.get() {
289            encoder.encode_info(&LabelSetWrapper(value))
290        } else {
291            Ok(())
292        }
293    }
294
295    fn metric_type(&self) -> MetricType {
296        MetricType::Info
297    }
298}
299
300impl<S: EncodeLabelSet> TypedMetric for Info<S> {
301    const TYPE: MetricType = MetricType::Info;
302}
303
304impl<S: EncodeLabelSet> EncodeGroupedMetric for Info<S> {}
305
306/// Error returned from [`Info::set()`].
307#[derive(Debug)]
308pub struct SetInfoError<S>(S);
309
310impl<S> SetInfoError<S> {
311    /// Converts the error into the unsuccessfully set value.
312    pub fn into_inner(self) -> S {
313        self.0
314    }
315}
316
317impl<S> fmt::Display for SetInfoError<S> {
318    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
319        formatter.write_str("cannot set info metric value; it is already set")
320    }
321}
322
323pub(crate) struct FamilyInner<S, M: BuildMetric> {
324    map: FrozenMap<S, Box<M>>,
325    builder: M::Builder,
326}
327
328impl<S, M> fmt::Debug for FamilyInner<S, M>
329where
330    S: fmt::Debug + Clone + Eq + Hash,
331    M: BuildMetric + fmt::Debug,
332    M::Builder: fmt::Debug,
333{
334    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
335        let map_keys = self.map.keys_cloned();
336        let map_snapshot: HashMap<_, _> = map_keys
337            .iter()
338            .map(|key| (key, self.map.get(key).unwrap()))
339            .collect();
340
341        formatter
342            .debug_struct("Family")
343            .field("map", &map_snapshot)
344            .field("builder", &self.builder)
345            .finish()
346    }
347}
348
349impl<S: Eq + Hash, M: BuildMetric> FamilyInner<S, M> {
350    pub(crate) fn new(builder: M::Builder) -> Self {
351        Self {
352            map: FrozenMap::new(),
353            builder,
354        }
355    }
356
357    pub(crate) fn get_or_create<Q>(&self, labels: &Q) -> &M
358    where
359        S: Borrow<Q>,
360        Q: Eq + Hash + ?Sized + ToOwned<Owned = S>,
361    {
362        if let Some(metric) = self.map.get(labels) {
363            return metric;
364        }
365        self.map
366            .insert_with(labels.to_owned(), || Box::new(M::build(self.builder)))
367    }
368}
369
370impl<S, M> FamilyInner<S, M>
371where
372    S: Clone + Eq + Hash,
373    M: BuildMetric,
374{
375    pub(crate) fn to_entries(&self) -> impl ExactSizeIterator<Item = (S, &M)> + '_ {
376        let labels = self.map.keys_cloned();
377        labels.into_iter().map(|key| {
378            let metric = self.map.get(&key).unwrap();
379            (key, metric)
380        })
381    }
382}
383
384/// Family of metrics labelled by one or more labels.
385///
386/// Family members can be accessed by indexing.
387pub struct Family<S, M: BuildMetric, L = ()> {
388    inner: Arc<FamilyInner<S, M>>,
389    labels: L,
390}
391
392/// [`Family`] with separately specified label names.
393///
394/// Separately specifying labels allows to not define newtype wrappers for labels. Instead, labels
395/// (the first type param of `LabeledFamily`) can be specified as values (e.g., `&'static str`
396/// or `u8`), and the label names are provided separately using the `labels = [..]` attribute
397/// with the [`Metrics`](macro@crate::Metrics) derive macro.
398///
399/// - If there's a single label, its value type must be specified directly: `&'static str`.
400/// - If there are several labels, they must be specified as a tuple: `(&'static str, u16)`.
401/// - The number of labels must match the number of label names and the constant param of `LabeledFamily`
402///   (which is set to 1 by default). E.g., for two labels you should use `LabeledFamily<_, _, 2>`.
403///
404/// # Examples
405///
406/// ## Family with single label
407///
408/// ```
409/// use vise::{Counter, LabeledFamily, Metrics};
410/// # use vise::{Format, Registry};
411///
412/// #[derive(Debug, Metrics)]
413/// struct TestMetrics {
414///     #[metrics(labels = ["method"])]
415///     counters: LabeledFamily<&'static str, Counter>,
416/// }
417///
418/// // `counters` are keyed by a `&str`:
419/// let metrics = TestMetrics::default();
420/// metrics.counters[&"test"].inc();
421/// metrics.counters[&"another_test"].inc_by(3);
422/// // In the encoded metrics, these entries will be mentioned as follows:
423/// let entries = [
424///     r#"counters_total{method="test"} 1"#,
425///     r#"counters_total{method="another_test"} 3"#,
426/// ];
427/// # let mut registry = Registry::empty();
428/// # registry.register_metrics(&metrics);
429/// # let mut buffer = String::new();
430/// # registry.encode(&mut buffer, Format::OpenMetrics).unwrap();
431/// # for entry in entries {
432/// #     assert!(buffer.contains(&entry), "{buffer}");
433/// # }
434/// ```
435///
436/// ## Family with multiple labels
437///
438/// ```
439/// # use vise::{Buckets, Format, Histogram, LabeledFamily, Metrics, Registry};
440/// # use std::time::Duration;
441/// const LABELS: [&str; 2] = ["method", "code"];
442/// type Labels = (&'static str, u16);
443///
444/// #[derive(Debug, Metrics)]
445/// struct TestMetrics {
446///     #[metrics(labels = LABELS, buckets = Buckets::LATENCIES)]
447///     latencies: LabeledFamily<Labels, Histogram<Duration>, 2>,
448///     // ^ note that label names and type can be extracted elsewhere
449/// }
450///
451/// let metrics = TestMetrics::default();
452/// metrics.latencies[&("call", 200)].observe(Duration::from_millis(25));
453/// metrics.latencies[&("send", 502)].observe(Duration::from_secs(1));
454/// // In the encoded metrics, these entries will be mentioned as follows:
455/// let entries = [
456///     r#"latencies_sum{method="call",code="200"} 0.025"#,
457///     r#"latencies_sum{method="send",code="502"} 1.0"#,
458/// ];
459/// # let mut registry = Registry::empty();
460/// # registry.register_metrics(&metrics);
461/// # let mut buffer = String::new();
462/// # registry.encode(&mut buffer, Format::OpenMetrics).unwrap();
463/// # for entry in entries {
464/// #     assert!(buffer.contains(&entry), "{buffer}");
465/// # }
466/// ```
467pub type LabeledFamily<S, M, const N: usize = 1> = Family<S, M, [&'static str; N]>;
468
469impl<S, M, L> fmt::Debug for Family<S, M, L>
470where
471    S: fmt::Debug + Clone + Eq + Hash,
472    M: BuildMetric + fmt::Debug,
473    M::Builder: fmt::Debug,
474{
475    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
476        fmt::Debug::fmt(&self.inner, formatter)
477    }
478}
479
480impl<S, M: BuildMetric, L: Clone> Clone for Family<S, M, L> {
481    fn clone(&self) -> Self {
482        Self {
483            inner: Arc::clone(&self.inner),
484            labels: self.labels.clone(),
485        }
486    }
487}
488
489impl<S, M, L> Family<S, M, L>
490where
491    S: Clone + Eq + Hash,
492    M: BuildMetric,
493{
494    pub(crate) fn new(builder: M::Builder, labels: L) -> Self {
495        let inner = Arc::new(FamilyInner::new(builder));
496        Self { inner, labels }
497    }
498
499    /// Checks whether this family contains a metric with the specified labels. This is mostly useful
500    /// for testing.
501    pub fn contains(&self, labels: &S) -> bool {
502        self.inner.map.get(labels).is_some()
503    }
504
505    /// Gets a metric with the specified labels if it was reported previously. This is mostly useful
506    /// for testing; use indexing for reporting.
507    pub fn get(&self, labels: &S) -> Option<&M> {
508        self.inner.map.get(labels)
509    }
510
511    /// Gets or creates a metric with the specified labels *lazily* (i.e., on first access). This is useful
512    /// if the metric is updated conditionally and the condition is somewhat rare; in this case, indexing can
513    /// unnecessarily blow up the number of metrics in the family.
514    pub fn get_lazy(&self, labels: S) -> LazyItem<'_, S, M> {
515        LazyItem::new(&self.inner, labels)
516    }
517
518    /// Returns all metrics currently present in this family together with the corresponding labels.
519    /// This is inefficient and mostly useful for testing purposes.
520    pub fn to_entries(&self) -> impl ExactSizeIterator<Item = (S, &M)> + '_ {
521        self.inner.to_entries()
522    }
523}
524
525/// Will create a new metric with the specified labels if it's missing in the family.
526impl<S, M, L, Q> ops::Index<&Q> for Family<S, M, L>
527where
528    S: Borrow<Q> + Eq + Hash,
529    Q: Eq + Hash + ToOwned<Owned = S> + ?Sized,
530    M: BuildMetric,
531{
532    type Output = M;
533
534    fn index(&self, labels: &Q) -> &Self::Output {
535        self.inner.get_or_create(labels)
536    }
537}
538
539/// Lazily accessed member of a [`Family`] or [`MetricsFamily`](crate::MetricsFamily). Returned
540/// by `get_lazy()` methods.
541pub struct LazyItem<'a, S, M: BuildMetric> {
542    family: &'a FamilyInner<S, M>,
543    labels: S,
544}
545
546impl<'a, S, M: BuildMetric> LazyItem<'a, S, M> {
547    pub(crate) fn new(family: &'a FamilyInner<S, M>, labels: S) -> Self {
548        Self { family, labels }
549    }
550}
551
552impl<S, M> fmt::Debug for LazyItem<'_, S, M>
553where
554    S: fmt::Debug + Clone + Eq + Hash,
555    M: BuildMetric + fmt::Debug,
556    M::Builder: fmt::Debug,
557{
558    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
559        formatter
560            .debug_struct("LazyItem")
561            .field("family", self.family)
562            .field("labels", &self.labels)
563            .finish()
564    }
565}
566
567impl<S: Clone, M: BuildMetric> Clone for LazyItem<'_, S, M> {
568    fn clone(&self) -> Self {
569        Self {
570            family: self.family,
571            labels: self.labels.clone(),
572        }
573    }
574}
575
576impl<S, M> ops::Deref for LazyItem<'_, S, M>
577where
578    S: Clone + Eq + Hash,
579    M: BuildMetric,
580{
581    type Target = M;
582
583    fn deref(&self) -> &Self::Target {
584        self.family.get_or_create(&self.labels)
585    }
586}
587
588impl<S, M, L> EncodeMetric for Family<S, M, L>
589where
590    M: BuildMetric + EncodeMetric + TypedMetric,
591    S: Clone + Eq + Hash,
592    L: MapLabels<S>,
593{
594    fn encode(&self, mut encoder: MetricEncoder<'_>) -> fmt::Result {
595        for labels in &self.inner.map.keys_cloned() {
596            let metric = self.inner.map.get(labels).unwrap();
597            let mapped_labels = LabelSetWrapper(self.labels.map_labels(labels));
598            let encoder = encoder.encode_family(&mapped_labels)?;
599            metric.encode(encoder)?;
600        }
601        Ok(())
602    }
603
604    fn metric_type(&self) -> MetricType {
605        <Self as TypedMetric>::TYPE
606    }
607}
608
609impl<S, M: BuildMetric + TypedMetric, L> TypedMetric for Family<S, M, L> {
610    const TYPE: MetricType = <M as TypedMetric>::TYPE;
611}
612
613impl<S, M, L> EncodeGroupedMetric for Family<S, M, L>
614where
615    M: BuildMetric + EncodeMetric + TypedMetric,
616    S: Clone + Eq + Hash,
617    L: MapLabels<S>,
618{
619    fn encode_grouped(
620        &self,
621        group_labels: &dyn EncodeLabelSet,
622        encoder: &mut MetricEncoder<'_>,
623    ) -> fmt::Result {
624        for labels in &self.inner.map.keys_cloned() {
625            let metric = self.inner.map.get(labels).unwrap();
626            let mapped_labels = self.labels.map_labels(labels);
627            let all_labels = FullLabelSet::new(group_labels, &mapped_labels);
628            metric.encode(encoder.encode_family(&all_labels)?)?;
629        }
630        Ok(())
631    }
632}
633
634#[cfg(test)]
635mod tests {
636    use std::{sync::mpsc, thread};
637
638    use prometheus_client::metrics::family::Family as StandardFamily;
639
640    use super::*;
641    use crate::MetricBuilder;
642
643    type Label = (&'static str, &'static str);
644
645    #[test]
646    fn standard_family_is_easy_to_deadlock() {
647        let (stop_sender, stop_receiver) = mpsc::channel();
648        thread::spawn(move || {
649            let family = StandardFamily::<Label, Gauge>::default();
650            let first_metric = family.get_or_create(&("method", "test"));
651            let second_metric = family.get_or_create(&("method", "other"));
652            // ^ The second call will deadlock because of how `Family` is organized internally; its
653            // `get_or_create()` provides a read guard for the internal map, and creating a new metric
654            // requires a write lock on the same map.
655
656            first_metric.set(10);
657            second_metric.set(20);
658            stop_sender.send(()).ok();
659        });
660
661        let err = stop_receiver
662            .recv_timeout(Duration::from_millis(200))
663            .unwrap_err();
664        assert!(matches!(err, mpsc::RecvTimeoutError::Timeout));
665    }
666
667    #[test]
668    fn family_accesses_are_not_deadlocked() {
669        let family = Family::<Label, Gauge>::new(MetricBuilder::new(), ());
670        let first_metric = &family[&("method", "test")];
671        let second_metric = &family[&("method", "other")];
672        first_metric.set(10);
673        second_metric.set(20);
674
675        // We circumvent deadlocking problems by using a *frozen map* (one that can be updated via a shared ref).
676        // See its docs for more details. As an added bonus, we can use indexing notation instead of
677        // clunky methods!
678    }
679}