1use 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#[doc(hidden)] #[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#[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
76pub 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 pub fn inc_by(&self, v: V) -> V {
105 self.0.inc_by(v)
106 }
107
108 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 pub fn dec_by(&self, v: V) -> V {
125 self.0.dec_by(v)
126 }
127
128 pub fn set(&self, value: V) -> V {
130 self.0.set(value)
131 }
132
133 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#[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#[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 pub fn observe(&self, value: V) {
203 self.inner.observe(value.encode());
204 }
205}
206
207impl Histogram<Duration> {
208 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#[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 pub fn observe(self) -> Duration {
245 let elapsed = self.start.elapsed();
246 self.histogram.observe(elapsed);
247 elapsed
248 }
249}
250
251#[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 pub fn get(&self) -> Option<&S> {
273 self.0.get()
274 }
275
276 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#[derive(Debug)]
308pub struct SetInfoError<S>(S);
309
310impl<S> SetInfoError<S> {
311 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
384pub struct Family<S, M: BuildMetric, L = ()> {
388 inner: Arc<FamilyInner<S, M>>,
389 labels: L,
390}
391
392pub 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 pub fn contains(&self, labels: &S) -> bool {
502 self.inner.map.get(labels).is_some()
503 }
504
505 pub fn get(&self, labels: &S) -> Option<&M> {
508 self.inner.map.get(labels)
509 }
510
511 pub fn get_lazy(&self, labels: S) -> LazyItem<'_, S, M> {
515 LazyItem::new(&self.inner, labels)
516 }
517
518 pub fn to_entries(&self) -> impl ExactSizeIterator<Item = (S, &M)> + '_ {
521 self.inner.to_entries()
522 }
523}
524
525impl<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
539pub 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 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 }
679}