396 lines
14 KiB
Rust
396 lines
14 KiB
Rust
|
//! Forward the process's Prometheus metrics to the connected server over
|
||
|
|
//! the WebSocket transport (`metrics.donate`).
|
||
|
|
//!
|
||
|
|
//! A [`MetricDonationReporter`] periodically snapshots the default
|
||
|
|
//! Prometheus registry via [`prometheus::gather`], converts the
|
||
|
|
//! `MetricFamily` set to native OTLP metrics (Counter→Sum, Gauge→Gauge,
|
||
|
|
//! Histogram→Histogram, labels preserved, cumulative temporality), and
|
||
|
|
//! pumps the batch over the shared [`crate::donate_pump`]. Because it
|
||
|
|
//! gathers the whole registry, every current and future metric is
|
||
|
|
//! exported with zero per-metric wiring. Metrics are **process-aggregate**
|
||
|
|
//! — [`ToolServer::donate_metrics`] requires no bound session.
|
||
|
|
|
||
|
|
use std::sync::{Arc, LazyLock};
|
||
|
|
use std::time::Duration;
|
||
|
|
|
||
|
|
use arc_swap::ArcSwapOption;
|
||
|
|
use base64::Engine as _;
|
||
|
|
use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest;
|
||
|
|
use opentelemetry_proto::tonic::common::v1::KeyValue;
|
||
|
|
use opentelemetry_proto::tonic::metrics::v1::{
|
||
|
|
AggregationTemporality, Gauge, Histogram, HistogramDataPoint, Metric, NumberDataPoint,
|
||
|
|
ResourceMetrics, ScopeMetrics, Sum, metric, number_data_point,
|
||
|
|
};
|
||
|
|
use opentelemetry_proto::tonic::resource::v1::Resource;
|
||
|
|
use prometheus::proto::{MetricFamily, MetricType};
|
||
|
|
use prost::Message as _;
|
||
|
|
use tokio::sync::mpsc;
|
||
|
|
use tokio_util::sync::CancellationToken;
|
||
|
|
use xai_tool_protocol::{MAX_DONATION_BYTES, MAX_METRICS_PER_DONATION};
|
||
|
|
|
||
|
|
use crate::donate_pump::{
|
||
|
|
PENDING_FLUSHES, PumpMsg, make_resource, now_unix_nanos, run_pump, string_kv,
|
||
|
|
};
|
||
|
|
use crate::server::ToolServer;
|
||
|
|
|
||
|
|
/// How often the reporter snapshots the registry. The server re-stamps
|
||
|
|
/// attribution; cumulative temporality means missed ticks only delay
|
||
|
|
/// freshness, never lose monotonic state.
|
||
|
|
const DEFAULT_GATHER_INTERVAL: Duration = Duration::from_secs(60);
|
||
|
|
|
||
|
|
fn labels_to_kv(labels: &[prometheus::proto::LabelPair]) -> Vec<KeyValue> {
|
||
|
|
labels
|
||
|
|
.iter()
|
||
|
|
.map(|l| string_kv(l.name(), l.value().to_owned()))
|
||
|
|
.collect()
|
||
|
|
}
|
||
|
|
|
||
|
|
fn number_point(metric: &prometheus::proto::Metric, value: f64, now: u64) -> NumberDataPoint {
|
||
|
|
NumberDataPoint {
|
||
|
|
attributes: labels_to_kv(metric.get_label()),
|
||
|
|
time_unix_nano: now,
|
||
|
|
value: Some(number_data_point::Value::AsDouble(value)),
|
||
|
|
..Default::default()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Prometheus histogram buckets are **cumulative** (`le` counts); OTLP
|
||
|
|
/// wants per-bucket counts plus an implicit `+Inf` bucket, so the
|
||
|
|
/// cumulative counts are differenced here.
|
||
|
|
fn histogram_point(metric: &prometheus::proto::Metric, now: u64) -> HistogramDataPoint {
|
||
|
|
let hist = metric.get_histogram();
|
||
|
|
let mut bucket_counts = Vec::new();
|
||
|
|
let mut explicit_bounds = Vec::new();
|
||
|
|
let mut prev = 0u64;
|
||
|
|
for bucket in hist.get_bucket() {
|
||
|
|
let cumulative = bucket.cumulative_count();
|
||
|
|
bucket_counts.push(cumulative.saturating_sub(prev));
|
||
|
|
explicit_bounds.push(bucket.upper_bound());
|
||
|
|
prev = cumulative;
|
||
|
|
}
|
||
|
|
let total = hist.get_sample_count();
|
||
|
|
bucket_counts.push(total.saturating_sub(prev));
|
||
|
|
HistogramDataPoint {
|
||
|
|
attributes: labels_to_kv(metric.get_label()),
|
||
|
|
time_unix_nano: now,
|
||
|
|
count: total,
|
||
|
|
sum: Some(hist.get_sample_sum()),
|
||
|
|
bucket_counts,
|
||
|
|
explicit_bounds,
|
||
|
|
..Default::default()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Convert a gathered `MetricFamily` set to OTLP metrics. Summaries and
|
||
|
|
/// untyped families are skipped defensively (none registered today).
|
||
|
|
fn convert_families(families: &[MetricFamily]) -> Vec<Metric> {
|
||
|
|
let now = now_unix_nanos();
|
||
|
|
let cumulative = AggregationTemporality::Cumulative as i32;
|
||
|
|
let mut out = Vec::new();
|
||
|
|
for family in families {
|
||
|
|
let name = family.name().to_owned();
|
||
|
|
let data = match family.get_field_type() {
|
||
|
|
MetricType::COUNTER => metric::Data::Sum(Sum {
|
||
|
|
data_points: family
|
||
|
|
.get_metric()
|
||
|
|
.iter()
|
||
|
|
.map(|m| number_point(m, m.get_counter().value(), now))
|
||
|
|
.collect(),
|
||
|
|
aggregation_temporality: cumulative,
|
||
|
|
is_monotonic: true,
|
||
|
|
}),
|
||
|
|
MetricType::GAUGE => metric::Data::Gauge(Gauge {
|
||
|
|
data_points: family
|
||
|
|
.get_metric()
|
||
|
|
.iter()
|
||
|
|
.map(|m| number_point(m, m.get_gauge().value(), now))
|
||
|
|
.collect(),
|
||
|
|
}),
|
||
|
|
MetricType::HISTOGRAM => metric::Data::Histogram(Histogram {
|
||
|
|
data_points: family
|
||
|
|
.get_metric()
|
||
|
|
.iter()
|
||
|
|
.map(|m| histogram_point(m, now))
|
||
|
|
.collect(),
|
||
|
|
aggregation_temporality: cumulative,
|
||
|
|
}),
|
||
|
|
MetricType::SUMMARY | MetricType::UNTYPED => continue,
|
||
|
|
};
|
||
|
|
out.push(Metric {
|
||
|
|
name,
|
||
|
|
data: Some(data),
|
||
|
|
..Default::default()
|
||
|
|
});
|
||
|
|
}
|
||
|
|
out
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Encodes batches of OTLP metrics onto the pump channel. Chunks at
|
||
|
|
/// [`MAX_METRICS_PER_DONATION`], drops payloads over
|
||
|
|
/// [`MAX_DONATION_BYTES`], and never blocks.
|
||
|
|
#[derive(Clone)]
|
||
|
|
struct MetricExporter {
|
||
|
|
tx: mpsc::Sender<PumpMsg>,
|
||
|
|
resource: Resource,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl MetricExporter {
|
||
|
|
fn export(&self, mut metrics: Vec<Metric>) {
|
||
|
|
while !metrics.is_empty() {
|
||
|
|
let chunk = if metrics.len() > MAX_METRICS_PER_DONATION {
|
||
|
|
let rest = metrics.split_off(MAX_METRICS_PER_DONATION);
|
||
|
|
std::mem::replace(&mut metrics, rest)
|
||
|
|
} else {
|
||
|
|
std::mem::take(&mut metrics)
|
||
|
|
};
|
||
|
|
let request = ExportMetricsServiceRequest {
|
||
|
|
resource_metrics: vec![ResourceMetrics {
|
||
|
|
resource: Some(self.resource.clone()),
|
||
|
|
scope_metrics: vec![ScopeMetrics {
|
||
|
|
metrics: chunk,
|
||
|
|
..Default::default()
|
||
|
|
}],
|
||
|
|
schema_url: String::new(),
|
||
|
|
}],
|
||
|
|
};
|
||
|
|
let bytes = request.encode_to_vec();
|
||
|
|
if bytes.len() > MAX_DONATION_BYTES {
|
||
|
|
tracing::debug!(
|
||
|
|
len = bytes.len(),
|
||
|
|
"dropping oversized metric donation payload"
|
||
|
|
);
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
|
||
|
|
if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
|
||
|
|
tracing::debug!("metric donation queue full; dropping metric batch");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn gather_and_send(&self) {
|
||
|
|
let metrics = convert_families(&prometheus::gather());
|
||
|
|
if metrics.is_empty() {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
self.export(metrics);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Process-global handle to the active exporter so [`gather_and_send`]
|
||
|
|
/// can drive a final teardown gather without a reference.
|
||
|
|
static ACTIVE_METRIC_EXPORTER: LazyLock<ArcSwapOption<MetricExporter>> =
|
||
|
|
LazyLock::new(ArcSwapOption::empty);
|
||
|
|
|
||
|
|
/// Final registry gather onto the active metric pump. Called from
|
||
|
|
/// `ToolServer` teardown before the pump drain so a crash-y shutdown
|
||
|
|
/// captures the latest values.
|
||
|
|
pub(crate) fn gather_and_send() {
|
||
|
|
if let Some(exporter) = ACTIVE_METRIC_EXPORTER.load_full() {
|
||
|
|
exporter.gather_and_send();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Drop the process-global exporter on teardown so its pump `Sender` is released
|
||
|
|
/// and the metric pump can wind down. Called from `flush_donations_inner` after
|
||
|
|
/// the final [`gather_and_send`] (and alongside clearing the stored pump
|
||
|
|
/// senders), so a dropped `ToolServer` doesn't leak the pump task.
|
||
|
|
pub(crate) fn clear_active_exporter() {
|
||
|
|
ACTIVE_METRIC_EXPORTER.store(None);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Periodic registry gatherer spawned by
|
||
|
|
/// [`ToolServer::metric_donation_reporter`]. Internal: constructed and run
|
||
|
|
/// only by `metric_donation_reporter`; not part of the crate's public API.
|
||
|
|
pub(crate) struct MetricDonationReporter {
|
||
|
|
exporter: MetricExporter,
|
||
|
|
interval: Duration,
|
||
|
|
shutdown: CancellationToken,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl MetricDonationReporter {
|
||
|
|
async fn run(self) {
|
||
|
|
let mut ticker = tokio::time::interval(self.interval);
|
||
|
|
loop {
|
||
|
|
tokio::select! {
|
||
|
|
_ = ticker.tick() => self.exporter.gather_and_send(),
|
||
|
|
// Stop on teardown so this task (and the pump `tx` clone it
|
||
|
|
// holds) doesn't outlive `ToolServer::shutdown` and keep
|
||
|
|
// gathering/sending forever.
|
||
|
|
_ = self.shutdown.cancelled() => break,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Shutdown fence: drains queued metric donations before the connection
|
||
|
|
/// closes.
|
||
|
|
pub struct MetricDonationPump {
|
||
|
|
tx: mpsc::Sender<PumpMsg>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl MetricDonationPump {
|
||
|
|
/// Resolves once every payload queued before this call has had a
|
||
|
|
/// send attempt.
|
||
|
|
pub async fn drain(&self) {
|
||
|
|
crate::donate_pump::drain_via(&self.tx).await;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ToolServer {
|
||
|
|
/// Post-connect entry point: spawn the metric donation pump (wiring
|
||
|
|
/// [`ToolServer::donate_metrics`]) plus the periodic registry
|
||
|
|
/// gatherer, and return a drain handle. Activates on server presence
|
||
|
|
/// (no env flag). `service_name` must be server-allowlisted.
|
||
|
|
pub fn metric_donation_reporter(&self, service_name: impl Into<String>) -> MetricDonationPump {
|
||
|
|
let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
|
||
|
|
let server = self.downgrade();
|
||
|
|
tokio::spawn(run_pump(rx, move |payload: String| {
|
||
|
|
let server = server.clone();
|
||
|
|
async move {
|
||
|
|
let Some(server) = server.upgrade() else {
|
||
|
|
return (false, payload);
|
||
|
|
};
|
||
|
|
let ok = server.donate_metrics(&payload).await.is_ok();
|
||
|
|
(ok, payload)
|
||
|
|
}
|
||
|
|
}));
|
||
|
|
self.set_metric_donation_pump(tx.clone());
|
||
|
|
|
||
|
|
let exporter = MetricExporter {
|
||
|
|
tx: tx.clone(),
|
||
|
|
resource: make_resource(service_name.into()),
|
||
|
|
};
|
||
|
|
ACTIVE_METRIC_EXPORTER.store(Some(Arc::new(exporter.clone())));
|
||
|
|
tokio::spawn(
|
||
|
|
MetricDonationReporter {
|
||
|
|
exporter,
|
||
|
|
interval: DEFAULT_GATHER_INTERVAL,
|
||
|
|
shutdown: self.shutdown_token(),
|
||
|
|
}
|
||
|
|
.run(),
|
||
|
|
);
|
||
|
|
|
||
|
|
MetricDonationPump { tx }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use opentelemetry_proto::tonic::common::v1::any_value;
|
||
|
|
use prometheus::{Histogram, HistogramOpts, IntCounterVec, IntGauge, Opts, Registry};
|
||
|
|
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
fn decode(payload: String) -> ExportMetricsServiceRequest {
|
||
|
|
let bytes = base64::engine::general_purpose::STANDARD
|
||
|
|
.decode(payload)
|
||
|
|
.expect("payload must be base64");
|
||
|
|
ExportMetricsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP")
|
||
|
|
}
|
||
|
|
|
||
|
|
fn label_map(attrs: &[KeyValue]) -> std::collections::HashMap<String, String> {
|
||
|
|
attrs
|
||
|
|
.iter()
|
||
|
|
.filter_map(|kv| match kv.value.as_ref().and_then(|v| v.value.clone()) {
|
||
|
|
Some(any_value::Value::StringValue(s)) => Some((kv.key.clone(), s)),
|
||
|
|
_ => None,
|
||
|
|
})
|
||
|
|
.collect()
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn converts_counter_gauge_histogram_with_labels() {
|
||
|
|
let registry = Registry::new();
|
||
|
|
|
||
|
|
let counter =
|
||
|
|
IntCounterVec::new(Opts::new("grok_test_total", "help"), &["reason"]).unwrap();
|
||
|
|
registry.register(Box::new(counter.clone())).unwrap();
|
||
|
|
counter.with_label_values(&["zdr"]).inc_by(5);
|
||
|
|
|
||
|
|
let gauge = IntGauge::new("grok_test_pending", "help").unwrap();
|
||
|
|
registry.register(Box::new(gauge.clone())).unwrap();
|
||
|
|
gauge.set(7);
|
||
|
|
|
||
|
|
let hist = Histogram::with_opts(
|
||
|
|
HistogramOpts::new("grok_test_seconds", "help").buckets(vec![0.5, 1.0]),
|
||
|
|
)
|
||
|
|
.unwrap();
|
||
|
|
registry.register(Box::new(hist.clone())).unwrap();
|
||
|
|
hist.observe(0.25);
|
||
|
|
hist.observe(0.75);
|
||
|
|
hist.observe(5.0);
|
||
|
|
|
||
|
|
let metrics = convert_families(®istry.gather());
|
||
|
|
let by_name: std::collections::HashMap<_, _> =
|
||
|
|
metrics.iter().map(|m| (m.name.clone(), m)).collect();
|
||
|
|
|
||
|
|
// Counter -> Sum (monotonic, cumulative), label preserved.
|
||
|
|
let metric::Data::Sum(sum) = by_name["grok_test_total"].data.as_ref().unwrap() else {
|
||
|
|
panic!("counter must convert to Sum");
|
||
|
|
};
|
||
|
|
assert!(sum.is_monotonic);
|
||
|
|
assert_eq!(
|
||
|
|
sum.aggregation_temporality,
|
||
|
|
AggregationTemporality::Cumulative as i32
|
||
|
|
);
|
||
|
|
let dp = &sum.data_points[0];
|
||
|
|
assert_eq!(dp.value, Some(number_data_point::Value::AsDouble(5.0)));
|
||
|
|
assert_eq!(
|
||
|
|
label_map(&dp.attributes).get("reason").map(String::as_str),
|
||
|
|
Some("zdr")
|
||
|
|
);
|
||
|
|
|
||
|
|
// Gauge -> Gauge.
|
||
|
|
let metric::Data::Gauge(g) = by_name["grok_test_pending"].data.as_ref().unwrap() else {
|
||
|
|
panic!("gauge must convert to Gauge");
|
||
|
|
};
|
||
|
|
assert_eq!(
|
||
|
|
g.data_points[0].value,
|
||
|
|
Some(number_data_point::Value::AsDouble(7.0))
|
||
|
|
);
|
||
|
|
|
||
|
|
// Histogram -> Histogram with cumulative buckets differenced and
|
||
|
|
// a +Inf bucket appended.
|
||
|
|
let metric::Data::Histogram(h) = by_name["grok_test_seconds"].data.as_ref().unwrap() else {
|
||
|
|
panic!("histogram must convert to Histogram");
|
||
|
|
};
|
||
|
|
assert_eq!(
|
||
|
|
h.aggregation_temporality,
|
||
|
|
AggregationTemporality::Cumulative as i32
|
||
|
|
);
|
||
|
|
let hdp = &h.data_points[0];
|
||
|
|
assert_eq!(hdp.count, 3);
|
||
|
|
assert_eq!(hdp.sum, Some(6.0));
|
||
|
|
assert_eq!(hdp.explicit_bounds, vec![0.5, 1.0]);
|
||
|
|
// (<=0.5): 0.25 -> 1 ; (0.5,1.0]: 0.75 -> 1 ; (+Inf): 5.0 -> 1
|
||
|
|
assert_eq!(hdp.bucket_counts, vec![1, 1, 1]);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn exporter_chunks_at_max_metrics_per_donation() {
|
||
|
|
let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
|
||
|
|
let exporter = MetricExporter {
|
||
|
|
tx,
|
||
|
|
resource: make_resource("test-service".to_owned()),
|
||
|
|
};
|
||
|
|
let metrics = vec![Metric::default(); MAX_METRICS_PER_DONATION + 1];
|
||
|
|
exporter.export(metrics);
|
||
|
|
|
||
|
|
let mut payloads = 0;
|
||
|
|
let mut total = 0;
|
||
|
|
while let Ok(PumpMsg::Payload(p)) = rx.try_recv() {
|
||
|
|
payloads += 1;
|
||
|
|
total += decode(p).resource_metrics[0].scope_metrics[0].metrics.len();
|
||
|
|
}
|
||
|
|
assert_eq!(payloads, 2, "one full chunk + remainder");
|
||
|
|
assert_eq!(total, MAX_METRICS_PER_DONATION + 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn summary_and_untyped_families_are_skipped() {
|
||
|
|
// An empty registry gathers nothing; convert yields nothing.
|
||
|
|
let registry = Registry::new();
|
||
|
|
assert!(convert_families(®istry.gather()).is_empty());
|
||
|
|
}
|
||
|
|
}
|