From a9daff1f42e5c445eb511419e8e6db52fea603fd Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Tue, 20 Aug 2024 21:42:03 -0700
Subject: [PATCH] metrics: kcli: display p90 instead of avg

This commit adjusts the formatting of the json rendition of
histogram buckets, and then teaches kcli how to compute
quantiles over the buckets.

We use that to collect and show the p90 rather than the avg
latency in kcli top.

I want to add more histogram specific presentation options
to kcli in a follow up commit; this is just capturing
the quantile logic.
---
 crates/kcli/src/queue_summary.rs              | 167 ++++++++++++++++--
 crates/kcli/src/top.rs                        |  12 +-
 .../kumo-server-common/src/http_server/mod.rs |   7 +-
 3 files changed, 164 insertions(+), 22 deletions(-)

diff --git a/crates/kcli/src/queue_summary.rs b/crates/kcli/src/queue_summary.rs
index f3ae3347..e80ca2f4 100644
--- a/crates/kcli/src/queue_summary.rs
+++ b/crates/kcli/src/queue_summary.rs
@@ -142,11 +142,146 @@ pub struct LatencyIndividual {
 #[derive(Deserialize, Serialize)]
 pub struct LatencyEntry {
     pub avg: f64,
-    pub bucket: HashMap<String, u64>,
+    pub bucket: Vec<(f64, f64)>,
     pub count: u64,
     pub sum: f64,
 }
 
+impl LatencyEntry {
+    /// Given a quantile (eg: p90 would be q=0.9), returns the approximate
+    /// observed latency value that that percentage of samples would
+    /// have recorded.  The value is approximated through linear interpolation
+    /// across the range found in the last matching bucket.
+    ///
+    /// This logic is derived from the histogramQuantile function in prometheus
+    /// <https://github.com/prometheus/prometheus/blob/1435c8ae4aa1041592778018ba62fc3058a9ad3d/promql/quantile.go#L177>
+    pub fn quantile(&self, q: f64) -> f64 {
+        if q < 0.0 {
+            return f64::NEG_INFINITY;
+        }
+        if q > 1.0 {
+            return f64::INFINITY;
+        }
+
+        if self.count == 0 || q.is_nan() {
+            return f64::NAN;
+        }
+
+        #[derive(Debug, Clone, Copy, Default)]
+        struct Bucket {
+            lower_bound: f64,
+            upper_bound: f64,
+            count: f64,
+        }
+
+        let mut buckets = vec![];
+
+        let mut lower_bound = 0.0;
+        for &(upper_bound, cumulative_count) in &self.bucket {
+            buckets.push(Bucket {
+                lower_bound,
+                upper_bound,
+                count: cumulative_count,
+            });
+            lower_bound = upper_bound;
+        }
+
+        // Fixup cumulative counts to be the simple per-bucket counts.
+        // We do this by walking backwards and subtracting the earlier
+        // count from the current count
+        {
+            let mut iter = buckets.iter_mut().rev().peekable();
+            while let Some(b) = iter.next() {
+                if let Some(prev) = iter.peek() {
+                    b.count -= prev.count;
+                }
+            }
+        }
+
+        fn bucket_iter<'a>(
+            buckets: &'a [Bucket],
+            forward: bool,
+        ) -> Box<dyn Iterator<Item = &'a Bucket> + 'a> {
+            if forward {
+                Box::new(buckets.iter())
+            } else {
+                Box::new(buckets.iter().rev())
+            }
+        }
+
+        let forwards = self.sum.is_nan() || q < 0.5;
+
+        let (mut rank, iter) = if forwards {
+            (q * self.count as f64, bucket_iter(&buckets, true))
+        } else {
+            ((1.0 - q) * self.count as f64, bucket_iter(&buckets, false))
+        };
+
+        let mut count = 0.0;
+        let mut bucket = None;
+        for b in iter {
+            bucket.replace(b);
+            if b.count == 0.0 {
+                continue;
+            }
+            count += b.count;
+            if count >= rank {
+                break;
+            }
+        }
+
+        let Some(bucket) = bucket else {
+            return f64::NEG_INFINITY;
+        };
+
+        count = count.min(self.count as f64);
+        if count < rank {
+            return bucket.upper_bound;
+        }
+        if forwards {
+            rank -= count - bucket.count;
+        } else {
+            rank = count - rank;
+        }
+
+        bucket.lower_bound + (bucket.upper_bound - bucket.lower_bound) * (rank / bucket.count)
+    }
+}
+
+#[cfg(test)]
+mod histogram_test {
+    use super::*;
+
+    #[test]
+    fn basics() {
+        let entry = LatencyEntry {
+            avg: 0.03883320766938923,
+            bucket: vec![
+                (0.005, 148571.),
+                (0.01, 149185.),
+                (0.025, 201435.),
+                (0.05, 505005.),
+                (0.1, 611944.),
+                (0.25, 643205.),
+                (0.5, 643876.),
+                (1., 645492.),
+                (2.5, 646039.),
+                (5., 646039.),
+                (10., 646039.),
+            ],
+            count: 646039,
+            sum: 25087.76664952455,
+        };
+
+        assert_eq!(entry.quantile(1.0), 2.5);
+        assert_eq!(entry.quantile(0.99), 0.23259945299254658);
+        assert_eq!(entry.quantile(0.95), 0.10860361152874175);
+        assert_eq!(entry.quantile(0.9), 0.08573537250208063);
+        assert_eq!(entry.quantile(0.75), 0.04831375382942979);
+        assert_eq!(entry.quantile(0.5), 0.03501288829594493);
+    }
+}
+
 #[derive(Deserialize, Serialize)]
 pub struct Metrics {
     pub connection_count: Option<CounterGroup>,
@@ -182,9 +317,21 @@ pub struct Metrics {
 pub struct LatencyMetrics {
     pub name: String,
     pub avg: f64,
+    pub p90: f64,
     pub count: u64,
 }
 
+impl LatencyMetrics {
+    fn new(name: &str, entry: &LatencyEntry) -> Self {
+        Self {
+            name: name.to_string(),
+            avg: entry.avg,
+            p90: entry.quantile(0.9),
+            count: entry.count,
+        }
+    }
+}
+
 pub struct ThreadPoolMetrics {
     pub name: String,
     pub size: usize,
@@ -384,11 +531,7 @@ pub async fn obtain_metrics(endpoint: &Url, by_volume: bool) -> anyhow::Result<P
         .and_then(|group| group.value.as_ref())
     {
         for (event, entry) in &map.event {
-            latency.push(LatencyMetrics {
-                name: event.to_string(),
-                avg: entry.avg,
-                count: entry.count,
-            });
+            latency.push(LatencyMetrics::new(event, entry));
         }
     }
     if let Some(map) = result
@@ -397,11 +540,7 @@ pub async fn obtain_metrics(endpoint: &Url, by_volume: bool) -> anyhow::Result<P
         .and_then(|group| group.value.as_ref())
     {
         for (event, entry) in &map.logger {
-            latency.push(LatencyMetrics {
-                name: event.to_string(),
-                avg: entry.avg,
-                count: entry.count,
-            });
+            latency.push(LatencyMetrics::new(event, entry));
         }
     }
     for (name, entry) in [
@@ -436,11 +575,7 @@ pub async fn obtain_metrics(endpoint: &Url, by_volume: bool) -> anyhow::Result<P
         ),
     ] {
         if let Some(entry) = entry {
-            latency.push(LatencyMetrics {
-                name: name.to_string(),
-                avg: entry.value.avg,
-                count: entry.value.count,
-            });
+            latency.push(LatencyMetrics::new(name, &entry.value));
         }
     }
 
diff --git a/crates/kcli/src/top.rs b/crates/kcli/src/top.rs
index fce37d08..fab3b37c 100644
--- a/crates/kcli/src/top.rs
+++ b/crates/kcli/src/top.rs
@@ -117,6 +117,7 @@ struct State {
 
     thread_pools: BTreeMap<String, Vec<u64>>,
     latency_avg: BTreeMap<String, Vec<u64>>,
+    latency_p90: BTreeMap<String, Vec<u64>>,
     latency_count: BTreeMap<String, Vec<u64>>,
 }
 
@@ -229,6 +230,12 @@ impl State {
                         .entry(event.name.to_string())
                         .or_insert_with(Vec::new);
                     push_value(avg_entry, (event.avg * 1_000_000.0).ceil() as u64);
+
+                    let p90_entry = self
+                        .latency_p90
+                        .entry(event.name.to_string())
+                        .or_insert_with(Vec::new);
+                    push_value(p90_entry, (event.p90 * 1_000_000.0).ceil() as u64);
                 }
 
                 for pool in &metrics.thread_pools {
@@ -310,6 +317,9 @@ impl State {
                 for target in self.latency_avg.values_mut() {
                     push_value(target, 0);
                 }
+                for target in self.latency_p90.values_mut() {
+                    push_value(target, 0);
+                }
                 for target in self.latency_count.values_mut() {
                     push_value(target, 0);
                 }
@@ -454,7 +464,7 @@ impl State {
                 1,
             ));
         }
-        for (name, data) in self.latency_avg.iter() {
+        for (name, data) in self.latency_p90.iter() {
             if name == "init" || name == "pre_init" {
                 continue;
             }
diff --git a/crates/kumo-server-common/src/http_server/mod.rs b/crates/kumo-server-common/src/http_server/mod.rs
index fbe895f9..6619a84b 100644
--- a/crates/kumo-server-common/src/http_server/mod.rs
+++ b/crates/kumo-server-common/src/http_server/mod.rs
@@ -325,12 +325,9 @@ async fn report_metrics_json(_: TrustedIpRequired) -> Result<Json<serde_json::Va
                     let sum = hist.get_sample_sum();
                     let avg = if count != 0 { sum / count as f64 } else { 0. };
 
-                    let mut bucket = Map::new();
+                    let mut bucket = vec![];
                     for b in hist.get_bucket() {
-                        bucket.insert(
-                            b.get_upper_bound().to_string().into(),
-                            b.get_cumulative_count().into(),
-                        );
+                        bucket.push(vec![b.get_upper_bound(), b.get_cumulative_count() as f64]);
                     }
 
                     let hist_value = json!({