From 484e8aeb4ac04a456cfd1c0d932ed1e752ed6f99 Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Sun, 1 Sep 2024 05:54:28 -0700 Subject: [PATCH] kcli: adopt streaming metrics parser for top This makes performance super fast in the prescence of millions of queues. --- crates/kcli/src/queue_summary.rs | 522 +-------------------------- crates/kcli/src/top.rs | 275 +++++++------- crates/kumo-prometheus/src/parser.rs | 149 ++++++++ 3 files changed, 288 insertions(+), 658 deletions(-) diff --git a/crates/kcli/src/queue_summary.rs b/crates/kcli/src/queue_summary.rs index c901d3b6..7f826c24 100644 --- a/crates/kcli/src/queue_summary.rs +++ b/crates/kcli/src/queue_summary.rs @@ -5,11 +5,9 @@ use kumo_api_types::{BounceV1ListEntry, SuspendReadyQueueV1ListEntry, SuspendV1L use kumo_prometheus::parser::Metric; use lexicmp::natural_lexical_cmp; use message::message::QueueNameComponents; -use ordermap::OrderMap; use reqwest::Url; -use serde::{Deserialize, Serialize}; use std::cmp::Ordering; -use std::collections::{BTreeMap, HashMap}; +use std::collections::HashMap; use tabout::{Alignment, Column}; /// Prints a summary of the state of the queues, for a human to read. @@ -62,130 +60,6 @@ pub struct QueueSummaryCommand { domain: Option, } -#[derive(Deserialize, Serialize)] -pub struct IndividualCounter { - pub help: String, - #[serde(rename = "type")] - pub type_: String, - pub value: f64, -} - -#[derive(Deserialize, Serialize)] -pub struct CounterGroup { - pub help: String, - #[serde(rename = "type")] - pub type_: String, - pub value: ServiceMap, -} - -#[derive(Deserialize, Serialize)] -pub struct ServiceMap { - pub service: HashMap, -} - -#[derive(Deserialize, Serialize)] -pub struct QueueCounterGroup { - pub help: String, - #[serde(rename = "type")] - pub type_: String, - pub value: QueueMap, -} - -#[derive(Deserialize, Serialize)] -pub struct QueueMap { - pub queue: HashMap, -} - -#[derive(Deserialize, Serialize)] -pub struct ThreadPoolGroup { - pub help: String, - #[serde(rename = "type")] - pub type_: String, - pub value: ThreadPoolMap, -} - -#[derive(Deserialize, Serialize)] -pub struct ThreadPoolMap { - pub pool: HashMap, -} - -#[cfg(test)] -mod histogram_test { - use super::*; - - #[test] - fn basics() { - let entry = HistogramMetric { - 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, - pub ready_count: Option, - pub scheduled_count: Option, - pub total_connection_count: Option, - pub total_messages_delivered: Option, - pub total_messages_transfail: Option, - pub total_messages_fail: Option, - pub total_messages_received: Option, - pub message_count: Option, - pub message_data_resident_count: Option, - pub message_meta_resident_count: Option, - pub memory_usage: Option, - pub memory_limit: Option, - pub thread_pool_size: Option, - pub thread_pool_parked: Option, -} - -pub struct LatencyMetrics { - pub name: MetricName, - pub avg: f64, - pub p90: f64, - pub count: u64, -} - -impl LatencyMetrics { - fn new(name: &MetricName, entry: &HistogramMetric) -> Self { - Self { - name: name.clone(), - avg: entry.avg, - p90: entry.quantile(0.9), - count: entry.count, - } - } -} - -pub struct ThreadPoolMetrics { - pub name: String, - pub size: usize, - pub parked: usize, -} - #[derive(Default, Debug)] pub struct ReadyQueueMetrics { pub name: String, @@ -266,354 +140,6 @@ impl ScheduledQueueMetrics { } } -#[derive(Deserialize, Debug)] -#[serde(untagged)] -enum NumberEntry { - Single(f64), - Map(HashMap>), - List(Vec), -} - -#[derive(Deserialize, Debug)] -struct NumberListEntry { - #[serde(rename = "@")] - value: f64, - #[serde(flatten)] - labels: BTreeMap, -} - -impl NumberEntry { - fn normalize(self, label: String, map: &mut OrderMap) { - match self { - Self::Single(metric) => { - map.insert(MetricName::Label(label), metric); - } - Self::Map(map1) => { - for (label_name, map2) in map1 { - for (k, metric) in map2 { - let mut labels = BTreeMap::new(); - labels.insert(label_name.to_string(), k); - map.insert( - MetricName::Structured { - name: label.to_string(), - labels, - }, - metric, - ); - } - } - } - Self::List(entries) => { - for entry in entries { - map.insert( - MetricName::Structured { - name: label.to_string(), - labels: entry.labels, - }, - entry.value, - ); - } - } - } - } -} - -#[derive(Debug, Hash, PartialEq, Eq, Clone)] -pub enum MetricName { - Label(String), - Structured { - name: String, - labels: BTreeMap, - }, -} - -impl MetricName { - pub fn label(&self) -> String { - match self { - Self::Label(n) => n.to_string(), - Self::Structured { labels, .. } => { - if labels.len() == 1 { - labels.values().next().unwrap().to_string() - } else { - let mut result = String::new(); - for k in labels.values() { - if !result.is_empty() { - result.push(','); - } - result.push_str(k); - } - result - } - } - } - } -} - -impl PartialOrd for MetricName { - fn partial_cmp(&self, other: &MetricName) -> Option { - Some(self.cmp(other)) - } -} -impl Ord for MetricName { - fn cmp(&self, other: &MetricName) -> Ordering { - match (self, other) { - (Self::Label(name_a), Self::Label(name_b)) => natural_lexical_cmp(name_a, name_b), - ( - Self::Structured { - name: name_a, - labels: _labels_a, - }, - Self::Label(name_b), - ) => natural_lexical_cmp(name_a, name_b), - ( - Self::Label(name_a), - Self::Structured { - name: name_b, - labels: _labels_b, - }, - ) => natural_lexical_cmp(name_a, name_b), - ( - Self::Structured { - name: name_a, - labels: labels_a, - }, - Self::Structured { - name: name_b, - labels: labels_b, - }, - ) => match natural_lexical_cmp(name_a, name_b) { - Ordering::Equal => { - for key_a in labels_a.keys() { - match labels_b.get(key_a) { - Some(key_b) => match natural_lexical_cmp(key_a, key_b) { - Ordering::Equal => continue, - ordering => return ordering, - }, - None => return Ordering::Less, - } - } - Ordering::Greater - } - result => result, - }, - } - } -} - -#[derive(Deserialize, Debug)] -#[serde(untagged)] -enum HistogramEntry { - Single(HistogramMetric), - Map(HashMap>), - List(Vec), -} - -#[derive(Deserialize, Debug)] -struct HistogramListEntry { - #[serde(rename = "@")] - value: HistogramMetric, - #[serde(flatten)] - labels: BTreeMap, -} - -impl HistogramEntry { - fn normalize(self, label: String, map: &mut OrderMap) { - match self { - Self::Single(metric) => { - map.insert(MetricName::Label(label), metric); - } - Self::Map(map1) => { - for (label_name, map2) in map1 { - for (k, metric) in map2 { - let mut labels = BTreeMap::new(); - labels.insert(label_name.to_string(), k); - map.insert( - MetricName::Structured { - name: label.to_string(), - labels, - }, - metric, - ); - } - } - } - Self::List(entries) => { - for entry in entries { - map.insert( - MetricName::Structured { - name: label.to_string(), - labels: entry.labels, - }, - entry.value, - ); - } - } - } - } -} - -#[derive(Deserialize, Debug)] -pub struct HistogramMetric { - pub avg: f64, - pub bucket: Vec<(f64, f64)>, - pub count: u64, - pub sum: f64, -} - -impl HistogramMetric { - /// 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 - /// - 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 + '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) - } -} - -#[derive(Deserialize, Debug)] -#[serde(tag = "type", rename_all = "lowercase")] -enum MetricEntry { - Gauge { - #[allow(unused)] - help: String, - value: NumberEntry, - }, - Counter { - #[allow(unused)] - help: String, - value: NumberEntry, - }, - Histogram { - #[allow(unused)] - help: String, - value: HistogramEntry, - }, -} - -#[derive(Default, Debug)] -pub struct DynamicMetrics { - pub gauges: OrderMap, - pub counters: OrderMap, - pub histograms: OrderMap, -} - -fn parse_dynamic(metrics: serde_json::Value) -> anyhow::Result { - let result: HashMap = serde_json::from_value(metrics)?; - - let mut metrics = DynamicMetrics::default(); - - for (label, v) in result { - match v { - MetricEntry::Counter { help: _, value } => { - value.normalize(label, &mut metrics.counters); - } - MetricEntry::Gauge { help: _, value } => { - value.normalize(label, &mut metrics.gauges); - } - MetricEntry::Histogram { help: _, value } => { - value.normalize(label, &mut metrics.histograms); - } - } - } - - metrics.counters.sort_keys(); - metrics.gauges.sort_keys(); - metrics.histograms.sort_keys(); - - Ok(metrics) -} - pub async fn get_metrics Option>( endpoint: &Url, mut filter_map: F, @@ -735,52 +261,6 @@ impl QueueMetrics { } } -pub struct ProcessedMetrics { - pub thread_pools: Vec, - pub latency: Vec, - pub raw: Metrics, - #[allow(unused)] - pub dynamic: DynamicMetrics, -} - -pub async fn obtain_metrics(endpoint: &Url) -> anyhow::Result { - let result: serde_json::Value = crate::request_with_json_response( - reqwest::Method::GET, - endpoint.join("/metrics.json")?, - &(), - ) - .await?; - - let dynamic = parse_dynamic(result.clone())?; - let result: Metrics = serde_json::from_value(result)?; - - let thread_pools = match (&result.thread_pool_size, &result.thread_pool_parked) { - (Some(sizes), Some(values)) => sizes - .value - .pool - .iter() - .map(|(name, size)| ThreadPoolMetrics { - name: name.to_string(), - size: *size as usize, - parked: values.value.pool.get(name).copied().unwrap_or(0.) as usize, - }) - .collect(), - _ => vec![], - }; - - let mut latency = vec![]; - for (name, histo) in &dynamic.histograms { - latency.push(LatencyMetrics::new(name, histo)); - } - - Ok(ProcessedMetrics { - raw: result, - thread_pools, - latency, - dynamic, - }) -} - impl QueueSummaryCommand { pub async fn run(&self, endpoint: &Url) -> anyhow::Result<()> { let suspended_domains: Vec = crate::request_with_json_response( diff --git a/crates/kcli/src/top.rs b/crates/kcli/src/top.rs index dbfe4e23..9d7bca73 100644 --- a/crates/kcli/src/top.rs +++ b/crates/kcli/src/top.rs @@ -1,7 +1,9 @@ +use crate::queue_summary::get_metrics; use clap::Parser; use crossterm::event::{Event, KeyCode, KeyEventKind}; use futures::StreamExt; use human_bytes::human_bytes; +use kumo_prometheus::parser::Metric; use num_format::{Locale, ToFormattedString}; use ratatui::prelude::*; use ratatui::symbols::bar::NINE_LEVELS; @@ -88,8 +90,8 @@ impl TopCommand { } } -fn push_value(target: &mut Vec, value: u64) { - target.insert(0, value); +fn push_value(target: &mut Vec, value: f64) { + target.insert(0, value as u64); target.truncate(1024); } @@ -130,134 +132,125 @@ struct DiffState { impl State { async fn update_metrics(&mut self, endpoint: &Url) -> anyhow::Result<()> { - match crate::queue_summary::obtain_metrics(endpoint).await { - Ok(metrics) => { + let mut delivered = 0.; + let mut transfail = 0.; + let mut fail = 0.; + let mut received = 0.; + let mut latency = BTreeMap::new(); + let mut ready = 0.; + + #[derive(Default)] + struct ThreadPoolMetrics { + size: f64, + parked: f64, + } + let mut thread_pools = BTreeMap::::new(); + + match get_metrics::, _>(endpoint, |m| { + let name = m.name().as_str(); + let value = m.value(); + + match name { + "message_count" => push_value(&mut self.message_count, value), + "message_data_resident_count" => push_value(&mut self.message_data_resident, value), + "message_meta_resident_count" => push_value(&mut self.message_meta_resident, value), + "memory_usage" => push_value(&mut self.memory, value), + "total_messages_delivered" => { + if m.label_is("service", "smtp_client") { + delivered = value; + } + } + "total_messages_transfail" => { + if m.label_is("service", "smtp_client") { + transfail = value; + } + } + "total_messages_fail" => { + if m.label_is("service", "smtp_client") { + fail = value; + } + } + "total_messages_received" => { + if m.label_is("service", "esmtp_listener") + || m.label_is("service", "http_listener") + { + received += value; + } + } + "scheduled_count_total" => { + push_value(&mut self.scheduled, value); + } + "ready_count" => { + ready += value; + } + "connection_count" => { + if m.label_is("service", "esmtp_listener") { + push_value(&mut self.listener_conns, value); + } + if m.label_is("service", "smtp_client") { + push_value(&mut self.smtp_conns, value); + } + } + "thread_pool_size" => { + if let Some(pool) = m.labels().get("pool") { + thread_pools.entry(pool.to_string()).or_default().size = value; + } + } + "thread_pool_parked" => { + if let Some(pool) = m.labels().get("pool") { + thread_pools.entry(pool.to_string()).or_default().parked = value; + } + } + _ => match &m { + Metric::Histogram(h) => { + let label = match m.labels().values().next() { + Some(l) => l.as_str().to_string(), + None => h.name.as_str().to_string(), + }; + latency.insert(label.clone(), h.count); + + let avg_entry = self + .latency_avg + .entry(label.clone()) + .or_insert_with(Vec::new); + push_value(avg_entry, (h.sum / h.count * 1_000_000.0).ceil()); + + let p90_entry = self + .latency_p90 + .entry(label.clone()) + .or_insert_with(Vec::new); + push_value(p90_entry, (h.quantile(0.9) * 1_000_000.0).ceil()); + } + _ => {} + }, + }; + + None + }) + .await + { + Ok(_) => { self.error.clear(); + push_value(&mut self.ready, ready); - for (target, source) in &mut [ - (&mut self.message_count, &metrics.raw.message_count), - ( - &mut self.message_data_resident, - &metrics.raw.message_data_resident_count, - ), - ( - &mut self.message_meta_resident, - &metrics.raw.message_meta_resident_count, - ), - (&mut self.memory, &metrics.raw.memory_usage), - ] { - push_value(target, source.as_ref().map_or(0, |m| m.value as u64)); - target.truncate(1024); - } - - let new_state = DiffState { - when: Instant::now(), - delivered: metrics - .raw - .total_messages_delivered - .as_ref() - .and_then(|m| m.value.service.get("smtp_client")) - .copied() - .unwrap_or(0.), - transfail: metrics - .raw - .total_messages_transfail - .as_ref() - .and_then(|m| m.value.service.get("smtp_client")) - .copied() - .unwrap_or(0.), - fail: metrics - .raw - .total_messages_fail - .as_ref() - .and_then(|m| m.value.service.get("smtp_client")) - .copied() - .unwrap_or(0.), - received: metrics - .raw - .total_messages_received - .as_ref() - .map(|m| { - m.value.service.get("esmtp_listener").copied().unwrap_or(0.) - + m.value.service.get("http_listener").copied().unwrap_or(0.) - }) - .unwrap_or(0.), - latency: metrics - .latency - .iter() - .map(|entry| (entry.name.label().to_string(), entry.count as f64)) - .collect(), - }; - let scheduled = metrics - .raw - .scheduled_count - .as_ref() - .map(|m| m.value.queue.values().copied().sum()) - .unwrap_or(0.); - let ready = metrics - .raw - .ready_count - .as_ref() - .map(|m| m.value.service.values().copied().sum()) - .unwrap_or(0.); - let listener_conns = metrics - .raw - .connection_count - .as_ref() - .and_then(|g| g.value.service.get("esmtp_listener").copied()) - .unwrap_or(0.); - let smtp_conns = metrics - .raw - .connection_count - .as_ref() - .and_then(|g| g.value.service.get("smtp_client").copied()) - .unwrap_or(0.); - - for (target, value) in [ - (&mut self.scheduled, scheduled), - (&mut self.ready, ready), - (&mut self.listener_conns, listener_conns), - (&mut self.smtp_conns, smtp_conns), - ] { - push_value(target, value as u64); - } - - for event in &metrics.latency { - let avg_entry = self - .latency_avg - .entry(event.name.label().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.label().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 { + for (name, pool) in &thread_pools { let entry = self .thread_pools - .entry(pool.name.to_string()) + .entry(name.clone()) .or_insert_with(Vec::new); - let utilization_percent = (100 * (pool.size - pool.parked)) / pool.size; - push_value(entry, utilization_percent as u64); + let utilization_percent = (100. * (pool.size - pool.parked)) / pool.size; + push_value(entry, utilization_percent); } + let mut dead_pools = vec![]; for (key, entry) in self.thread_pools.iter_mut() { - if metrics - .thread_pools - .iter() - .find(|entry| entry.name == *key) - .is_none() - { + if !thread_pools.contains_key(key) { // The pool has gone away. // This can happen for eg: the spoolin pool once // it has completed its work. // We'll treat this as clocking a 0 through. // Once all the data is zero, we'll remove it - push_value(entry, 0); + push_value(entry, 0.0); if entry.iter().sum::() == 0 { dead_pools.push(key.to_string()); } @@ -268,6 +261,14 @@ impl State { self.thread_pools.remove(&name); } + let new_state = DiffState { + when: Instant::now(), + delivered, + transfail, + fail, + received, + latency, + }; if let Some(prior) = self.diff_state.take() { let elapsed = prior.when.elapsed().as_secs_f64(); @@ -278,10 +279,10 @@ impl State { let received = (new_state.received - prior.received) / elapsed; // and add to historical data - push_value(&mut self.delivered, delivered as u64); - push_value(&mut self.transfail, transfail as u64); - push_value(&mut self.fail, fail as u64); - push_value(&mut self.received, received as u64); + push_value(&mut self.delivered, delivered); + push_value(&mut self.transfail, transfail); + push_value(&mut self.fail, fail); + push_value(&mut self.received, received); for (name, new_value) in &new_state.latency { let rate = @@ -290,7 +291,7 @@ impl State { .latency_count .entry(name.to_string()) .or_insert_with(Vec::new); - push_value(entry, rate as u64); + push_value(entry, rate); } } self.diff_state.replace(new_state); @@ -298,28 +299,28 @@ impl State { Err(err) => { self.error = format!("{err:#}"); self.diff_state.take(); - push_value(&mut self.memory, 0); - push_value(&mut self.message_count, 0); - push_value(&mut self.message_data_resident, 0); - push_value(&mut self.scheduled, 0); - push_value(&mut self.ready, 0); - push_value(&mut self.listener_conns, 0); - push_value(&mut self.smtp_conns, 0); - push_value(&mut self.delivered, 0); - push_value(&mut self.transfail, 0); - push_value(&mut self.fail, 0); - push_value(&mut self.received, 0); + push_value(&mut self.memory, 0.0); + push_value(&mut self.message_count, 0.0); + push_value(&mut self.message_data_resident, 0.0); + push_value(&mut self.scheduled, 0.0); + push_value(&mut self.ready, 0.0); + push_value(&mut self.listener_conns, 0.0); + push_value(&mut self.smtp_conns, 0.0); + push_value(&mut self.delivered, 0.0); + push_value(&mut self.transfail, 0.0); + push_value(&mut self.fail, 0.0); + push_value(&mut self.received, 0.0); for target in self.thread_pools.values_mut() { - push_value(target, 0); + push_value(target, 0.0); } for target in self.latency_avg.values_mut() { - push_value(target, 0); + push_value(target, 0.0); } for target in self.latency_p90.values_mut() { - push_value(target, 0); + push_value(target, 0.0); } for target in self.latency_count.values_mut() { - push_value(target, 0); + push_value(target, 0.0); } } } diff --git a/crates/kumo-prometheus/src/parser.rs b/crates/kumo-prometheus/src/parser.rs index ccfaef07..516395ee 100644 --- a/crates/kumo-prometheus/src/parser.rs +++ b/crates/kumo-prometheus/src/parser.rs @@ -315,6 +315,13 @@ impl Metric { } } + pub fn label_is(&self, key: &str, value: &str) -> bool { + self.labels() + .get(key) + .map(|v| v.as_str() == value) + .unwrap_or(false) + } + pub fn labels(&self) -> &Map { match self { Self::Counter(c) => &c.labels, @@ -339,6 +346,17 @@ impl Metric { } key } + + pub fn is_histogram(&self) -> bool { + matches!(self, Self::Histogram(_)) + } + + pub fn as_histogram(&self) -> &HistogramMetric { + match self { + Self::Histogram(h) => h, + _ => panic!("as_histogram called on non-histogram"), + } + } } #[derive(Debug, PartialEq)] @@ -364,6 +382,107 @@ pub struct HistogramMetric { pub bucket: Vec<(f64, f64)>, } +impl HistogramMetric { + /// 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 + /// + 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.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 + '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 test { use super::*; @@ -605,4 +724,34 @@ disk_free_bytes{name="meta spool"} 1540683988992 ] ); } + + #[test] + fn histogram_quantile() { + let entry = HistogramMetric { + name: InternString::new("boop"), + labels: Map::new(), + 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.0, + 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); + } }