From 120cfb0e97e3f7853a4ac8cdd1f8e4418e533f65 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Sun, 23 Oct 2022 22:47:36 +1100
Subject: [PATCH] Alternative implementation of `sum` encoding

---
 src/encoding.rs          |  8 ++++++++
 src/encoding/proto.rs    | 23 +++++++++++++++++++++++
 src/encoding/text.rs     | 10 ++++++++++
 src/metrics.rs           |  2 +-
 src/metrics/histogram.rs |  2 +-
 5 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/encoding.rs b/src/encoding.rs
index cc8006ee..73d775b1 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -6,6 +6,7 @@ pub use prometheus_client_derive_encode::*;
 use crate::metrics::exemplar::Exemplar;
 use crate::metrics::MetricType;
 use std::borrow::Cow;
+use std::fmt;
 use std::fmt::Write;
 use std::ops::Deref;
 
@@ -76,6 +77,13 @@ macro_rules! for_both {
 }
 
 impl<'a> MetricEncoder<'a> {
+    pub fn encode_sum(&mut self, value: f64) -> fmt::Result {
+        match self {
+            MetricEncoder::Text(inner) => inner.encode_sum(value),
+            MetricEncoder::Protobuf(inner) => inner.encode_sum(value),
+        }
+    }
+
     pub fn encode_suffix(
         &mut self,
         suffix: &'static str,
diff --git a/src/encoding/proto.rs b/src/encoding/proto.rs
index 4f358e72..9e817ce4 100644
--- a/src/encoding/proto.rs
+++ b/src/encoding/proto.rs
@@ -107,6 +107,29 @@ impl<'a> MetricEncoder<'a> {
         self.no_suffix()
     }
 
+    pub fn encode_sum(&mut self, value: f64) -> Result<(), std::fmt::Error> {
+        if self.metric_type != MetricType::Histogram {
+            panic!("Bad usage")
+        }
+
+        self.family.push(openmetrics_data_model::Metric {
+            labels: self.labels.clone(),
+            metric_points: vec![openmetrics_data_model::MetricPoint {
+                value: Some(openmetrics_data_model::metric_point::Value::HistogramValue(
+                    openmetrics_data_model::HistogramValue {
+                        sum: Some(openmetrics_data_model::histogram_value::Sum::DoubleValue(
+                            value,
+                        )),
+                        ..openmetrics_data_model::HistogramValue::default()
+                    },
+                )),
+                ..Default::default()
+            }],
+        });
+
+        Ok(())
+    }
+
     // TODO: This should consume it, no?
     pub fn no_suffix(&mut self) -> Result<HistogramEncoder, std::fmt::Error> {
         let metric_point = match self.metric_type {
diff --git a/src/encoding/text.rs b/src/encoding/text.rs
index a0a6670c..fd8e5579 100644
--- a/src/encoding/text.rs
+++ b/src/encoding/text.rs
@@ -174,6 +174,16 @@ impl<'a> Encoder<'a> {
         self.encode_labels()
     }
 
+    pub fn encode_sum(&mut self, value: f64) -> Result<(), std::fmt::Error> {
+        self.write_name_and_unit()?;
+
+        self.writer.write_str("_sum")?;
+
+        self.encode_labels()?.encode_sum_f64(value)?;
+
+        Ok(())
+    }
+
     /// Signal that the metric has no suffix.
     pub fn no_suffix(&mut self) -> Result<HistogramEncoder, std::fmt::Error> {
         self.write_name_and_unit()?;
diff --git a/src/metrics.rs b/src/metrics.rs
index 647fa5c7..480086e1 100644
--- a/src/metrics.rs
+++ b/src/metrics.rs
@@ -14,7 +14,7 @@ pub trait TypedMetric {
 }
 
 /// OpenMetrics metric type.
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
 #[allow(missing_docs)]
 pub enum MetricType {
     Counter,
diff --git a/src/metrics/histogram.rs b/src/metrics/histogram.rs
index 01f2228e..ba0a879c 100644
--- a/src/metrics/histogram.rs
+++ b/src/metrics/histogram.rs
@@ -164,7 +164,7 @@ fn encode_histogram_with_maybe_exemplars<S: EncodeLabelSet>(
     exemplars: Option<&HashMap<usize, Exemplar<S, f64>>>,
     mut encoder: MetricEncoder,
 ) -> Result<(), std::fmt::Error> {
-    encoder.encode_suffix("sum")?.encode_sum_f64(sum)?;
+    encoder.encode_sum(sum)?;
     encoder.encode_suffix("count")?.encode_count(count)?;
 
     let mut cummulative = 0;