From dc2cba3ae7762949d5e4ca06e6cb0a86f763bf6d Mon Sep 17 00:00:00 2001
From: Palash Nigam <npalash25@gmail.com>
Date: Fri, 14 Oct 2022 22:57:23 +0530
Subject: [PATCH] Add encoding for Summary metric

Signed-off-by: Palash Nigam <npalash25@gmail.com>
---
 src/encoding/text.rs   | 76 ++++++++++++++++++++++++++++++++++++++++++
 src/metrics/summary.rs |  4 +--
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/src/encoding/text.rs b/src/encoding/text.rs
index e033e76a..7646bcb0 100644
--- a/src/encoding/text.rs
+++ b/src/encoding/text.rs
@@ -29,6 +29,7 @@ use crate::metrics::exemplar::{CounterWithExemplar, Exemplar, HistogramWithExemp
 use crate::metrics::family::{Family, MetricConstructor};
 use crate::metrics::gauge::{self, Gauge};
 use crate::metrics::histogram::Histogram;
+use crate::metrics::summary::Summary;
 use crate::metrics::info::Info;
 use crate::metrics::{MetricType, TypedMetric};
 use crate::registry::{Registry, Unit};
@@ -332,6 +333,23 @@ impl<'a> BucketEncoder<'a> {
         })
     }
 
+    /// Encode a quantile. Used for the [`Summary`] metric type.
+    pub fn encode_quantile(&mut self, quantile: f64) -> Result<ValueEncoder, std::io::Error> {
+        if self.opened_curly_brackets {
+            self.writer.write_all(b",")?;
+        } else {
+            self.writer.write_all(b"{")?;
+        }
+
+        self.writer.write_all(b"quantile=\"")?;
+        quantile.encode(self.writer)?;
+        self.writer.write_all(b"\"}")?;
+
+        Ok(ValueEncoder {
+            writer: self.writer,
+        })
+    }
+
     /// Signal that the metric type has no bucket.
     pub fn no_bucket(&mut self) -> Result<ValueEncoder, std::io::Error> {
         if self.opened_curly_brackets {
@@ -581,6 +599,40 @@ fn encode_histogram_with_maybe_exemplars<S: Encode>(
     Ok(())
 }
 
+/////////////////////////////////////////////////////////////////////////////////
+// Summary
+
+impl EncodeMetric for Summary {
+    fn encode(&self, mut encoder: Encoder) -> Result<(), std::io::Error> {
+        let (sum, count, quantiles) = self.get();
+
+        encoder
+        .encode_suffix("sum")?
+        .no_bucket()?
+        .encode_value(sum)?
+        .no_exemplar()?;
+        encoder
+        .encode_suffix("count")?
+        .no_bucket()?
+        .encode_value(count)?
+        .no_exemplar()?;
+
+        for (_, (quantile, result)) in quantiles.iter().enumerate() {
+            let mut bucket_encoder = encoder.no_suffix()?;
+            let mut value_encoder = bucket_encoder.encode_quantile(*quantile)?;
+            let mut exemplar_encoder = value_encoder.encode_value(*result)?;
+            exemplar_encoder.no_exemplar()?
+        }
+
+        Result::Ok(())
+    }
+
+    fn metric_type(&self) -> MetricType {
+        Self::TYPE
+    }
+}
+
+
 /////////////////////////////////////////////////////////////////////////////////
 // Info
 
@@ -821,6 +873,30 @@ mod tests {
         parse_with_python_client(String::from_utf8(encoded).unwrap());
     }
 
+    #[test]
+    fn encode_summary() {
+        let mut registry = Registry::default();
+        let summary = Summary::new(3, 10, vec![0.5, 0.9, 0.99], 0.0);
+        registry.register("my_summary", "My summary", summary.clone());
+        summary.observe(0.10);
+        summary.observe(0.20);
+        summary.observe(0.30);
+
+        let mut encoded = Vec::new();
+
+        encode(&mut encoded, &registry).unwrap();
+
+        let expected = "# HELP my_summary My summary.\n".to_owned()
+            + "# TYPE my_summary summary\n"
+            + "my_summary_sum 0.6000000000000001\n"
+            + "my_summary_count 3\n"
+            + "my_summary{quantile=\"0.5\"} 0.2\n"
+            + "my_summary{quantile=\"0.9\"} 0.3\n"
+            + "my_summary{quantile=\"0.99\"} 0.3\n"
+            + "# EOF\n";
+        assert_eq!(expected, String::from_utf8(encoded.clone()).unwrap());
+    }
+
     fn parse_with_python_client(input: String) {
         pyo3::prepare_freethreaded_python();
 
diff --git a/src/metrics/summary.rs b/src/metrics/summary.rs
index 2507dc83..3488406c 100644
--- a/src/metrics/summary.rs
+++ b/src/metrics/summary.rs
@@ -60,7 +60,7 @@ impl Summary {
         }
     }
 
-    pub fn observe(&mut self, v: f64) {
+    pub fn observe(&self, v: f64) {
         let mut inner = self.inner.lock().unwrap();
         inner.sum += v;
         inner.count += 1;
@@ -101,7 +101,7 @@ mod tests {
 
     #[test]
     fn basic() {
-        let mut summary = Summary::new(5, 10, vec![0.5, 0.9, 0.99], 0.01);
+        let summary = Summary::new(5, 10, vec![0.5, 0.9, 0.99], 0.01);
         summary.observe(5.0);
         summary.observe(15.0);
         summary.observe(25.0);