From bdd24286ba7d753675bcd646903c116c57b447fd Mon Sep 17 00:00:00 2001
From: Mihir Nanavati <code@mihirnanavati.com>
Date: Mon, 7 Aug 2023 18:46:13 -0400
Subject: [PATCH] output: use library's compact histogram format

The sparse histogram format used for serialization is now
available in the library so remove the local implementation
and use the library's version of the structure.
---
 Cargo.lock        |  6 ++++--
 Cargo.toml        |  1 +
 src/output/mod.rs | 51 ++++++-----------------------------------------
 3 files changed, 11 insertions(+), 47 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 08cb8043..8d5cc651 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -798,10 +798,11 @@ checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286"
 
 [[package]]
 name = "histogram"
-version = "0.7.3"
+version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d0978bb4ae7b21dded5037bc688271ec01443b6eb2c7713aad75fce2cf670923"
+checksum = "e673d137229619d5c2c8903b6ed5852b43636c0017ff2e66b1aafb8ccf04b80b"
 dependencies = [
+ "serde",
  "thiserror",
 ]
 
@@ -1893,6 +1894,7 @@ dependencies = [
  "clocksource",
  "foreign-types-shared",
  "futures",
+ "histogram",
  "http-body-util",
  "humantime",
  "hyper 1.0.0-rc.3",
diff --git a/Cargo.toml b/Cargo.toml
index 3e1464ea..acc58405 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,6 +21,7 @@ foreign-types-shared = "0.3.1"
 futures = "0.3.28"
 http-body-util = "0.1.0-rc.2"
 hyper = { version = "1.0.0-rc.3", features = ["http1", "http2", "client"]}
+histogram = "0.7.4"
 humantime = "2.1.0"
 momento = "0.31.0"
 metriken = "0.2.1"
diff --git a/src/output/mod.rs b/src/output/mod.rs
index 9f9ba63b..871c888a 100644
--- a/src/output/mod.rs
+++ b/src/output/mod.rs
@@ -4,6 +4,8 @@ use ratelimit::Ratelimiter;
 use serde::Serialize;
 use std::io::{BufWriter, Write};
 
+use histogram::CompactHistogram;
+
 #[macro_export]
 macro_rules! output {
     () => {
@@ -255,28 +257,12 @@ struct Responses {
     miss: u64,
 }
 
-#[derive(Serialize)]
-/// Sparse histogram in column-order, i.e., storing a single vector
-/// per-field of non-zero buckets. Assuming index[0] = n,
-/// (index[0], count[0]) corresponds to the nth bucket.
-struct RequestLatencies {
-    /// parameters representing the resolution and the range of
-    /// the histogram tracking request latencies
-    m: u32,
-    r: u32,
-    n: u32,
-    /// indices for the non-zero buckets in the histogram
-    index: Vec<usize>,
-    /// histogram bucket counts corresponding to the indices
-    count: Vec<u32>,
-}
-
 #[derive(Serialize)]
 struct Client {
     connections: Connections,
     requests: Requests,
     responses: Responses,
-    request_latency: RequestLatencies,
+    request_latency: CompactHistogram,
 }
 
 #[derive(Serialize)]
@@ -308,7 +294,7 @@ struct JsonSnapshot {
 }
 
 // gets the non-zero buckets for the most recent window in the heatmap
-fn heatmap_to_buckets(heatmap: &Heatmap) -> RequestLatencies {
+fn heatmap_to_buckets(heatmap: &Heatmap) -> CompactHistogram {
     // XXX: The heatmap corrects for wraparound and fixes indices once
     // the heatmap is full so this returns the histogram for the last
     // completed epoch, assuming a heatmap with a total of 60 valid
@@ -316,35 +302,10 @@ fn heatmap_to_buckets(heatmap: &Heatmap) -> RequestLatencies {
     // has been populated, so for the first minute, no histograms
     // are returned (the histogram at offset 59 is still invalid).
     if let Some(Some(histogram)) = heatmap.iter().map(|mut i| i.nth(59)) {
-        let p = histogram.parameters();
-        let mut index = Vec::new();
-        let mut count = Vec::new();
-
-        for (i, bucket) in histogram
-            .into_iter()
-            .enumerate()
-            .filter(|(_i, bucket)| bucket.count() != 0)
-        {
-            index.push(i);
-            count.push(bucket.count());
-        }
-
-        RequestLatencies {
-            m: p.0,
-            r: p.1,
-            n: p.2,
-            index,
-            count,
-        }
+        CompactHistogram::from(histogram)
     } else {
         trace!("no histogram");
-        RequestLatencies {
-            m: 0,
-            r: 0,
-            n: 0,
-            index: vec![],
-            count: vec![],
-        }
+        CompactHistogram::default()
     }
 }