From 12923ca39c6392dcc7d0e81af1acc3788c4235d4 Mon Sep 17 00:00:00 2001
From: Ivan Babrou <github@ivan.computer>
Date: Wed, 16 Oct 2024 03:40:16 -0700
Subject: [PATCH] fix(metrics/gauge): implement Atomic<u64> for AtomicU64
 (#226)

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>
---
 CHANGELOG.md             |  5 +++++
 src/encoding.rs          | 13 +++++++++++++
 src/encoding/protobuf.rs | 38 ++++++++++++++++++++++++++++++++++++++
 src/metrics/gauge.rs     | 27 +++++++++++++++++++++++++++
 4 files changed, 83 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 66eb579..c398e07 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.23.0] - unreleased
 
+
 ### Changed
 
 - `ConstCounter::new` now requires specifying the type of literal arguments, like this: `ConstCounter::new(42u64);`.
@@ -14,6 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Update `prost` dependencies to `v0.12`.
   See [PR 198].
 
+- Implement `Atomic<u64>` for `AtomicU64` for gauges.
+  See [PR 226].
+
+[PR 226]: https://github.com/prometheus/client_rust/pull/198
 [PR 198]: https://github.com/prometheus/client_rust/pull/198
 
 ### Added
diff --git a/src/encoding.rs b/src/encoding.rs
index c644f82..7f77efa 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -568,6 +568,19 @@ impl EncodeGaugeValue for i64 {
     }
 }
 
+impl EncodeGaugeValue for u64 {
+    fn encode(&self, encoder: &mut GaugeValueEncoder) -> Result<(), std::fmt::Error> {
+        // 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.
+        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..095a3db 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,43 @@ 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);
+
+        // This expected to fail as protobuf uses i64 and u64::MAX does not fit into it.
+        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 {