From 8b3b539fd7a0b4e5fc531dc7b9543d186201cea9 Mon Sep 17 00:00:00 2001
From: Ivan Babrou <github@ivan.computer>
Date: Fri, 20 Sep 2024 21:51:24 -0700
Subject: [PATCH] Implement Atomic<u64> for AtomicU64 for gauges

Between forcing end users to do endless `as i64` for things that are
clearly `u64` and having one error case for rarely used protobuf when
a gauge is set to `u64::MAX`, the latter seems like the right choice.

Signed-off-by: Ivan Babrou <github@ivan.computer>
---
 src/encoding.rs          | 10 ++++++++++
 src/encoding/protobuf.rs | 37 +++++++++++++++++++++++++++++++++++++
 src/metrics/gauge.rs     | 27 +++++++++++++++++++++++++++
 3 files changed, 74 insertions(+)

diff --git a/src/encoding.rs b/src/encoding.rs
index c644f82..e1a3c0b 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -568,6 +568,16 @@ impl EncodeGaugeValue for i64 {
     }
 }
 
+impl EncodeGaugeValue for u64 {
+    fn encode(&self, encoder: &mut GaugeValueEncoder) -> Result<(), std::fmt::Error> {
+        if *self == u64::MAX {
+            return Err(std::fmt::Error);
+        }
+
+        encoder.encode_i64(*self as i64)
+    }
+}
+
 impl EncodeGaugeValue for f64 {
     fn encode(&self, encoder: &mut GaugeValueEncoder) -> Result<(), std::fmt::Error> {
         encoder.encode_f64(*self)
diff --git a/src/encoding/protobuf.rs b/src/encoding/protobuf.rs
index eef2731..6ace229 100644
--- a/src/encoding/protobuf.rs
+++ b/src/encoding/protobuf.rs
@@ -453,6 +453,7 @@ mod tests {
     use std::borrow::Cow;
     use std::collections::HashSet;
     use std::sync::atomic::AtomicI64;
+    use std::sync::atomic::AtomicU64;
 
     #[test]
     fn encode_counter_int() {
@@ -600,6 +601,42 @@ mod tests {
         }
     }
 
+    #[test]
+    fn encode_gauge_u64_normal() {
+        let mut registry = Registry::default();
+        let gauge = Gauge::<u64, AtomicU64>::default();
+        registry.register("my_gauge", "My gauge", gauge.clone());
+        gauge.set(12345);
+
+        let metric_set = encode(&registry).unwrap();
+        let family = metric_set.metric_families.first().unwrap();
+        assert_eq!("my_gauge", family.name);
+        assert_eq!("My gauge.", family.help);
+
+        assert_eq!(
+            openmetrics_data_model::MetricType::Gauge as i32,
+            extract_metric_type(&metric_set)
+        );
+
+        match extract_metric_point_value(&metric_set) {
+            openmetrics_data_model::metric_point::Value::GaugeValue(value) => {
+                let expected = openmetrics_data_model::gauge_value::Value::IntValue(12345);
+                assert_eq!(Some(expected), value.value);
+            }
+            _ => panic!("wrong value type"),
+        }
+    }
+
+    #[test]
+    fn encode_gauge_u64_max() {
+        let mut registry = Registry::default();
+        let gauge = Gauge::<u64, AtomicU64>::default();
+        registry.register("my_gauge", "My gauge", gauge.clone());
+        gauge.set(u64::MAX);
+
+        assert!(encode(&registry).is_err());
+    }
+
     #[test]
     fn encode_counter_family() {
         let mut registry = Registry::default();
diff --git a/src/metrics/gauge.rs b/src/metrics/gauge.rs
index 5c8bef9..abbc8ce 100644
--- a/src/metrics/gauge.rs
+++ b/src/metrics/gauge.rs
@@ -213,6 +213,33 @@ impl Atomic<i64> for AtomicI64 {
     }
 }
 
+#[cfg(target_has_atomic = "64")]
+impl Atomic<u64> for AtomicU64 {
+    fn inc(&self) -> u64 {
+        self.inc_by(1)
+    }
+
+    fn inc_by(&self, v: u64) -> u64 {
+        self.fetch_add(v, Ordering::Relaxed)
+    }
+
+    fn dec(&self) -> u64 {
+        self.dec_by(1)
+    }
+
+    fn dec_by(&self, v: u64) -> u64 {
+        self.fetch_sub(v, Ordering::Relaxed)
+    }
+
+    fn set(&self, v: u64) -> u64 {
+        self.swap(v, Ordering::Relaxed)
+    }
+
+    fn get(&self) -> u64 {
+        self.load(Ordering::Relaxed)
+    }
+}
+
 #[cfg(target_has_atomic = "64")]
 impl Atomic<f64> for AtomicU64 {
     fn inc(&self) -> f64 {