From f255dc50652224931ba51c5bed5c6cf92ca1be36 Mon Sep 17 00:00:00 2001
From: Darksome <dev@darkso.me>
Date: Mon, 1 Jul 2024 17:11:52 +0000
Subject: [PATCH] feat: Optional label

---
 crates/metrics/src/label.rs | 45 ++++++++++++++++++++++++++++++++++++-
 crates/metrics/src/lib.rs   | 24 ++++++++++++++------
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/crates/metrics/src/label.rs b/crates/metrics/src/label.rs
index 9d9b143..5e3d9e8 100644
--- a/crates/metrics/src/label.rs
+++ b/crates/metrics/src/label.rs
@@ -8,7 +8,7 @@ use {
     enum_ordinalize::Ordinalize,
     parking_lot::Mutex,
     smallvec::SmallVec,
-    std::{borrow::Borrow, collections::HashMap, sync::Arc},
+    std::{borrow::Borrow, collections::HashMap, marker::PhantomData, sync::Arc},
 };
 
 pub type DynamicLabels = SmallVec<[Label; 4]>;
@@ -343,6 +343,49 @@ where
     }
 }
 
+/// Makes any other label optional by accepting `Option` instead of the actual
+/// label value during the label resolution.
+pub struct Optional<T>(PhantomData<T>);
+
+impl<T, M> DynamicLabel<M> for Optional<T>
+where
+    T: DynamicLabel<M>,
+{
+    type MetricCollection = (M, WithLabel<T, M>);
+}
+
+impl<T, M> Metric for WithLabel<Optional<T>, M>
+where
+    T: DynamicLabel<M>,
+    M: Metric + 'static,
+    WithLabel<T, M>: Metric,
+{
+    fn register(attrs: &Attrs) -> Self {
+        Self {
+            collection: (
+                M::register(attrs),
+                <WithLabel<T, M> as Metric>::register(attrs),
+            ),
+        }
+    }
+}
+
+impl<T, U, M> ResolveLabels<(Option<U>,)> for WithLabel<Optional<T>, M>
+where
+    T: DynamicLabel<M>,
+    M: Metric,
+    WithLabel<T, M>: ResolveLabels<(U,), Target = M>,
+{
+    type Target = M;
+
+    fn resolve_labels(&self, (label,): (Option<U>,)) -> &M {
+        match label {
+            Some(l) => self.collection.1.resolve_labels((l,)),
+            None => &self.collection.0,
+        }
+    }
+}
+
 // TODO: macro to autogenerate these
 
 impl<L, M, A, B> ResolveLabels<(A, B)> for WithLabel<L, M>
diff --git a/crates/metrics/src/lib.rs b/crates/metrics/src/lib.rs
index 6e6d4c0..5e4a0e6 100644
--- a/crates/metrics/src/lib.rs
+++ b/crates/metrics/src/lib.rs
@@ -25,6 +25,7 @@
 //!     LabeledGauge3,
 //!     LabeledHistogram,
 //!     Lazy,
+//!     OptionalLabel,
 //!     StringLabel,
 //! };
 //!
@@ -63,7 +64,7 @@
 //!         .with_description("My labeled counter")
 //!         .build();
 //!
-//! static GAUGE_B: Lazy<LabeledGauge3<MyU8StringLabel, MyEnumLabel, MyBoolLabel>> =
+//! static GAUGE_B: Lazy<LabeledGauge3<MyU8StringLabel, MyEnumLabel, OptionalLabel<MyBoolLabel>>> =
 //!     metrics::new("gauge_b");
 //!
 //! COUNTER_A.increment(1);
@@ -71,18 +72,23 @@
 //! HISTOGRAM_A.record(1000, (MyEnumLabel::new(MyEnum::A),));
 //! COUNTER_B.increment(2u64, (MyStringLabel::new("test"), MyBoolLabel::new(false)));
 //!
-//! let labels = (
-//!     MyU8StringLabel::new(&42),
-//!     MyEnumLabel::new(MyEnum::B),
-//!     MyBoolLabel::new(true),
-//! );
+//! let labels = (MyU8StringLabel::new(&42), MyEnumLabel::new(MyEnum::B), None);
 //! GAUGE_B.decrement(2, labels);
 //! ```
 
 pub use {
     backend,
     enum_ordinalize,
-    label::{label_name, BoolLabel, Enum, EnumLabel, LabelName, StringLabel, WithLabel},
+    label::{
+        label_name,
+        BoolLabel,
+        Enum,
+        EnumLabel,
+        LabelName,
+        Optional as OptionalLabel,
+        StringLabel,
+        WithLabel,
+    },
     lazy::Lazy,
 };
 use {
@@ -310,3 +316,7 @@ pub type LabeledFutureMetrics2<A, B> = Labeled2<FutureMetrics, A, B>;
 pub type LabeledFutureMetrics3<A, B, C> = Labeled3<FutureMetrics, A, B, C>;
 #[cfg(feature = "future")]
 pub type LabeledFutureMetrics4<A, B, C, D> = Labeled4<FutureMetrics, A, B, C, D>;
+
+pub type OptionalEnumLabel<const NAME: LabelName, T> = OptionalLabel<EnumLabel<NAME, T>>;
+pub type OptionalBoolLabel<const NAME: LabelName> = OptionalLabel<BoolLabel<NAME>>;
+pub type OptionalStringLabel<const NAME: LabelName, T> = OptionalLabel<StringLabel<NAME, T>>;