Skip to main content

av_core/
metrics.rs

1//! Dependency-free metrics registry rendering Prometheus text exposition
2//! format. Counters and histograms only (what the SLA criteria need), all
3//! lock-free on the hot path (atomics; registration takes a short mutex).
4
5use parking_lot::Mutex;
6use std::collections::BTreeMap;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9
10/// A monotonically increasing counter.
11#[derive(Debug, Default)]
12pub struct Counter {
13    value: AtomicU64,
14}
15
16impl Counter {
17    /// Increment by 1.
18    pub fn inc(&self) {
19        self.value.fetch_add(1, Ordering::Relaxed);
20    }
21
22    /// Increment by `n`.
23    pub fn add(&self, n: u64) {
24        self.value.fetch_add(n, Ordering::Relaxed);
25    }
26
27    /// Current value.
28    pub fn get(&self) -> u64 {
29        self.value.load(Ordering::Relaxed)
30    }
31}
32
33/// A signed gauge whose value may go up, down, or be `set` to an
34/// absolute number. Backed by `AtomicI64` because Prometheus gauges
35/// are unrestricted-sign floats; we clamp to i64 here because every
36/// data-plane gauge in the workspace (open sessions, queue depth,
37/// spool bytes) is a non-negative count that fits comfortably in
38/// 63 bits, and the signed representation lets `sub_saturating`
39/// bottom out at 0 rather than wrap.
40#[derive(Debug, Default)]
41pub struct Gauge {
42    value: std::sync::atomic::AtomicI64,
43}
44
45impl Gauge {
46    /// Set the gauge to an absolute value. Truncates to `i64::MAX`
47    /// on overflow — no data-plane count in the workspace is even
48    /// close to that ceiling, and silently overflowing to a
49    /// negative value would produce misleading dashboards.
50    pub fn set(&self, value: u64) {
51        let clamped = i64::try_from(value).unwrap_or(i64::MAX);
52        self.value.store(clamped, Ordering::Relaxed);
53    }
54
55    /// Increment by 1. Saturates at `i64::MAX`.
56    pub fn inc(&self) {
57        self.add(1);
58    }
59
60    /// Decrement by 1. Saturates at 0 — a gauge going negative is
61    /// almost always a lifecycle bookkeeping bug (dec called twice
62    /// for one inc), and rendering `-1 open_sessions` would confuse
63    /// every dashboard downstream.
64    pub fn dec(&self) {
65        self.sub(1);
66    }
67
68    /// Increment by `n`. Saturates at `i64::MAX`.
69    pub fn add(&self, n: u64) {
70        let delta = i64::try_from(n).unwrap_or(i64::MAX);
71        // A signed `fetch_add` with a positive delta saturates
72        // silently to `i64::MIN` on overflow (two's complement
73        // wrap). Loop-CAS to a saturated add so overflow lands at
74        // `i64::MAX` — never negative.
75        let mut current = self.value.load(Ordering::Relaxed);
76        loop {
77            let next = current.saturating_add(delta);
78            match self
79                .value
80                .compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed)
81            {
82                Ok(_) => return,
83                Err(observed) => current = observed,
84            }
85        }
86    }
87
88    /// Decrement by `n`. Saturates at 0.
89    pub fn sub(&self, n: u64) {
90        let delta = i64::try_from(n).unwrap_or(i64::MAX);
91        let mut current = self.value.load(Ordering::Relaxed);
92        loop {
93            let next = current.saturating_sub(delta).max(0);
94            match self
95                .value
96                .compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed)
97            {
98                Ok(_) => return,
99                Err(observed) => current = observed,
100            }
101        }
102    }
103
104    /// Current value as an unsigned integer — the saturating math
105    /// above guarantees non-negativity so the cast never truncates
106    /// a real observation.
107    pub fn get(&self) -> u64 {
108        self.value.load(Ordering::Relaxed).max(0) as u64
109    }
110}
111
112/// Fixed-bucket latency histogram (microsecond samples).
113///
114/// Buckets are cumulative on render, per Prometheus convention.
115#[derive(Debug)]
116pub struct Histogram {
117    bounds_us: Vec<u64>,
118    buckets: Vec<AtomicU64>,
119    count: AtomicU64,
120    sum_us: AtomicU64,
121    /// Samples above every declared bucket. Represents the implicit
122    /// `le="+Inf"` bucket in Prometheus terms and lets `quantile_us`
123    /// distinguish "quantile equals the top bound" from "quantile lies
124    /// beyond the top bound".
125    overflow: AtomicU64,
126}
127
128/// Default latency bucket upper bounds in microseconds: 50µs … 10s.
129///
130/// Sized for fast internal stages (identity check, quota, sanitize,
131/// dashboard endpoints, receipt sign). Long-running histograms — the
132/// upstream `dispatch` call which includes provider streaming, and the
133/// filesystem-scan reconciler / finalize paths — must use
134/// [`WIDE_LATENCY_BOUNDS_US`] instead, or every observation lands in
135/// the `+Inf` overflow bucket and p95/p99 renders as `u64::MAX`.
136pub const DEFAULT_LATENCY_BOUNDS_US: &[u64] = &[
137    50, 100, 250, 500, 1_000, 2_000, 5_000, 8_000, 10_000, 25_000, 50_000, 100_000, 1_000_000, 10_000_000,
138];
139
140/// Wide latency bucket upper bounds in microseconds: 1ms … 300s.
141///
142/// Use for histograms that measure spans dominated by network I/O
143/// (upstream LLM streaming, reconciler recovery scans, finalisation
144/// under load). Provider p99 for GPT-4o / Claude regularly sits in the
145/// 15–90 s band; the top bound of 300 s covers pathological long
146/// contexts without saturating for realistic operator SLA use.
147pub const WIDE_LATENCY_BOUNDS_US: &[u64] = &[
148    1_000,
149    5_000,
150    10_000,
151    100_000,
152    500_000,
153    1_000_000,
154    5_000_000,
155    10_000_000,
156    30_000_000,
157    60_000_000,
158    120_000_000,
159    300_000_000,
160];
161
162impl Histogram {
163    /// Create a histogram with the given bucket upper bounds (µs, ascending).
164    pub fn new(bounds_us: &[u64]) -> Self {
165        Self {
166            bounds_us: bounds_us.to_vec(),
167            buckets: bounds_us.iter().map(|_| AtomicU64::new(0)).collect(),
168            count: AtomicU64::new(0),
169            sum_us: AtomicU64::new(0),
170            overflow: AtomicU64::new(0),
171        }
172    }
173
174    /// Record a sample in microseconds.
175    pub fn observe_us(&self, us: u64) {
176        self.count.fetch_add(1, Ordering::Relaxed);
177        self.sum_us.fetch_add(us, Ordering::Relaxed);
178        for (i, b) in self.bounds_us.iter().enumerate() {
179            if us <= *b {
180                if let Some(bucket) = self.buckets.get(i) {
181                    bucket.fetch_add(1, Ordering::Relaxed);
182                }
183                return;
184            }
185        }
186        // Sample landed above every declared bucket. Track it so the
187        // renderer can emit the standard Prometheus `le="+Inf"` bucket
188        // (bounded-bucket sum otherwise equals `count` minus the tail,
189        // which is invalid Prometheus text) and so operators can spot a
190        // regime where the true P99 is beyond the top bound.
191        self.overflow.fetch_add(1, Ordering::Relaxed);
192    }
193
194    /// Total number of samples.
195    pub fn count(&self) -> u64 {
196        self.count.load(Ordering::Relaxed)
197    }
198
199    /// Approximate quantile (µs) from bucket boundaries. Returns the upper
200    /// bound of the bucket containing quantile `q` (0.0–1.0); a target that
201    /// lands above every bounded bucket returns the sentinel below.
202    ///
203    /// A `q` whose target is served only by samples in the implicit +Inf
204    /// bucket returns `u64::MAX` as a sentinel: this signals "beyond top
205    /// bucket" instead of silently under-reporting `bounds_us.last()`.
206    pub fn quantile_us(&self, q: f64) -> u64 {
207        let total = self.count();
208        if total == 0 {
209            return 0;
210        }
211        #[allow(
212            clippy::cast_possible_truncation,
213            clippy::cast_sign_loss,
214            clippy::cast_precision_loss
215        )]
216        // Clamp to >= 1: with q = 0.0 the raw target is 0 and `cum >= 0`
217        // holds before any sample is accumulated, which would report the
218        // first bucket bound even when that bucket is empty. p0 must be
219        // the first bucket that actually contains a sample.
220        let target = (((total as f64) * q.clamp(0.0, 1.0)).ceil() as u64).max(1);
221        let mut cum = 0u64;
222        for (i, b) in self.bounds_us.iter().enumerate() {
223            cum += self.buckets.get(i).map_or(0, |x| x.load(Ordering::Relaxed));
224            if cum >= target {
225                return *b;
226            }
227        }
228        // Bounded buckets could not reach the target: the target must sit in
229        // the +Inf overflow bucket. Surface that with u64::MAX rather than
230        // returning bounds_us.last() and silently under-reporting.
231        if self.overflow.load(Ordering::Relaxed) > 0 {
232            return u64::MAX;
233        }
234        self.bounds_us.last().copied().unwrap_or(0)
235    }
236}
237
238enum Metric {
239    Counter(Arc<Counter>),
240    Gauge(Arc<Gauge>),
241    Histogram(Arc<Histogram>),
242}
243
244/// A registry mapping metric names (+ optional fixed labels) to metrics.
245#[derive(Default)]
246pub struct Registry {
247    metrics: Mutex<BTreeMap<String, (String, Metric)>>,
248    /// Track the metric KIND (counter vs histogram) per base name — the
249    /// full-key type-collision guard in `counter`/`histogram_with_bounds`
250    /// only catches same-key clashes, but Prometheus text exposition
251    /// requires that ALL variants of a base name (across every label
252    /// combination) share the same type. Two contributors registering
253    /// `av_foo_total{stage="a"}` as a counter and
254    /// `av_foo_total{stage="b"}` as a histogram would produce an
255    /// invalid `# TYPE` header and Prometheus's parser would reject
256    /// the whole scrape.
257    base_kinds: Mutex<BTreeMap<String, MetricKind>>,
258}
259
260#[derive(Debug, Clone, Copy, PartialEq, Eq)]
261enum MetricKind {
262    Counter,
263    Gauge,
264    Histogram,
265    /// Not a registrable kind: a name reserved because a histogram's
266    /// text exposition generates it (`<base>_bucket`, `<base>_sum`,
267    /// `<base>_count`). Conflicts with EVERY registration — including
268    /// another histogram — since any metric under this name would
269    /// emit series/TYPE lines colliding with the generating
270    /// histogram's rendered samples.
271    HistogramReserved,
272}
273
274impl MetricKind {
275    fn label(self) -> &'static str {
276        match self {
277            Self::Counter => "counter",
278            Self::Gauge => "gauge",
279            Self::Histogram => "histogram",
280            Self::HistogramReserved => "histogram-generated series (_bucket/_sum/_count)",
281        }
282    }
283}
284
285impl Registry {
286    /// Create an empty registry.
287    pub fn new() -> Self {
288        Self::default()
289    }
290
291    /// Register (or fetch) a counter. `key` must be a valid Prometheus series
292    /// name with optional `{label="v"}` suffix.
293    ///
294    /// Panics if `key` is already registered as a histogram; a silent
295    /// overwrite would detach the existing metric from the registry and
296    /// lose every observation collected against it. Also panics if `key`
297    /// contains a byte that would corrupt Prometheus text exposition
298    /// (cross-line bytes anywhere; `"` or `\\` in the base name — both
299    /// stay legal inside the label section) — this catches accidental
300    /// interpolation of
301    /// attacker-controlled strings into a metric key at registration
302    /// rather than surfacing corrupt scrape output later.
303    #[allow(clippy::panic)]
304    pub fn counter(&self, key: &str, help: &str) -> Arc<Counter> {
305        validate_metric_key(key);
306        self.reserve_base_kind(key, MetricKind::Counter);
307        let mut m = self.metrics.lock();
308        match m.get(key) {
309            Some((existing_help, Metric::Counter(c))) => {
310                // R73 F4 (landed R74): warn on HELP-string drift
311                // across registration sites. `MetricRegistry` is
312                // first-wins on HELP, so a rich forensic HELP
313                // registered at the actual observation site
314                // (e.g. `reconciler.rs`) is silently shadowed by
315                // an earlier terse HELP at boot (e.g.
316                // `pipeline.rs`), quietly defeating alerting
317                // playbooks that key on the descriptive text.
318                // Debug-only warn (not panic) so a drifting
319                // refactor is visible in tracing output but doesn't
320                // break the boot path — the correct fix is to
321                // unify the HELP text across the two sites (extract
322                // to a `const`), and the warn surfaces exactly
323                // which pairs still drift.
324                if cfg!(debug_assertions) && existing_help != help {
325                    tracing::warn!(
326                        target: "av_core::metrics",
327                        metric = key,
328                        existing_help = %existing_help,
329                        new_help = %help,
330                        "counter registered with divergent HELP; first-wins semantic — extract HELP to a shared const to eliminate drift"
331                    );
332                }
333                Arc::clone(c)
334            }
335            Some((_, Metric::Gauge(_))) => panic!(
336                "metric name conflict: {key:?} is already registered as a gauge, \
337                 cannot register as a counter",
338            ),
339            Some((_, Metric::Histogram(_))) => panic!(
340                "metric name conflict: {key:?} is already registered as a histogram, \
341                 cannot register as a counter",
342            ),
343            None => {
344                let c = Arc::new(Counter::default());
345                m.insert(key.to_owned(), (help.to_owned(), Metric::Counter(Arc::clone(&c))));
346                c
347            }
348        }
349    }
350
351    /// Register (or fetch) a gauge. Same key/label conventions as
352    /// [`Self::counter`]; gauges may go up, down, and be reset via
353    /// [`Gauge::set`]. Panics on kind conflict with an existing
354    /// counter or histogram registered under the same base name.
355    #[allow(clippy::panic)]
356    pub fn gauge(&self, key: &str, help: &str) -> Arc<Gauge> {
357        validate_metric_key(key);
358        self.reserve_base_kind(key, MetricKind::Gauge);
359        let mut m = self.metrics.lock();
360        match m.get(key) {
361            Some((existing_help, Metric::Gauge(g))) => {
362                if cfg!(debug_assertions) && existing_help != help {
363                    tracing::warn!(
364                        target: "av_core::metrics",
365                        metric = key,
366                        existing_help = %existing_help,
367                        new_help = %help,
368                        "gauge registered with divergent HELP; first-wins semantic — extract HELP to a shared const to eliminate drift"
369                    );
370                }
371                Arc::clone(g)
372            }
373            Some((_, Metric::Counter(_))) => panic!(
374                "metric name conflict: {key:?} is already registered as a counter, \
375                 cannot register as a gauge",
376            ),
377            Some((_, Metric::Histogram(_))) => panic!(
378                "metric name conflict: {key:?} is already registered as a histogram, \
379                 cannot register as a gauge",
380            ),
381            None => {
382                let g = Arc::new(Gauge::default());
383                m.insert(key.to_owned(), (help.to_owned(), Metric::Gauge(Arc::clone(&g))));
384                g
385            }
386        }
387    }
388
389    /// Register (or fetch) a histogram with default latency buckets.
390    ///
391    /// Panics if `key` is already registered as a counter; a silent overwrite
392    /// would detach the existing metric from the registry and lose every
393    /// observation collected against it. Same key-byte guard as
394    /// [`Self::counter`].
395    #[allow(clippy::panic)]
396    pub fn histogram(&self, key: &str, help: &str) -> Arc<Histogram> {
397        self.histogram_with_bounds(key, help, DEFAULT_LATENCY_BOUNDS_US)
398    }
399
400    /// Register (or fetch) a histogram with explicit bucket bounds. Use
401    /// [`WIDE_LATENCY_BOUNDS_US`] for spans dominated by network I/O
402    /// (upstream LLM streaming, reconciler filesystem scans); use
403    /// [`DEFAULT_LATENCY_BOUNDS_US`] for fast internal stages.
404    ///
405    /// # Panics
406    ///
407    /// Panics if `key` is already registered as a counter; a silent
408    /// overwrite would detach the existing metric from the registry
409    /// and lose every observation collected against it. Same
410    /// key-byte guard as [`Self::counter`] applies.
411    #[allow(clippy::panic)]
412    pub fn histogram_with_bounds(&self, key: &str, help: &str, bounds_us: &[u64]) -> Arc<Histogram> {
413        validate_metric_key(key);
414        self.reserve_base_kind(key, MetricKind::Histogram);
415        let mut m = self.metrics.lock();
416        match m.get(key) {
417            Some((existing_help, Metric::Histogram(h))) => {
418                if cfg!(debug_assertions) && existing_help != help {
419                    tracing::warn!(
420                        target: "av_core::metrics",
421                        metric = key,
422                        existing_help = %existing_help,
423                        new_help = %help,
424                        "histogram registered with divergent HELP; first-wins semantic — extract HELP to a shared const to eliminate drift"
425                    );
426                }
427                Arc::clone(h)
428            }
429            Some((_, Metric::Counter(_))) => panic!(
430                "metric name conflict: {key:?} is already registered as a counter, \
431                 cannot register as a histogram",
432            ),
433            Some((_, Metric::Gauge(_))) => panic!(
434                "metric name conflict: {key:?} is already registered as a gauge, \
435                 cannot register as a histogram",
436            ),
437            None => {
438                let h = Arc::new(Histogram::new(bounds_us));
439                m.insert(
440                    key.to_owned(),
441                    (help.to_owned(), Metric::Histogram(Arc::clone(&h))),
442                );
443                h
444            }
445        }
446    }
447
448    /// Enforce that a base name (the metric name up to `{`) is only
449    /// ever registered as one metric kind across every label
450    /// combination — a Prometheus text-exposition invariant. Panics
451    /// with a clear message on mismatch. Called from `counter` and
452    /// `histogram_with_bounds` before the per-key type check.
453    #[allow(clippy::panic)]
454    fn reserve_base_kind(&self, key: &str, kind: MetricKind) {
455        let (base, _labels) = split_key(key);
456        let mut kinds = self.base_kinds.lock();
457        // A histogram's exposition OWNS three generated series names
458        // beyond its base: `<base>_bucket`, `<base>_sum`,
459        // `<base>_count`. A counter registered later as, say,
460        // `latency_seconds_count` would render a second bare
461        // `latency_seconds_count` sample under its own TYPE header —
462        // duplicate series + conflicting TYPE claims in one scrape,
463        // which Prometheus rejects wholesale. Reserve the generated
464        // names alongside the base; the same loop symmetrically
465        // panics when a histogram's generated names collide with
466        // anything already registered (either order fails fast).
467        let reserved: Vec<(String, MetricKind)> = match kind {
468            MetricKind::Histogram => vec![
469                (base.to_owned(), kind),
470                (format!("{base}_bucket"), MetricKind::HistogramReserved),
471                (format!("{base}_sum"), MetricKind::HistogramReserved),
472                (format!("{base}_count"), MetricKind::HistogramReserved),
473            ],
474            _ => vec![(base.to_owned(), kind)],
475        };
476        for (name, reserve_as) in reserved {
477            match kinds.get(name.as_str()) {
478                Some(existing) if *existing != reserve_as => panic!(
479                    "metric base-name kind conflict: {name:?} is already registered as \
480                     `{}`; cannot register {key:?} as `{}`. \
481                     Prometheus rejects mismatched TYPE headers across variants of the \
482                     same base name and the whole scrape becomes invalid text exposition.",
483                    existing.label(),
484                    reserve_as.label(),
485                ),
486                _ => {
487                    kinds.insert(name, reserve_as);
488                }
489            }
490        }
491    }
492
493    /// Render the registry in Prometheus text exposition format.
494    ///
495    /// Histogram observations are stored internally in microseconds but
496    /// rendered in seconds (`le` bounds and `_sum`), per Prometheus base-unit
497    /// convention. Histogram metric names should therefore end in `_seconds`.
498    ///
499    /// HELP text is escaped per the Prometheus text format
500    /// spec (`\` → `\\`, LF → `\n`). A future counter/histogram
501    /// registration whose HELP contained a newline would otherwise
502    /// silently corrupt the scrape response — Prometheus would parse
503    /// the remainder of the HELP text as metric samples and fail.
504    pub fn render(&self) -> String {
505        let m = self.metrics.lock();
506        let mut out = String::new();
507        let mut declared = std::collections::BTreeSet::new();
508        for (key, (help, metric)) in m.iter() {
509            let (base, labels) = split_key(key);
510            let help = escape_prom_help(help);
511            match metric {
512                Metric::Counter(c) => {
513                    if declared.insert(base.to_owned()) {
514                        out.push_str(&format!("# HELP {base} {help}\n# TYPE {base} counter\n"));
515                    }
516                    out.push_str(&format!("{key} {}\n", c.get()));
517                }
518                Metric::Gauge(g) => {
519                    if declared.insert(base.to_owned()) {
520                        out.push_str(&format!("# HELP {base} {help}\n# TYPE {base} gauge\n"));
521                    }
522                    out.push_str(&format!("{key} {}\n", g.get()));
523                }
524                Metric::Histogram(h) => {
525                    if declared.insert(base.to_owned()) {
526                        out.push_str(&format!("# HELP {base} {help}\n# TYPE {base} histogram\n"));
527                    }
528                    // Snapshot the count ONCE before rendering either
529                    // the `+Inf` bucket or the `_count` line, AND
530                    // clamp the cumulative bucket sum against the
531                    // snapshot. Two invariants strict OpenMetrics
532                    // parsers enforce:
533                    //
534                    //   (a) `_count == cum(bucket{le="+Inf"})`
535                    //   (b) `bucket{le=X} <= bucket{le=Y}` for X < Y,
536                    //       and `bucket{le=X} <= _count` for all X
537                    //
538                    // `observe_us` writes `count.fetch_add` BEFORE
539                    // `bucket.fetch_add` (metrics.rs:176 → :181), so
540                    // a naive render that reads count first and
541                    // buckets after can see MORE bucket increments
542                    // than count increments — a concurrent observe
543                    // that completed both fetches between the
544                    // count snapshot and the bucket load. The
545                    // previous fix (R24 review) snapshotted count
546                    // once but didn't clamp; that closed invariant
547                    // (a) but re-opened invariant (b), because
548                    // `cum` could exceed `count_snapshot` on the
549                    // very next bucket iteration and get emitted
550                    // verbatim.
551                    //
552                    // Clamping cum against count_snapshot preserves
553                    // both invariants: bucket lines are monotone
554                    // non-decreasing (cum is monotone, `.min(K)` of
555                    // a monotone sequence is monotone), and every
556                    // bucket line is `<= count_snapshot` by
557                    // construction. Under a race, a bucket that
558                    // truly holds `N+1` observations may render as
559                    // `N` for one scrape; the missing observation
560                    // shows up on the next scrape. This is the same
561                    // "eventually consistent under concurrent
562                    // observation" contract every Prometheus client
563                    // library the ecosystem trusts uses.
564                    let count_snapshot = h.count();
565                    let mut cum = 0u64;
566                    for (i, b) in h.bounds_us.iter().enumerate() {
567                        cum += h.buckets.get(i).map_or(0, |x| x.load(Ordering::Relaxed));
568                        let clamped = cum.min(count_snapshot);
569                        let le = (*b as f64) / 1_000_000.0;
570                        out.push_str(&format!(
571                            "{base}_bucket{{{}le=\"{le}\"}} {clamped}\n",
572                            join_labels(labels)
573                        ));
574                    }
575                    out.push_str(&format!(
576                        "{base}_bucket{{{}le=\"+Inf\"}} {count_snapshot}\n",
577                        join_labels(labels),
578                    ));
579                    let sum_s = (h.sum_us.load(Ordering::Relaxed) as f64) / 1_000_000.0;
580                    out.push_str(&format!(
581                        "{base}_sum{labels_block} {sum_s}\n",
582                        labels_block = labels_suffix(labels)
583                    ));
584                    out.push_str(&format!(
585                        "{base}_count{labels_block} {count_snapshot}\n",
586                        labels_block = labels_suffix(labels)
587                    ));
588                }
589            }
590        }
591        out
592    }
593}
594
595/// Escape a HELP text per the Prometheus text exposition
596/// format spec. Backslash and line-feed are the only two chars the
597/// format reserves in HELP lines. A future counter/histogram
598/// registration whose HELP contained a newline would otherwise
599/// silently corrupt the scrape response — Prometheus would parse
600/// the remainder of the HELP text as metric samples and fail.
601fn escape_prom_help(help: &str) -> String {
602    let mut out = String::with_capacity(help.len());
603    for c in help.chars() {
604        match c {
605            '\\' => out.push_str("\\\\"),
606            '\n' => out.push_str("\\n"),
607            // Replace CR
608            // and every other C0 control (0x00–0x1F except LF
609            // which we just escaped) with a literal space.
610            // `validate_metric_key` already refuses these bytes
611            // in metric keys; the HELP-side didn't. NUL trips
612            // promtool lint / grafana-agent; ESC (0x1B) lets a
613            // caller who controls HELP text inject ANSI codes
614            // into operator terminals via `avctl` piped
615            // `/metrics`. The substitution also covers DEL (0x7F) and the C1 range (0x80..=0x9F) —
616            // CSI (0x9B) is a valid single-byte ANSI escape
617            // prefix under 8-bit terminal emulation, so a HELP
618            // text carrying a `\u{9b}...` sequence renders the
619            // same way in an 8-bit-clean terminal as the C0
620            // ESC+`[` prefix we already scrub. Space is the
621            // same convention prometheus_client-python's
622            // `_ESCAPE_RE` uses.
623            c if (c as u32) < 0x20 || (c as u32) == 0x7f || (0x80..=0x9f).contains(&(c as u32)) => {
624                out.push(' ');
625            }
626            other => out.push(other),
627        }
628    }
629    out
630}
631
632/// Reject metric keys that would corrupt the Prometheus text exposition
633/// format. Legitimate keys carry label values like `{stage="identity"}` and
634/// must be allowed to contain double quotes and `=`, but a newline,
635/// carriage return, backslash, or NUL byte inside a key would split the
636/// scrape line and produce invalid text. These bytes arrive only from
637/// callers accidentally interpolating attacker-influenced strings into a
638/// series name; catch that at registration.
639#[allow(clippy::panic)]
640fn validate_metric_key(key: &str) {
641    // Split the key into base + labels first, then apply
642    // strict byte checks to the base name only. Label values are
643    // enclosed in double quotes by convention (`{stage="worker_queue"}`)
644    // so a global `"` ban would panic on every labelled counter
645    // registration (see the labelled-counter tests). The
646    // real hazard is a base-name that carries `"` / `\` / `\n` / `\r`
647    // / NUL — those would produce unbalanced quotes or split-line
648    // output that Prometheus rejects as invalid text exposition.
649    // Label values carry no render-time escaper: cross-line bytes
650    // (`\n` / `\r` / NUL) anywhere in the labels section are refused
651    // at registration by the loop below; this validator is the only
652    // guard, and its job covers both the base and the labels bytes.
653    let (base, labels) = split_key(key);
654    for byte in base.bytes() {
655        if matches!(byte, b'\n' | b'\r' | b'\\' | b'"' | 0x00) {
656            panic!(
657                "metric base name {base:?} contains a byte (0x{byte:02x}) that would corrupt \
658                 Prometheus text exposition; interpolating attacker-controlled strings \
659                 into metric names is unsafe",
660            );
661        }
662    }
663    // Also refuse the raw cross-line bytes anywhere in the labels
664    // section — Prometheus parses one metric per line, so any bare
665    // `\n`, `\r`, or NUL slips a synthetic line into the scrape.
666    // Backslash and double quote stay legal here because the quotes
667    // are structural to the `l="v"` convention and label sections are
668    // built from code-controlled constants (never attacker input);
669    // labels render verbatim — there is no render-time escaper for
670    // them (only HELP text gets `escape_prom_help`).
671    for byte in labels.bytes() {
672        if matches!(byte, b'\n' | b'\r' | 0x00) {
673            panic!(
674                "metric label section {labels:?} contains a byte (0x{byte:02x}) that would \
675                 corrupt Prometheus text exposition; interpolating attacker-controlled \
676                 strings into metric labels is unsafe",
677            );
678        }
679    }
680}
681
682/// Split `name{l="v"}` into (`name`, `l="v"`); no labels → (`key`, ``).
683fn split_key(key: &str) -> (&str, &str) {
684    match key.find('{') {
685        Some(i) => {
686            let base = key.get(..i).unwrap_or(key);
687            let labels = key.get(i + 1..key.len().saturating_sub(1)).unwrap_or("");
688            (base, labels)
689        }
690        None => (key, ""),
691    }
692}
693
694fn join_labels(labels: &str) -> String {
695    if labels.is_empty() {
696        String::new()
697    } else {
698        format!("{labels},")
699    }
700}
701
702fn labels_suffix(labels: &str) -> String {
703    if labels.is_empty() {
704        String::new()
705    } else {
706        format!("{{{labels}}}")
707    }
708}
709
710#[cfg(test)]
711mod tests {
712    use super::*;
713
714    #[test]
715    fn counter_roundtrip() {
716        let r = Registry::new();
717        let c = r.counter("av_test_total", "test counter");
718        c.inc();
719        c.add(4);
720        assert_eq!(c.get(), 5);
721        let text = r.render();
722        assert!(text.contains("# TYPE av_test_total counter"), "{text}");
723        assert!(text.contains("av_test_total 5"), "{text}");
724    }
725
726    /// Mutation-run hardening: `validate_metric_key` replaced with a no-op
727    /// survived — nothing pinned that a base name carrying a scrape-
728    /// corrupting byte actually panics at registration.
729    #[test]
730    #[should_panic(expected = "corrupt")]
731    fn metric_key_with_newline_panics_at_registration() {
732        let r = Registry::new();
733        let _ = r.counter("av_bad\nname_total", "bad");
734    }
735
736    /// Mutation-run hardening: `escape_prom_help`'s C0/DEL/C1 replacement
737    /// boundary was untested beyond `\n` — ESC (terminal injection via
738    /// piped /metrics), NUL (promtool lint), CR and CSI must all become
739    /// spaces, while printable ASCII from 0x20 up passes through.
740    #[test]
741    fn help_text_scrubs_all_control_bytes_to_spaces() {
742        let scrubbed = escape_prom_help("a\u{1b}[2Jb\u{0}c\rd\u{9b}e");
743        assert_eq!(scrubbed, "a [2Jb c d e");
744        assert_eq!(escape_prom_help(" printable ~"), " printable ~");
745        assert_eq!(escape_prom_help("line\nbreak\\slash"), "line\\nbreak\\\\slash");
746    }
747
748    /// Mutation-run hardening: the `_sum` line converts accumulated
749    /// microseconds to SECONDS via division; a `/`→`*` mutant survived
750    /// because no test pinned the rendered sum value.
751    #[test]
752    fn histogram_sum_renders_in_seconds() {
753        let r = Registry::new();
754        let h = r.histogram("av_test_lat_seconds", "test latency");
755        h.observe_us(1_500_000);
756        h.observe_us(500_000);
757        let text = r.render();
758        assert!(
759            text.contains("av_test_lat_seconds_sum 2\n"),
760            "2_000_000 us must render as 2 seconds: {text}"
761        );
762    }
763
764    #[test]
765    fn gauge_add_sub_set_and_render() {
766        let r = Registry::new();
767        let g = r.gauge("av_test_gauge", "test gauge");
768        g.set(10);
769        g.inc();
770        g.add(3);
771        g.dec();
772        g.sub(2);
773        assert_eq!(g.get(), 11);
774        let text = r.render();
775        assert!(text.contains("# TYPE av_test_gauge gauge"), "{text}");
776        assert!(text.contains("av_test_gauge 11"), "{text}");
777    }
778
779    #[test]
780    fn gauge_saturates_at_zero_not_negative() {
781        let g = Gauge::default();
782        g.set(3);
783        g.sub(100);
784        assert_eq!(
785            g.get(),
786            0,
787            "gauge went negative — dashboards would be misleading and the review's data-plane counters would report bogus values"
788        );
789    }
790
791    #[test]
792    fn gauge_saturates_at_i64_max_not_wrap() {
793        let g = Gauge::default();
794        g.add(u64::MAX);
795        g.add(u64::MAX);
796        // Any positive number is a passing signal (proves no wrap);
797        // the specific ceiling is i64::MAX under the current impl.
798        assert!(g.get() > 0, "gauge wrapped through overflow");
799    }
800
801    #[test]
802    #[should_panic(expected = "already registered as `counter`")]
803    fn gauge_refuses_conflict_with_counter() {
804        let r = Registry::new();
805        let _c = r.counter("av_conflict", "conflict");
806        let _g = r.gauge("av_conflict", "conflict");
807    }
808
809    #[test]
810    #[should_panic(expected = "already registered as `gauge`")]
811    fn counter_refuses_conflict_with_gauge() {
812        let r = Registry::new();
813        let _g = r.gauge("av_conflict2", "conflict");
814        let _c = r.counter("av_conflict2", "conflict");
815    }
816
817    /// HELP text with an embedded newline must be escaped
818    /// so the Prometheus text-format scraper does not interpret the
819    /// second half as a metric sample. Backslash must also be
820    /// escaped per the spec.
821    #[test]
822    #[allow(clippy::expect_used)]
823    fn render_escapes_newlines_and_backslashes_in_help_text() {
824        let r = Registry::new();
825        r.counter("av_dangerous_total", "line1\nline2\\path");
826        let text = r.render();
827        // A rendered HELP line must not contain the raw newline —
828        // any newline in the HELP must appear as literal `\n`.
829        let help_line = text
830            .lines()
831            .find(|line| line.starts_with("# HELP av_dangerous_total"))
832            .expect("HELP line present");
833        assert!(help_line.contains(r"\n"), "help was not escaped: {help_line:?}");
834        assert!(
835            help_line.contains(r"\\"),
836            "backslash was not escaped: {help_line:?}"
837        );
838        // And the "line2" fragment must not be on its own line.
839        assert!(
840            !text.contains("\nline2\\path"),
841            "unescaped fragment leaked into scrape: {text}"
842        );
843    }
844
845    #[test]
846    fn counter_with_labels_renders_base_name_in_type() {
847        let r = Registry::new();
848        let c = r.counter("av_drops_total{stage=\"publish\"}", "drops");
849        c.inc();
850        let text = r.render();
851        assert!(text.contains("# TYPE av_drops_total counter"), "{text}");
852        assert!(text.contains("av_drops_total{stage=\"publish\"} 1"), "{text}");
853    }
854
855    #[test]
856    fn labeled_family_is_declared_once() {
857        let r = Registry::new();
858        r.histogram("av_stage_duration_seconds{stage=\"identity\"}", "stage latency");
859        r.histogram("av_stage_duration_seconds{stage=\"quota\"}", "stage latency");
860        let text = r.render();
861        assert_eq!(
862            text.matches("# HELP av_stage_duration_seconds ").count(),
863            1,
864            "{text}"
865        );
866        assert_eq!(
867            text.matches("# TYPE av_stage_duration_seconds histogram").count(),
868            1,
869            "{text}"
870        );
871    }
872
873    #[test]
874    fn histogram_quantiles() {
875        let h = Histogram::new(DEFAULT_LATENCY_BOUNDS_US);
876        for _ in 0..99 {
877            h.observe_us(80); // ≤ 100µs bucket
878        }
879        h.observe_us(9_000); // ≤ 10ms bucket
880        assert_eq!(h.count(), 100);
881        assert_eq!(h.quantile_us(0.5), 100);
882        assert_eq!(h.quantile_us(0.99), 100);
883        assert_eq!(h.quantile_us(1.0), 10_000);
884    }
885
886    #[test]
887    fn histogram_renders_cumulative_buckets() {
888        let r = Registry::new();
889        let h = r.histogram("av_lat", "latency");
890        h.observe_us(60);
891        h.observe_us(60);
892        h.observe_us(600);
893        let text = r.render();
894        // 50µs bucket: 0, 100µs bucket: 2, ..., 1ms bucket: 3
895        assert!(text.contains("av_lat_bucket{le=\"0.0001\"} 2"), "{text}");
896        assert!(text.contains("av_lat_bucket{le=\"0.001\"} 3"), "{text}");
897        assert!(text.contains("av_lat_bucket{le=\"+Inf\"} 3"), "{text}");
898        assert!(text.contains("av_lat_count 3"), "{text}");
899    }
900
901    /// Prometheus text format requires `_count == cum(bucket{le="+Inf"})`.
902    /// Pre-fix, the render read `h.count()` TWICE — once for the `+Inf`
903    /// bucket, once for the `_count` line — with `sum_us.load(...)` in
904    /// between. A concurrent `observe_us` completing between the two
905    /// loads left the scrape with `_count = N+1` and `+Inf = N`, which
906    /// strict OpenMetrics parsers reject and `histogram_quantile`
907    /// interprets as off-by-one quantile artifacts. Hammer both a
908    /// writer and a reader in parallel and assert the invariant on
909    /// EVERY scrape.
910    /// Prometheus text format enforces two invariants strict
911    /// OpenMetrics parsers reject on:
912    ///
913    ///   (a) `_count == cum(bucket{le="+Inf"})`
914    ///   (b) `bucket{le=X} <= bucket{le=Y}` for X < Y, and
915    ///       `bucket{le=X} <= _count` for all X
916    ///
917    /// R25 closed (a) by snapshotting count once, but LEFT (b) open
918    /// because `observe_us` writes count-then-bucket — a snapshot-
919    /// count-then-read-buckets order could see MORE bucket
920    /// increments than count increments (a concurrent observe that
921    /// completed both fetches between the count snapshot and the
922    /// bucket load). The current fix ALSO clamps cum against the
923    /// snapshot to preserve (b). Hammer a writer + reader in
924    /// parallel and assert BOTH invariants on EVERY scrape.
925    #[test]
926    #[allow(clippy::unwrap_used, clippy::expect_used)]
927    fn histogram_scrape_preserves_all_invariants_under_concurrent_observe() {
928        use std::sync::atomic::{AtomicBool, Ordering};
929        use std::sync::Arc;
930        use std::thread;
931
932        let r = Arc::new(Registry::new());
933        // Diverse bucket bounds so a busy writer spreads observations
934        // across several buckets — increases the chance of catching a
935        // mid-loop race between the bucket loop and the count snapshot.
936        let h = r.histogram_with_bounds(
937            "av_race",
938            "concurrent scrape probe",
939            &[10, 100, 1_000, 10_000, 100_000],
940        );
941        let stop = Arc::new(AtomicBool::new(false));
942
943        let writer = {
944            let h = h.clone();
945            let stop = Arc::clone(&stop);
946            thread::spawn(move || {
947                // Rotate through values that hit different buckets.
948                let mut i: u64 = 0;
949                while !stop.load(Ordering::Relaxed) {
950                    let us = (i % 5 + 1) * 50; // 50, 100, 150, 200, 250
951                    h.observe_us(us);
952                    i = i.wrapping_add(1);
953                }
954            })
955        };
956
957        let start = std::time::Instant::now();
958        let mut scrapes = 0usize;
959        while start.elapsed() < std::time::Duration::from_millis(300) {
960            let text = r.render();
961            let mut buckets: Vec<(f64, u64)> = Vec::new();
962            let mut count: Option<u64> = None;
963            for line in text.lines() {
964                if line.starts_with("av_race_bucket") {
965                    let le_start = line.find("le=\"").unwrap() + 4;
966                    let le_end = le_start + line[le_start..].find('"').unwrap();
967                    let le_str = &line[le_start..le_end];
968                    let val: u64 = line.rsplit(' ').next().and_then(|s| s.parse().ok()).unwrap();
969                    let le_f = if le_str == "+Inf" {
970                        f64::INFINITY
971                    } else {
972                        le_str.parse().unwrap()
973                    };
974                    buckets.push((le_f, val));
975                } else if line.starts_with("av_race_count") {
976                    count = Some(line.rsplit(' ').next().and_then(|s| s.parse().ok()).unwrap());
977                }
978            }
979            let count = count.expect("_count line missing");
980            buckets.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
981            let mut prev = 0u64;
982            for (le, val) in &buckets {
983                // Invariant (b1): buckets monotone non-decreasing.
984                assert!(
985                    *val >= prev,
986                    "monotonicity violated at le={le}: {val} < prev {prev} on scrape {scrapes}"
987                );
988                // Invariant (b2): every bucket <= _count.
989                assert!(
990                    *val <= count,
991                    "bucket at le={le} = {val} > _count = {count} on scrape {scrapes}"
992                );
993                prev = *val;
994            }
995            // Invariant (a): +Inf bucket == _count.
996            let inf = buckets
997                .iter()
998                .find(|(le, _)| le.is_infinite())
999                .expect("+Inf bucket missing")
1000                .1;
1001            assert_eq!(
1002                inf, count,
1003                "invariant (a) violated: +Inf = {inf} != _count = {count} on scrape {scrapes}"
1004            );
1005            scrapes = scrapes.saturating_add(1);
1006        }
1007        stop.store(true, Ordering::Relaxed);
1008        writer.join().unwrap();
1009        assert!(scrapes > 100, "should have taken > 100 scrapes; got {scrapes}");
1010    }
1011
1012    #[test]
1013    fn same_key_returns_same_metric() {
1014        let r = Registry::new();
1015        let a = r.counter("x_total", "x");
1016        let b = r.counter("x_total", "x");
1017        a.inc();
1018        assert_eq!(b.get(), 1);
1019    }
1020
1021    #[test]
1022    fn labeled_histogram_renders_labels_on_every_series_line() {
1023        // Catches `split_key`, `join_labels`, and `labels_suffix` stubs.
1024        // A labeled histogram must emit label pairs on _bucket, _sum, and
1025        // _count — the label prefix must land BEFORE the `le=` in buckets
1026        // and the entire `{labels}` block must land after _sum/_count.
1027        let r = Registry::new();
1028        let h = r.histogram("av_lat{route=\"chat\"}", "lat");
1029        h.observe_us(60);
1030        let text = r.render();
1031        assert!(
1032            text.contains("av_lat_bucket{route=\"chat\",le=\"0.0001\"} 1"),
1033            "join_labels lost the route label: {text}"
1034        );
1035        assert!(
1036            text.contains("av_lat_sum{route=\"chat\"} "),
1037            "labels_suffix lost the route label on _sum: {text}"
1038        );
1039        assert!(
1040            text.contains("av_lat_count{route=\"chat\"} 1"),
1041            "labels_suffix lost the route label on _count: {text}"
1042        );
1043    }
1044
1045    /// Vicious registry bug: `counter()` used to silently
1046    /// insert a fresh Counter over an existing Histogram at the same key,
1047    /// detaching the histogram from the registry and losing every prior
1048    /// observation. A metric name registered as one type must never be
1049    /// silently repurposed as the other — panic loudly instead.
1050    #[test]
1051    #[should_panic(expected = "kind conflict")]
1052    fn counter_over_existing_histogram_panics_instead_of_overwriting() {
1053        let r = Registry::new();
1054        r.histogram("av_metric", "help");
1055        r.counter("av_metric", "help");
1056    }
1057
1058    #[test]
1059    #[should_panic(expected = "kind conflict")]
1060    fn histogram_over_existing_counter_panics_instead_of_overwriting() {
1061        let r = Registry::new();
1062        r.counter("av_metric", "help");
1063        r.histogram("av_metric", "help");
1064    }
1065
1066    /// A histogram's text exposition generates `<base>_bucket`,
1067    /// `<base>_sum` and `<base>_count` series. Registering a counter
1068    /// under one of those names produced TWO samples named
1069    /// `<base>_count` under conflicting TYPE headers — Prometheus
1070    /// rejects the whole scrape as invalid exposition. Both
1071    /// registration orders must fail fast.
1072    #[test]
1073    #[should_panic(expected = "kind conflict")]
1074    fn counter_named_like_histogram_count_series_panics() {
1075        let r = Registry::new();
1076        r.histogram("av_lat_seconds", "help");
1077        r.counter("av_lat_seconds_count", "help");
1078    }
1079
1080    #[test]
1081    #[should_panic(expected = "kind conflict")]
1082    fn histogram_whose_generated_series_collide_with_counter_panics() {
1083        let r = Registry::new();
1084        r.counter("av_lat_seconds_sum", "help");
1085        r.histogram("av_lat_seconds", "help");
1086    }
1087
1088    /// The reservation must not break legitimate reuse: the same
1089    /// histogram re-registered (other label sets) and unrelated
1090    /// suffix-shaped names with no histogram sibling stay valid.
1091    #[test]
1092    fn histogram_reservation_permits_reregistration_and_unrelated_suffix_names() {
1093        let r = Registry::new();
1094        r.histogram("av_lat_seconds{route=\"a\"}", "help");
1095        r.histogram("av_lat_seconds{route=\"b\"}", "help");
1096        r.counter("av_unrelated_count", "help");
1097        assert!(r.render().contains("av_unrelated_count"));
1098    }
1099
1100    /// Cross-label type collision on the same base name: `foo{a="1"}`
1101    /// registered as counter, `foo{a="2"}` registered as histogram.
1102    /// Prometheus TYPE header would be ambiguous — reject at
1103    /// registration.
1104    #[test]
1105    #[should_panic(expected = "kind conflict")]
1106    fn different_labels_same_base_name_type_collision_panics() {
1107        let r = Registry::new();
1108        r.counter("av_metric{shard=\"a\"}", "help");
1109        r.histogram("av_metric{shard=\"b\"}", "help");
1110    }
1111
1112    /// Same-type re-registration must still work (idempotent) — the panic
1113    /// guard only fires on genuine type conflicts.
1114    #[test]
1115    fn same_type_re_registration_returns_existing_arc() {
1116        let r = Registry::new();
1117        let c1 = r.counter("av_metric", "help");
1118        let c2 = r.counter("av_metric", "help");
1119        c1.inc();
1120        assert_eq!(c2.get(), 1, "must return the same underlying counter");
1121
1122        let h1 = r.histogram("av_other", "help");
1123        let h2 = r.histogram("av_other", "help");
1124        h1.observe_us(50);
1125        assert_eq!(h2.count(), 1, "must return the same underlying histogram");
1126    }
1127
1128    /// HELP text escaper substitutes DEL (0x7F) and every
1129    /// C1 control (0x80..=0x9F) with a literal space. C0 controls
1130    /// were handled first; C1 was left through and CSI (0x9B) is a
1131    /// valid single-byte ANSI escape prefix under 8-bit terminal
1132    /// emulation. A future counter registration with operator-
1133    /// influenced HELP text (charter name / model name / SSE error
1134    /// snippet) that happened to include a C1 byte could otherwise
1135    /// inject terminal-escape sequences through `curl /metrics | less`.
1136    #[test]
1137    fn escape_prom_help_scrubs_del_and_c1_controls() {
1138        let dangerous = "help\u{7f}with\u{9b}csi\u{80}c1\u{9f}end";
1139        let out = escape_prom_help(dangerous);
1140        for c in ['\u{7f}', '\u{9b}', '\u{80}', '\u{9f}'] {
1141            assert!(!out.contains(c), "escape_prom_help left {c:?} in {out:?}");
1142        }
1143        // The safe chars are preserved.
1144        assert!(out.contains("help"));
1145        assert!(out.contains("csi"));
1146        assert!(out.contains("end"));
1147    }
1148}
1149
1150#[cfg(test)]
1151mod histogram_boundary_tests {
1152    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
1153
1154    use super::*;
1155
1156    /// Mutation-run hardening: the overflow-bucket branch of
1157    /// `quantile_us` and the seconds conversion in `_sum` rendering were
1158    /// unpinned. A sample past the last bound must surface as u64::MAX
1159    /// (never silently under-report as the last bound), and the rendered
1160    /// sum must be the microsecond total divided by exactly 1e6.
1161    #[test]
1162    fn quantile_overflow_reports_max_and_sum_renders_in_seconds() {
1163        let r = Registry::new();
1164        let h = r.histogram_with_bounds("av_overflow_probe", "probe", &[10, 100]);
1165        h.observe_us(1_000_000); // beyond the last bound: overflow bucket
1166        assert_eq!(h.quantile_us(0.99), u64::MAX);
1167        h.observe_us(2_000_000);
1168        let text = r.render();
1169        assert!(
1170            text.contains("av_overflow_probe_sum 3"),
1171            "sum must be 3 seconds (3_000_000 us / 1e6), got:\n{text}"
1172        );
1173    }
1174}