use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
pub struct Summary {
name: String,
help: String,
quantiles: Vec<f64>,
samples: Arc<RwLock<Vec<f64>>>,
sum: Arc<RwLock<f64>>,
count: Arc<RwLock<u64>>,
}
impl Summary {
pub fn new(name: impl Into<String>, help: impl Into<String>, quantiles: Vec<f64>) -> Self {
Self {
name: name.into(),
help: help.into(),
quantiles,
samples: Arc::new(RwLock::new(Vec::new())),
sum: Arc::new(RwLock::new(0.0)),
count: Arc::new(RwLock::new(0)),
}
}
pub fn observe(&self, value: f64) {
let mut samples = self.samples.write();
let pos = samples.partition_point(|&v| v < value);
samples.insert(pos, value);
let mut sum = self.sum.write();
*sum += value;
let mut count = self.count.write();
*count += 1;
}
pub fn quantile(&self, q: f64) -> Option<f64> {
if !(0.0..=1.0).contains(&q) {
return None;
}
let samples = self.samples.read();
if samples.is_empty() {
return None;
}
let n = samples.len();
let rank = ((q * n as f64).ceil() as usize).max(1).min(n);
Some(samples[rank - 1])
}
pub fn quantiles(&self) -> Vec<(f64, Option<f64>)> {
self.quantiles
.iter()
.map(|&q| (q, self.quantile(q)))
.collect()
}
pub fn count(&self) -> u64 {
*self.count.read()
}
pub fn sum(&self) -> f64 {
*self.sum.read()
}
pub fn name(&self) -> &str {
&self.name
}
pub fn render(&self) -> String {
let samples = self.samples.read();
let sum = *self.sum.read();
let count = *self.count.read();
let mut output = String::new();
output.push_str(&format!("# HELP {} {}\n", self.name, self.help));
output.push_str(&format!("# TYPE {} summary\n", self.name));
for &q in &self.quantiles {
let value = if samples.is_empty() {
0.0
} else {
let n = samples.len();
let rank = ((q * n as f64).ceil() as usize).max(1).min(n);
samples[rank - 1]
};
output.push_str(&format!("{}{{quantile=\"{}\"}} {}\n", self.name, q, value));
}
output.push_str(&format!("{}_sum {}\n", self.name, sum));
output.push_str(&format!("{}_count {}\n", self.name, count));
output
}
pub fn reset(&self) {
let mut samples = self.samples.write();
samples.clear();
*self.sum.write() = 0.0;
*self.count.write() = 0;
}
}
pub struct LabeledHistogram {
name: String,
help: String,
buckets: Vec<f64>,
series: RwLock<HashMap<String, LabeledSeries>>,
}
#[derive(Debug, Clone)]
struct LabeledSeries {
labels: Vec<(String, String)>,
counts: Vec<u64>,
sum: f64,
count: u64,
}
impl LabeledSeries {
fn new(labels: Vec<(String, String)>, bucket_count: usize) -> Self {
Self {
labels,
counts: vec![0; bucket_count],
sum: 0.0,
count: 0,
}
}
fn label_key(labels: &[(String, String)]) -> String {
labels
.iter()
.map(|(k, v)| format!("{}=\"{}\"", k, v.replace('"', "\\\"")))
.collect::<Vec<_>>()
.join(",")
}
}
impl LabeledHistogram {
pub fn new(name: impl Into<String>, help: impl Into<String>, mut buckets: Vec<f64>) -> Self {
buckets.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if !buckets.contains(&f64::INFINITY) {
buckets.push(f64::INFINITY);
}
Self {
name: name.into(),
help: help.into(),
buckets,
series: RwLock::new(HashMap::new()),
}
}
pub fn observe(&self, labels: &HashMap<String, String>, value: f64) {
let mut sorted_labels: Vec<(String, String)> =
labels.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
sorted_labels.sort_by(|a, b| a.0.cmp(&b.0));
let key = LabeledSeries::label_key(&sorted_labels);
let mut series = self.series.write();
let entry = series
.entry(key)
.or_insert_with(|| LabeledSeries::new(sorted_labels.clone(), self.buckets.len()));
for (i, bucket) in self.buckets.iter().enumerate() {
if value <= *bucket {
entry.counts[i] += 1;
}
}
entry.sum += value;
entry.count += 1;
}
pub fn count(&self, labels: &HashMap<String, String>) -> u64 {
let mut sorted_labels: Vec<(String, String)> =
labels.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
sorted_labels.sort_by(|a, b| a.0.cmp(&b.0));
let key = LabeledSeries::label_key(&sorted_labels);
self.series.read().get(&key).map(|s| s.count).unwrap_or(0)
}
pub fn label_combination_count(&self) -> usize {
self.series.read().len()
}
pub fn render(&self) -> String {
let series = self.series.read();
let mut output = String::new();
output.push_str(&format!("# HELP {} {}\n", self.name, self.help));
output.push_str(&format!("# TYPE {} histogram\n", self.name));
for s in series.values() {
let label_str = LabeledSeries::label_key(&s.labels);
for (i, bucket) in self.buckets.iter().enumerate() {
if *bucket == f64::INFINITY {
output.push_str(&format!(
"{}_bucket{{{},le=\"+Inf\"}} {}\n",
self.name, label_str, s.counts[i]
));
} else {
output.push_str(&format!(
"{}_bucket{{{},le=\"{}\"}} {}\n",
self.name, label_str, bucket, s.counts[i]
));
}
}
output.push_str(&format!("{}_sum{{{}}} {}\n", self.name, label_str, s.sum));
output.push_str(&format!(
"{}_count{{{}}} {}\n",
self.name, label_str, s.count
));
}
output
}
}
#[derive(Debug, Clone)]
pub struct PushgatewayConfig {
pub endpoint: String,
pub job: String,
pub instance: Option<String>,
}
impl Default for PushgatewayConfig {
fn default() -> Self {
Self {
endpoint: "http://localhost:9091".to_string(),
job: "sz-orm".to_string(),
instance: None,
}
}
}
pub struct PushgatewayExporter {
config: PushgatewayConfig,
pushed: RwLock<Vec<PushSnapshot>>,
}
#[derive(Debug, Clone)]
pub struct PushSnapshot {
pub timestamp_ms: i64,
pub metrics_text: String,
pub job: String,
pub instance: Option<String>,
}
impl PushgatewayExporter {
pub fn new(config: PushgatewayConfig) -> Self {
Self {
config,
pushed: RwLock::new(Vec::new()),
}
}
pub fn push(&self, metrics_text: impl Into<String>) -> Result<(), String> {
let text = metrics_text.into();
let snapshot = PushSnapshot {
timestamp_ms: current_timestamp_ms(),
metrics_text: text.clone(),
job: self.config.job.clone(),
instance: self.config.instance.clone(),
};
self.pushed.write().push(snapshot);
#[cfg(feature = "push-gateway")]
{
self.push_http(&text)
}
#[cfg(not(feature = "push-gateway"))]
Ok(())
}
#[cfg(feature = "push-gateway")]
fn push_http(&self, text: &str) -> Result<(), String> {
let mut url = format!(
"{}/metrics/job/{}",
self.config.endpoint.trim_end_matches('/'),
url_encode(&self.config.job)
);
if let Some(ref instance) = self.config.instance {
url.push_str(&format!("/instance/{}", url_encode(instance)));
}
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()
.map_err(|e| format!("reqwest client build failed: {}", e))?;
let resp = client
.put(&url)
.header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
.body(text.to_string())
.send()
.map_err(|e| format!("push to {} failed: {}", url, e))?;
let status = resp.status();
if status.is_success() {
Ok(())
} else {
let body = resp.text().unwrap_or_default();
Err(format!(
"push to {} returned non-2xx status {}: {}",
url, status, body
))
}
}
pub fn push_from_registry(&self, registry: &crate::MetricsRegistry) -> Result<(), String> {
let text = registry.render();
self.push(text)
}
pub fn snapshots(&self) -> Vec<PushSnapshot> {
self.pushed.read().clone()
}
pub fn push_count(&self) -> usize {
self.pushed.read().len()
}
pub fn clear(&self) {
self.pushed.write().clear();
}
pub fn config(&self) -> &PushgatewayConfig {
&self.config
}
}
fn current_timestamp_ms() -> i64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as i64
}
#[allow(dead_code)]
fn url_encode(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for byte in s.bytes() {
match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
out.push(byte as char);
}
_ => {
out.push_str(&format!("%{:02X}", byte));
}
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_url_encode_alphanumeric() {
assert_eq!(url_encode("job1"), "job1");
assert_eq!(url_encode("my-job"), "my-job");
assert_eq!(url_encode("job.test"), "job.test");
assert_eq!(url_encode("job_test"), "job_test");
assert_eq!(url_encode("job~test"), "job~test");
}
#[test]
fn test_url_encode_special_chars() {
assert_eq!(url_encode("my job"), "my%20job");
assert_eq!(url_encode("a/b"), "a%2Fb");
assert_eq!(url_encode("任务"), "%E4%BB%BB%E5%8A%A1");
}
#[test]
fn test_url_encode_empty() {
assert_eq!(url_encode(""), "");
}
#[test]
fn test_pushgateway_memory_mode_records_snapshot() {
let exporter = PushgatewayExporter::new(PushgatewayConfig::default());
let result = exporter.push("# HELP test_metric\n");
assert!(result.is_ok());
assert_eq!(exporter.push_count(), 1);
let snap = &exporter.snapshots()[0];
assert_eq!(snap.metrics_text, "# HELP test_metric\n");
assert_eq!(snap.job, "sz-orm");
}
#[test]
fn test_summary_new_empty() {
let s = Summary::new("latency", "latency summary", vec![0.5, 0.9, 0.99]);
assert_eq!(s.count(), 0);
assert_eq!(s.sum(), 0.0);
assert!(s.quantile(0.5).is_none());
}
#[test]
fn test_summary_observe_single() {
let s = Summary::new("latency", "help", vec![0.5]);
s.observe(1.5);
assert_eq!(s.count(), 1);
assert!((s.sum() - 1.5).abs() < 1e-9);
assert!((s.quantile(0.5).unwrap() - 1.5).abs() < 1e-9);
}
#[test]
fn test_summary_observe_multiple_p50() {
let s = Summary::new("latency", "help", vec![0.5]);
for v in [1.0, 2.0, 3.0, 4.0, 5.0] {
s.observe(v);
}
assert!((s.quantile(0.5).unwrap() - 3.0).abs() < 1e-9);
}
#[test]
fn test_summary_observe_multiple_p99() {
let s = Summary::new("latency", "help", vec![0.99]);
for v in [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 100.0] {
s.observe(v);
}
assert!((s.quantile(0.99).unwrap() - 100.0).abs() < 1e-9);
}
#[test]
fn test_summary_quantile_out_of_range() {
let s = Summary::new("latency", "help", vec![0.5]);
s.observe(1.0);
assert!(s.quantile(-0.1).is_none());
assert!(s.quantile(1.1).is_none());
}
#[test]
fn test_summary_quantile_empty() {
let s = Summary::new("latency", "help", vec![0.5]);
assert!(s.quantile(0.5).is_none());
}
#[test]
fn test_summary_quantile_p0_and_p1() {
let s = Summary::new("latency", "help", vec![]);
for v in [10.0, 20.0, 30.0] {
s.observe(v);
}
assert!((s.quantile(0.0).unwrap() - 10.0).abs() < 1e-9);
assert!((s.quantile(1.0).unwrap() - 30.0).abs() < 1e-9);
}
#[test]
fn test_summary_quantiles_all() {
let s = Summary::new("latency", "help", vec![0.5, 0.9, 0.99]);
for v in [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0] {
s.observe(v);
}
let qs = s.quantiles();
assert_eq!(qs.len(), 3);
assert!(qs.iter().all(|(_, v)| v.is_some()));
let qmap: Vec<(f64, f64)> = qs
.iter()
.filter_map(|(q, v)| v.map(|val| (*q, val)))
.collect();
let lookup = |target: f64| -> f64 {
qmap.iter()
.find(|(q, _)| (*q - target).abs() < 1e-9)
.map(|(_, v)| *v)
.unwrap_or(f64::NAN)
};
let p50 = lookup(0.5);
assert!(
(5.0..=6.0).contains(&p50),
"p50 应在 5-6 之间,实际: {}",
p50
);
let p90 = lookup(0.9);
assert!(
(9.0..=10.0).contains(&p90),
"p90 应在 9-10 之间,实际: {}",
p90
);
let p99 = lookup(0.99);
assert!(
(9.0..=10.0).contains(&p99),
"p99 应在 9-10 之间,实际: {}",
p99
);
}
#[test]
fn test_summary_unsorted_input_stays_sorted() {
let s = Summary::new("latency", "help", vec![0.5]);
s.observe(50.0);
s.observe(10.0);
s.observe(30.0);
assert!((s.quantile(0.5).unwrap() - 30.0).abs() < 1e-9);
}
#[test]
fn test_summary_render_contains_type() {
let s = Summary::new("latency", "latency help", vec![0.5, 0.99]);
s.observe(1.0);
let output = s.render();
assert!(output.contains("# HELP latency latency help"));
assert!(output.contains("# TYPE latency summary"));
assert!(output.contains("latency{quantile=\"0.5\"}"));
assert!(output.contains("latency{quantile=\"0.99\"}"));
assert!(output.contains("latency_sum"));
assert!(output.contains("latency_count"));
}
#[test]
fn test_summary_render_empty_shows_zero() {
let s = Summary::new("latency", "help", vec![0.5]);
let output = s.render();
assert!(output.contains("latency{quantile=\"0.5\"} 0"));
assert!(output.contains("latency_count 0"));
}
#[test]
fn test_summary_reset() {
let s = Summary::new("latency", "help", vec![0.5]);
s.observe(1.0);
s.observe(2.0);
assert_eq!(s.count(), 2);
s.reset();
assert_eq!(s.count(), 0);
assert!((s.sum() - 0.0).abs() < 1e-9);
assert!(s.quantile(0.5).is_none());
}
#[test]
fn test_summary_name() {
let s = Summary::new("my_metric", "help", vec![0.5]);
assert_eq!(s.name(), "my_metric");
}
#[test]
fn test_labeled_histogram_new() {
let h = LabeledHistogram::new("requests", "help", vec![0.1, 0.5, 1.0]);
assert_eq!(h.label_combination_count(), 0);
}
#[test]
fn test_labeled_histogram_observe_single_label() {
let h = LabeledHistogram::new("requests", "help", vec![0.1, 0.5, 1.0]);
let mut labels = HashMap::new();
labels.insert("method".to_string(), "GET".to_string());
h.observe(&labels, 0.3);
assert_eq!(h.count(&labels), 1);
assert_eq!(h.label_combination_count(), 1);
}
#[test]
fn test_labeled_histogram_observe_multiple_labels() {
let h = LabeledHistogram::new("requests", "help", vec![0.1, 0.5, 1.0]);
let mut get_labels = HashMap::new();
get_labels.insert("method".to_string(), "GET".to_string());
let mut post_labels = HashMap::new();
post_labels.insert("method".to_string(), "POST".to_string());
h.observe(&get_labels, 0.1);
h.observe(&get_labels, 0.2);
h.observe(&post_labels, 0.5);
assert_eq!(h.count(&get_labels), 2);
assert_eq!(h.count(&post_labels), 1);
assert_eq!(h.label_combination_count(), 2);
}
#[test]
fn test_labeled_histogram_label_order_independent() {
let h = LabeledHistogram::new("requests", "help", vec![0.1, 1.0]);
let mut labels1 = HashMap::new();
labels1.insert("a".to_string(), "1".to_string());
labels1.insert("b".to_string(), "2".to_string());
let mut labels2 = HashMap::new();
labels2.insert("b".to_string(), "2".to_string());
labels2.insert("a".to_string(), "1".to_string());
h.observe(&labels1, 0.5);
assert_eq!(h.count(&labels2), 1);
assert_eq!(h.label_combination_count(), 1);
}
#[test]
fn test_labeled_histogram_count_missing_labels() {
let h = LabeledHistogram::new("requests", "help", vec![0.1, 1.0]);
let labels = HashMap::new();
assert_eq!(h.count(&labels), 0);
}
#[test]
fn test_labeled_histogram_render_contains_labels() {
let h = LabeledHistogram::new("requests", "request help", vec![0.1, 1.0]);
let mut labels = HashMap::new();
labels.insert("method".to_string(), "GET".to_string());
h.observe(&labels, 0.05);
let output = h.render();
assert!(output.contains("# HELP requests request help"));
assert!(output.contains("# TYPE requests histogram"));
assert!(output.contains("method=\"GET\""));
assert!(output.contains("requests_count"));
assert!(output.contains("requests_sum"));
}
#[test]
fn test_labeled_histogram_render_inf_bucket() {
let h = LabeledHistogram::new("req", "help", vec![0.1]);
let labels = HashMap::new();
h.observe(&labels, 0.05);
h.observe(&labels, 5.0);
let output = h.render();
assert!(output.contains("le=\"+Inf\""));
}
#[test]
fn test_pushgateway_config_default() {
let config = PushgatewayConfig::default();
assert_eq!(config.endpoint, "http://localhost:9091");
assert_eq!(config.job, "sz-orm");
assert!(config.instance.is_none());
}
#[test]
fn test_pushgateway_exporter_new() {
let exporter = PushgatewayExporter::new(PushgatewayConfig::default());
assert_eq!(exporter.push_count(), 0);
assert!(exporter.snapshots().is_empty());
}
#[test]
fn test_pushgateway_push_text() {
let exporter = PushgatewayExporter::new(PushgatewayConfig::default());
exporter.push("metric1 1\n").unwrap();
exporter.push("metric2 2\n").unwrap();
assert_eq!(exporter.push_count(), 2);
let snaps = exporter.snapshots();
assert_eq!(snaps.len(), 2);
assert_eq!(snaps[0].metrics_text, "metric1 1\n");
assert_eq!(snaps[1].metrics_text, "metric2 2\n");
}
#[test]
fn test_pushgateway_push_from_registry() {
let registry = crate::MetricsRegistry::new();
let counter = registry.register_counter("test_total", "test");
counter.inc();
let exporter = PushgatewayExporter::new(PushgatewayConfig::default());
exporter.push_from_registry(®istry).unwrap();
assert_eq!(exporter.push_count(), 1);
let snap = &exporter.snapshots()[0];
assert!(snap.metrics_text.contains("test_total"));
}
#[test]
fn test_pushgateway_snapshot_has_metadata() {
let config = PushgatewayConfig {
endpoint: "http://push:9091".to_string(),
job: "myjob".to_string(),
instance: Some("inst1".to_string()),
};
let exporter = PushgatewayExporter::new(config);
exporter.push("m 1\n").unwrap();
let snap = &exporter.snapshots()[0];
assert_eq!(snap.job, "myjob");
assert_eq!(snap.instance, Some("inst1".to_string()));
assert!(snap.timestamp_ms > 0);
}
#[test]
fn test_pushgateway_clear() {
let exporter = PushgatewayExporter::new(PushgatewayConfig::default());
exporter.push("m 1\n").unwrap();
assert_eq!(exporter.push_count(), 1);
exporter.clear();
assert_eq!(exporter.push_count(), 0);
}
#[test]
fn test_pushgateway_config_access() {
let config = PushgatewayConfig {
job: "custom".to_string(),
..Default::default()
};
let exporter = PushgatewayExporter::new(config);
assert_eq!(exporter.config().job, "custom");
}
}