fix(gateway): rewrite gauge_histogram.rs for zero-allocation hot path (#16878)

This commit is contained in:
Simo Lin
2026-01-10 18:15:21 -08:00
committed by GitHub
parent ff97814232
commit 7c25687c9b

View File

@@ -1,194 +1,583 @@
//! Gauge with gt/le bucket labels for Grafana heatmap visualization.
//! Non-cumulative gauge histogram for Grafana heatmap visualization.
//!
//! Unlike Prometheus Histogram which uses cumulative buckets, this uses
//! non-cumulative buckets (gt < value <= le) suitable for heatmap display.
//! Unlike Prometheus Histogram which uses cumulative `le` buckets, this emits
//! non-cumulative bucket counts with `(gt, le]` ranges suitable for heatmaps.
//!
//! Note: Keep in sync with Python implementation in
//! python/sglang/srt/utils/gauge_histogram.py
//! # Design: True Zero-Allocation Hot Path
//!
//! The key insight is that `gauge!` returns a `Gauge` handle that can be stored.
//! By pre-registering all gauge handles at startup, the hot path becomes just
//! N+1 atomic `gauge.set()` calls with zero allocations.
//!
//! # Performance Characteristics
//!
//! Setup (once per label combination):
//! - `register()`: N+1 gauge registrations, N+1 String allocations for gt/le
//!
//! Hot path (`set_counts()`):
//! - **Zero heap allocations**
//! - **Zero key lookups** (handles are pre-registered)
//! - N+1 atomic `gauge.set()` calls
//!
//! # Example
//!
//! ```ignore
//! use crate::observability::gauge_histogram::{BucketBounds, GaugeHistogramVec};
//!
//! // Define at module level
//! static BOUNDS: BucketBounds<10> = BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
//! static HISTOGRAM: GaugeHistogramVec<10> = GaugeHistogramVec::new("smg_request_dist", &BOUNDS);
//!
//! // At startup: register for each label combination
//! let handle = HISTOGRAM.register(&[("router", "round_robin"), ("model", "llama")]);
//!
//! // Pre-allocate counts buffer
//! let mut counts = vec![0usize; BOUNDS.bucket_count()];
//!
//! // Hot path: TRUE zero allocation
//! fn update(handle: &GaugeHistogramHandle, counts: &mut [usize], observations: &[u64]) {
//! BOUNDS.compute_counts_into(counts, observations);
//! handle.set_counts(counts); // Just N+1 atomic gauge.set() calls!
//! }
//! ```
use metrics::gauge;
use std::sync::Arc;
pub struct BucketLabels {
upper_bounds: &'static [u64],
le_labels: Vec<&'static str>,
gt_labels: Vec<&'static str>,
use dashmap::DashMap;
use metrics::{gauge, Label};
// =============================================================================
// BUCKET BOUNDS
// =============================================================================
/// Static bucket boundary configuration.
///
/// Uses const generics to define bucket bounds at compile time with validation.
/// The bounds define `N` upper limits, creating `N + 1` buckets:
/// `(0, b[0]], (b[0], b[1]], ..., (b[N-1], +Inf]`.
#[derive(Debug)]
pub struct BucketBounds<const N: usize> {
bounds: [u64; N],
}
impl BucketLabels {
pub fn new(upper_bounds: &'static [u64]) -> Self {
let leak_str = |n: u64| Box::leak(n.to_string().into_boxed_str()) as &'static str;
impl<const N: usize> BucketBounds<N> {
/// Create new bucket bounds from a sorted array of upper limits.
///
/// # Panics
///
/// Panics at compile time (in const context) or runtime if bounds are not
/// strictly ascending.
#[must_use]
pub const fn new(bounds: [u64; N]) -> Self {
let mut i = 1;
while i < N {
assert!(
bounds[i] > bounds[i - 1],
"bucket bounds must be strictly ascending"
);
i += 1;
}
Self { bounds }
}
let mut le_labels: Vec<&'static str> = upper_bounds.iter().map(|&b| leak_str(b)).collect();
le_labels.push("+Inf");
/// Returns the number of buckets (one more than the number of bounds).
#[inline]
#[must_use]
pub const fn bucket_count(&self) -> usize {
N + 1
}
let mut gt_labels: Vec<&'static str> = vec!["0"];
gt_labels.extend(upper_bounds.iter().map(|&b| leak_str(b)));
/// Returns the number of bounds.
#[inline]
#[must_use]
pub const fn bound_count(&self) -> usize {
N
}
Self {
upper_bounds,
le_labels,
gt_labels,
/// Get the bounds array.
#[inline]
#[must_use]
pub const fn bounds(&self) -> &[u64; N] {
&self.bounds
}
/// Find the bucket index for a value. O(log N).
#[inline]
#[must_use]
pub fn bucket_index(&self, value: u64) -> usize {
self.bounds.partition_point(|&bound| bound < value)
}
/// Get the upper bound for a bucket index, or None for the +Inf bucket.
#[inline]
#[must_use]
pub const fn upper_bound(&self, idx: usize) -> Option<u64> {
if idx < N {
Some(self.bounds[idx])
} else {
None
}
}
pub fn len(&self) -> usize {
self.le_labels.len()
/// Get the lower bound for a bucket index (0 for the first bucket).
#[inline]
#[must_use]
pub const fn lower_bound(&self, idx: usize) -> u64 {
if idx == 0 {
0
} else {
self.bounds[idx - 1]
}
}
pub fn iter(&self) -> impl Iterator<Item = (&'static str, &'static str)> + '_ {
std::iter::zip(&self.gt_labels, &self.le_labels).map(|(&gt, &le)| (gt, le))
}
/// Compute bucket counts from observations.
pub fn compute_bucket_counts(&self, observations: &[u64]) -> Vec<usize> {
let mut counts = vec![0usize; self.len()];
/// Compute bucket counts into a pre-allocated buffer. **Zero allocation.**
///
/// # Panics
///
/// Panics if `counts.len() < bucket_count()`.
#[inline]
pub fn compute_counts_into(&self, counts: &mut [usize], observations: &[u64]) {
debug_assert!(
counts.len() >= self.bucket_count(),
"counts buffer too small"
);
counts[..self.bucket_count()].fill(0);
for &value in observations {
// Equivalent to Python's bisect.bisect_left. O(log n).
let idx = self.upper_bounds.partition_point(|&bound| bound < value);
let idx = self.bucket_index(value);
counts[idx] += 1;
}
}
/// Compute bucket counts, allocating a new Vec.
///
/// Prefer `compute_counts_into` in hot paths to avoid allocation.
#[must_use]
pub fn compute_counts(&self, observations: &[u64]) -> Vec<usize> {
let mut counts = vec![0usize; self.bucket_count()];
self.compute_counts_into(&mut counts, observations);
counts
}
}
pub struct GaugeHistogram {
name: &'static str,
buckets: &'static BucketLabels,
// =============================================================================
// GAUGE HISTOGRAM HANDLE (pre-registered, zero-alloc hot path)
// =============================================================================
/// Pre-registered gauge handles for a histogram with specific labels.
///
/// This is what you use in the hot path. Calling `set_counts()` does only
/// N+1 atomic `gauge.set()` operations - zero allocations, zero lookups.
#[derive(Clone)]
pub struct GaugeHistogramHandle {
gauges: Vec<metrics::Gauge>,
}
impl GaugeHistogram {
pub const fn new(name: &'static str, buckets: &'static BucketLabels) -> Self {
Self { name, buckets }
}
/// Set bucket counts directly.
pub fn set_raw(&self, values: &[usize]) {
for ((gt, le), &count) in self.buckets.iter().zip(values.iter()) {
gauge!(self.name, "gt" => gt, "le" => le).set(count as f64);
impl GaugeHistogramHandle {
/// Set bucket counts. **TRUE zero allocation.**
///
/// Just N+1 atomic `gauge.set()` calls - no key lookup, no allocation.
#[inline]
pub fn set_counts(&self, counts: &[usize]) {
debug_assert_eq!(
counts.len(),
self.gauges.len(),
"counts length must match bucket count"
);
for (gauge, &count) in self.gauges.iter().zip(counts.iter()) {
gauge.set(count as f64);
}
}
/// Compute bucket counts from observations and set them.
pub fn set_by_current_observations(&self, observations: &[u64]) {
let counts = self.buckets.compute_bucket_counts(observations);
self.set_raw(&counts);
/// Number of buckets.
#[inline]
pub fn bucket_count(&self) -> usize {
self.gauges.len()
}
pub fn buckets(&self) -> &BucketLabels {
self.buckets
/// Zero out all gauges. **Zero allocation.**
#[inline]
pub fn zero_counts(&self) {
for gauge in &self.gauges {
gauge.set(0.0);
}
}
}
// =============================================================================
// GAUGE HISTOGRAM VEC (factory for registering handles)
// =============================================================================
/// Factory for creating pre-registered histogram handles.
///
/// Define as a static constant, then call `register()` for each label combination
/// you need. The returned `GaugeHistogramHandle` provides zero-allocation updates.
#[derive(Debug)]
pub struct GaugeHistogramVec<const N: usize> {
name: &'static str,
bounds: &'static BucketBounds<N>,
}
impl<const N: usize> GaugeHistogramVec<N> {
/// Create a new gauge histogram factory.
///
/// This just stores the name and bounds - no allocation or registration yet.
#[must_use]
pub const fn new(name: &'static str, bounds: &'static BucketBounds<N>) -> Self {
Self { name, bounds }
}
#[inline]
pub const fn name(&self) -> &'static str {
self.name
}
#[inline]
pub const fn bounds(&self) -> &BucketBounds<N> {
self.bounds
}
/// Register gauges for a specific label combination.
///
/// Call this once per unique label combination (at startup or when first seen).
/// The returned handle can be cloned cheaply (just Arc clones internally).
///
/// # Arguments
///
/// - `labels`: Static key-value label pairs for this histogram instance
///
/// # Example
///
/// ```ignore
/// let handle = HISTOGRAM.register(&[("router", "round_robin"), ("model", "llama")]);
/// ```
pub fn register(&self, labels: &[(&'static str, &str)]) -> GaugeHistogramHandle {
let bucket_count = self.bounds.bucket_count();
let mut gauges = Vec::with_capacity(bucket_count);
for i in 0..bucket_count {
// Build gt/le labels for this bucket
let gt_str = if i == 0 {
"0".to_string()
} else {
self.bounds.bounds[i - 1].to_string()
};
let le_str = if i < N {
self.bounds.bounds[i].to_string()
} else {
"+Inf".to_string()
};
// Build complete label set
let mut all_labels: Vec<Label> = Vec::with_capacity(labels.len() + 2);
for &(k, v) in labels {
all_labels.push(Label::new(k, v.to_string()));
}
all_labels.push(Label::new("gt", gt_str));
all_labels.push(Label::new("le", le_str));
// Register and store the gauge handle
let g = gauge!(self.name, all_labels);
gauges.push(g);
}
GaugeHistogramHandle { gauges }
}
/// Register with no additional labels (just gt/le).
pub fn register_no_labels(&self) -> GaugeHistogramHandle {
self.register(&[])
}
}
// =============================================================================
// CACHED GAUGE HISTOGRAM (for dynamic labels discovered at runtime)
// =============================================================================
/// A gauge histogram with automatic handle caching for dynamic labels.
///
/// Use this when label values (like worker names) are discovered at runtime.
/// Handles are registered on first use and cached for subsequent calls.
///
/// # Example
///
/// ```ignore
/// static BOUNDS: BucketBounds<10> = BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
/// static HISTOGRAM: GaugeHistogramVec<10> = GaugeHistogramVec::new("smg_worker_dist", &BOUNDS);
///
/// // Create cached wrapper (do this once, store in your router state)
/// let cached = CachedGaugeHistogram::new(&HISTOGRAM);
///
/// // Hot path - first call registers, subsequent calls use cached handle
/// cached.observe("worker-1", &request_counts);
/// cached.observe("worker-2", &request_counts);
/// cached.observe("worker-1", &request_counts); // Uses cached handle
/// ```
pub struct CachedGaugeHistogram<const N: usize> {
histogram: &'static GaugeHistogramVec<N>,
/// Cache of label value -> (handle, counts_buffer)
cache: DashMap<Arc<str>, (GaugeHistogramHandle, Vec<usize>)>,
/// Static label key (e.g., "worker", "model")
label_key: &'static str,
}
impl<const N: usize> CachedGaugeHistogram<N> {
/// Create a new cached histogram for a single dynamic label.
///
/// # Arguments
///
/// - `histogram`: The static histogram factory
/// - `label_key`: The label key for the dynamic value (e.g., "worker")
pub fn new(histogram: &'static GaugeHistogramVec<N>, label_key: &'static str) -> Self {
Self {
histogram,
cache: DashMap::new(),
label_key,
}
}
/// Get or create a handle for the given label value.
///
/// First call for a label value registers the gauges (allocates).
/// Subsequent calls return the cached handle (no allocation).
/// Thread-safe: uses entry API to avoid race conditions.
pub fn get_or_register(
&self,
label_value: &str,
) -> dashmap::mapref::one::Ref<'_, Arc<str>, (GaugeHistogramHandle, Vec<usize>)> {
// Fast path: already cached
if let Some(entry) = self.cache.get(label_value) {
return entry;
}
// Slow path: use entry API to handle concurrent inserts atomically
self.cache.entry(Arc::from(label_value)).or_insert_with(|| {
let handle = self.histogram.register(&[(self.label_key, label_value)]);
let counts_buf = vec![0usize; self.histogram.bounds.bucket_count()];
(handle, counts_buf)
});
self.cache.get(label_value).unwrap()
}
/// Observe a distribution for a label value. **Zero allocation after first call.**
///
/// First call for a new label value registers gauges (allocates).
/// All subsequent calls are zero-allocation.
/// Thread-safe: uses entry API to avoid race conditions.
pub fn observe(&self, label_value: &str, observations: &[u64]) {
// Fast path: existing entry
if let Some(mut entry) = self.cache.get_mut(label_value) {
let (ref handle, ref mut counts_buf) = entry.value_mut();
self.histogram
.bounds
.compute_counts_into(counts_buf, observations);
handle.set_counts(counts_buf);
return;
}
// Slow path: use entry API to handle concurrent inserts atomically
let mut entry = self.cache.entry(Arc::from(label_value)).or_insert_with(|| {
let handle = self.histogram.register(&[(self.label_key, label_value)]);
let counts_buf = vec![0usize; self.histogram.bounds.bucket_count()];
(handle, counts_buf)
});
let (ref handle, ref mut counts_buf) = entry.value_mut();
self.histogram
.bounds
.compute_counts_into(counts_buf, observations);
handle.set_counts(counts_buf);
}
/// Number of cached label combinations.
pub fn cache_size(&self) -> usize {
self.cache.len()
}
/// Remove a worker and zero out its metrics. **Zero allocation.**
///
/// Call this when a worker is removed from the pool.
/// Sets all bucket counts to 0 (so Grafana shows it as empty).
///
/// Note: The gauge handles remain in the Prometheus registry (the `metrics`
/// crate doesn't support unregistering). But memory in our cache is freed.
pub fn remove(&self, label_value: &str) {
if let Some((_, (handle, _))) = self.cache.remove(label_value) {
handle.zero_counts();
}
}
/// Remove workers not in the provided set.
///
/// Call this periodically with your current active workers to clean up stale entries.
/// Uses `DashMap::retain` for atomic operation without intermediate allocation.
///
/// # Example
///
/// ```ignore
/// let active: HashSet<&str> = workers.iter().map(|w| w.name.as_str()).collect();
/// cached.retain_only(&active);
/// ```
pub fn retain_only<S: std::borrow::Borrow<str> + std::hash::Hash + Eq>(
&self,
active_labels: &std::collections::HashSet<S>,
) {
self.cache.retain(|key, (handle, _)| {
if active_labels.contains(key.as_ref()) {
true
} else {
handle.zero_counts();
false
}
});
}
/// Get all currently tracked label values.
pub fn tracked_labels(&self) -> Vec<Arc<str>> {
self.cache.iter().map(|e| Arc::clone(e.key())).collect()
}
}
// =============================================================================
// CONVENIENCE CONSTANTS
// =============================================================================
/// Common bucket bounds for request counts.
pub static REQUEST_COUNT_BOUNDS: BucketBounds<10> =
BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
// =============================================================================
// TESTS
// =============================================================================
#[cfg(test)]
mod tests {
use super::*;
// --- Bucket Labels Tests ---
#[test]
fn test_bucket_bounds_creation() {
let bounds = BucketBounds::new([10, 30, 60]);
assert_eq!(bounds.bucket_count(), 4);
assert_eq!(bounds.bound_count(), 3);
}
#[test]
fn test_bucket_labels_basic() {
let buckets = BucketLabels::new(&[10, 30, 60]);
let pairs: Vec<_> = buckets.iter().collect();
fn test_bucket_bounds_const_creation() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
assert_eq!(BOUNDS.bucket_count(), 4);
}
#[test]
#[should_panic(expected = "bucket bounds must be strictly ascending")]
fn test_bucket_bounds_not_ascending_panics() {
let _ = BucketBounds::new([10, 5, 60]);
}
#[test]
fn test_bucket_index() {
let bounds = BucketBounds::new([10, 30, 60]);
assert_eq!(bounds.bucket_index(0), 0);
assert_eq!(bounds.bucket_index(10), 0);
assert_eq!(bounds.bucket_index(11), 1);
assert_eq!(bounds.bucket_index(30), 1);
assert_eq!(bounds.bucket_index(31), 2);
assert_eq!(bounds.bucket_index(60), 2);
assert_eq!(bounds.bucket_index(61), 3);
}
#[test]
fn test_compute_counts() {
let bounds = BucketBounds::new([10, 30, 60]);
assert_eq!(
pairs,
vec![("0", "10"), ("10", "30"), ("30", "60"), ("60", "+Inf")]
);
}
#[test]
fn test_bucket_labels_single_bound() {
let buckets = BucketLabels::new(&[100]);
let pairs: Vec<_> = buckets.iter().collect();
assert_eq!(pairs, vec![("0", "100"), ("100", "+Inf")]);
}
#[test]
fn test_bucket_labels_many_bounds() {
let buckets = BucketLabels::new(&[1, 2, 5, 10]);
let pairs: Vec<_> = buckets.iter().collect();
assert_eq!(
pairs,
vec![
("0", "1"),
("1", "2"),
("2", "5"),
("5", "10"),
("10", "+Inf")
]
);
}
#[test]
fn test_bucket_labels_len() {
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(buckets.len(), 4);
}
// --- Bucket Counts Tests ---
#[test]
fn test_compute_bucket_counts_empty() {
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(buckets.compute_bucket_counts(&[]), vec![0, 0, 0, 0]);
}
#[test]
fn test_compute_bucket_counts_single_value_first_bucket() {
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(buckets.compute_bucket_counts(&[5]), vec![1, 0, 0, 0]);
}
#[test]
fn test_compute_bucket_counts_single_value_last_bucket() {
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(buckets.compute_bucket_counts(&[100]), vec![0, 0, 0, 1]);
}
#[test]
fn test_compute_bucket_counts_exact_boundary_values() {
// Values at exact boundaries: 10 -> (0,10], 30 -> (10,30], 60 -> (30,60]
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(
buckets.compute_bucket_counts(&[10, 30, 60]),
vec![1, 1, 1, 0]
);
}
#[test]
fn test_compute_bucket_counts_just_above_boundary() {
// 11 -> (10,30], 31 -> (30,60], 61 -> (60,+Inf]
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(
buckets.compute_bucket_counts(&[11, 31, 61]),
vec![0, 1, 1, 1]
);
}
#[test]
fn test_compute_bucket_counts_multiple_values_same_bucket() {
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(
buckets.compute_bucket_counts(&[1, 2, 3, 4, 5]),
vec![5, 0, 0, 0]
);
}
#[test]
fn test_compute_bucket_counts_all_overflow() {
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(
buckets.compute_bucket_counts(&[100, 200, 300]),
vec![0, 0, 0, 3]
);
}
#[test]
fn test_compute_bucket_counts_distribution() {
// 5 (<=10), 10 (<=10), 15 (<=30), 40 (<=60), 100 (+Inf)
let buckets = BucketLabels::new(&[10, 30, 60]);
assert_eq!(
buckets.compute_bucket_counts(&[5, 10, 15, 40, 100]),
bounds.compute_counts(&[5, 10, 15, 40, 100]),
vec![2, 1, 1, 1]
);
}
#[test]
fn test_compute_counts_into() {
let bounds = BucketBounds::new([10, 30, 60]);
let mut counts = [0usize; 4];
bounds.compute_counts_into(&mut counts, &[5, 10, 15, 40, 100]);
assert_eq!(counts, [2, 1, 1, 1]);
}
#[test]
fn test_gauge_histogram_vec_creation() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_metric", &BOUNDS);
assert_eq!(HISTOGRAM.name(), "test_metric");
assert_eq!(HISTOGRAM.bounds().bucket_count(), 4);
}
#[test]
fn test_gauge_histogram_handle_registration() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_hist", &BOUNDS);
let handle = HISTOGRAM.register(&[("router", "rr")]);
assert_eq!(handle.bucket_count(), 4);
// This should be zero-allocation
handle.set_counts(&[1, 2, 3, 4]);
}
#[test]
fn test_request_count_bounds() {
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_count(), 11);
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(1), 0);
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(2), 1);
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(201), 10);
}
#[test]
fn test_cached_histogram() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_cached", &BOUNDS);
let cached = CachedGaugeHistogram::new(&HISTOGRAM, "worker");
// First call registers
cached.observe("worker-1", &[5, 15, 45, 100]);
assert_eq!(cached.cache_size(), 1);
// Second call uses cache
cached.observe("worker-1", &[1, 2, 3]);
assert_eq!(cached.cache_size(), 1);
// New worker registers
cached.observe("worker-2", &[10, 20, 30]);
assert_eq!(cached.cache_size(), 2);
}
#[test]
fn test_cached_histogram_removal() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> =
GaugeHistogramVec::new("test_cached_remove", &BOUNDS);
let cached = CachedGaugeHistogram::new(&HISTOGRAM, "worker");
// Add some workers
cached.observe("worker-1", &[5, 15]);
cached.observe("worker-2", &[10, 20]);
cached.observe("worker-3", &[1, 2]);
assert_eq!(cached.cache_size(), 3);
// Remove one
cached.remove("worker-2");
assert_eq!(cached.cache_size(), 2);
// retain_only
let active: std::collections::HashSet<&str> = ["worker-1"].into_iter().collect();
cached.retain_only(&active);
assert_eq!(cached.cache_size(), 1);
// Check tracked labels
let labels = cached.tracked_labels();
assert_eq!(labels.len(), 1);
assert_eq!(&*labels[0], "worker-1");
}
}