Skip to main content

systemprompt_analytics/snapshots/
histogram.rs

1//! Version one geometric microsecond bounds merge bucket counts without
2//! averaging percentiles.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeMap;
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, schemars::JsonSchema)]
11/// Versioned geometric latency buckets whose counts can be merged exactly.
12pub struct LatencyHistogram {
13    pub version: u32,
14    pub buckets: BTreeMap<u32, i64>,
15}
16impl Default for LatencyHistogram {
17    fn default() -> Self {
18        Self {
19            version: 1,
20            buckets: BTreeMap::new(),
21        }
22    }
23}
24impl LatencyHistogram {
25    pub fn record(&mut self, micros: u64) -> crate::Result<()> {
26        let bucket = if micros == 0 {
27            0
28        } else {
29            64 - micros.leading_zeros()
30        };
31        let count = self.buckets.entry(bucket).or_default();
32        *count = count
33            .checked_add(1)
34            .ok_or_else(|| super::invalid("Histogram overflow"))?;
35        Ok(())
36    }
37    pub fn merge(&mut self, other: &Self) -> crate::Result<()> {
38        if self.version != 1 || other.version != 1 {
39            return Err(super::invalid("Unsupported histogram version"));
40        }
41        for (bucket, count) in &other.buckets {
42            if *bucket > 64 || *count < 0 {
43                return Err(super::invalid("Invalid histogram bucket"));
44            }
45            let current = self.buckets.entry(*bucket).or_default();
46            *current = current
47                .checked_add(*count)
48                .ok_or_else(|| super::invalid("Histogram overflow"))?;
49        }
50        Ok(())
51    }
52    pub fn percentile_upper_bound_micros(&self, percentile: u32) -> crate::Result<Option<u64>> {
53        if !(1..=100).contains(&percentile) || self.version != 1 {
54            return Err(super::invalid("Invalid histogram percentile or version"));
55        }
56        let count = self
57            .buckets
58            .values()
59            .try_fold(0i64, |sum, value| sum.checked_add(*value))
60            .ok_or_else(|| super::invalid("Histogram overflow"))?;
61        if count == 0 {
62            return Ok(None);
63        }
64        let target = (i128::from(count) * i128::from(percentile) + 99) / 100;
65        let mut running = 0i128;
66        for (bucket, count) in &self.buckets {
67            if *count < 0 || *bucket > 64 {
68                return Err(super::invalid("Invalid histogram"));
69            }
70            running += i128::from(*count);
71            if running >= target {
72                return Ok(Some(if *bucket == 0 {
73                    0
74                } else {
75                    1u64.checked_shl(*bucket)
76                        .map_or(u64::MAX, |value| value - 1)
77                }));
78            }
79        }
80        Err(super::invalid("Invalid histogram count"))
81    }
82}