use std::sync::Arc;
use tokio::sync::RwLock;
use std::collections::HashMap;
use std::time::Duration;
#[derive(Clone)]
pub struct Metrics {
counters: Arc<RwLock<HashMap<String, Counter>>>,
gauges: Arc<RwLock<HashMap<String, Gauge>>>,
histograms: Arc<RwLock<HashMap<String, Histogram>>>,
}
impl Metrics {
pub fn new() -> Self {
Self {
counters: Arc::new(RwLock::new(HashMap::new())),
gauges: Arc::new(RwLock::new(HashMap::new())),
histograms: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn counter(&self, name: &str) -> Counter {
let mut counters = self.counters.write().await;
counters
.entry(name.to_string())
.or_insert_with(Counter::new)
.clone()
}
pub async fn gauge(&self, name: &str) -> Gauge {
let mut gauges = self.gauges.write().await;
gauges
.entry(name.to_string())
.or_insert_with(Gauge::new)
.clone()
}
pub async fn histogram(&self, name: &str) -> Histogram {
let mut histograms = self.histograms.write().await;
histograms
.entry(name.to_string())
.or_insert_with(Histogram::new)
.clone()
}
pub async fn inc_http_requests(&self, method: &str, path: &str, status: u16) {
let counter = self.counter("http_requests_total").await;
counter.inc_by(1.0, &[("method", method), ("path", path), ("status", &status.to_string())]).await;
}
pub async fn observe_http_duration(&self, method: &str, path: &str, duration: Duration) {
let histogram = self.histogram("http_request_duration_seconds").await;
histogram.observe(duration.as_secs_f64(), &[("method", method), ("path", path)]).await;
}
pub async fn set_active_connections(&self, count: usize) {
let gauge = self.gauge("active_connections").await;
gauge.set(count as f64, &[]).await;
}
pub async fn export(&self) -> String {
let mut output = String::new();
let counters = self.counters.read().await;
for (name, counter) in counters.iter() {
output.push_str(&format!("# HELP {} {}\n", name, name));
output.push_str(&format!("# TYPE {} counter\n", name));
for (labels, value) in counter.get_values().await {
output.push_str(&format!("{}{} {}\n", name, labels, value));
}
}
let gauges = self.gauges.read().await;
for (name, gauge) in gauges.iter() {
output.push_str(&format!("# HELP {} {}\n", name, name));
output.push_str(&format!("# TYPE {} gauge\n", name));
for (labels, value) in gauge.get_values().await {
output.push_str(&format!("{}{} {}\n", name, labels, value));
}
}
let histograms = self.histograms.read().await;
for (name, histogram) in histograms.iter() {
output.push_str(&format!("# HELP {} {}\n", name, name));
output.push_str(&format!("# TYPE {} histogram\n", name));
for (labels, bucket_data) in histogram.get_buckets().await {
let bucket_name = format!("{}_bucket", name);
for data in bucket_data {
output.push_str(&format!("{}{{{}le=\"{}\"}} {}\n", bucket_name, labels, data.le, data.count));
}
}
output.push_str(&format!("{}_sum{} {}\n", name, "", histogram.get_sum().await));
output.push_str(&format!("{}_count{} {}\n", name, "", histogram.get_count().await));
}
output
}
}
impl Default for Metrics {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct Counter {
values: Arc<RwLock<HashMap<String, f64>>>,
labels: Arc<RwLock<HashMap<String, Vec<(String, String)>>>>,
}
impl Counter {
pub fn new() -> Self {
Self {
values: Arc::new(RwLock::new(HashMap::new())),
labels: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn inc(&self, labels: &[(&str, &str)]) {
let key = Self::labels_to_key(labels);
let mut values = self.values.write().await;
*values.entry(key.clone()).or_insert(0.0) += 1.0;
let mut all_labels = self.labels.write().await;
all_labels.insert(key, labels.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect());
}
pub async fn inc_by(&self, value: f64, labels: &[(&str, &str)]) {
let key = Self::labels_to_key(labels);
let mut values = self.values.write().await;
*values.entry(key.clone()).or_insert(0.0) += value;
let mut all_labels = self.labels.write().await;
all_labels.insert(key, labels.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect());
}
pub async fn get_values(&self) -> Vec<(String, f64)> {
let values = self.values.read().await;
let labels = self.labels.read().await;
values
.iter()
.map(|(k, v)| {
let label_str = labels.get(k)
.map(|l| l.iter().map(|(k, v)| format!("{}=\"{}\"", k, v)).collect::<Vec<_>>().join(","))
.unwrap_or_default();
(format!("{{{}}}", label_str), *v)
})
.collect()
}
fn labels_to_key(labels: &[(&str, &str)]) -> String {
labels.iter().map(|(k, v)| format!("{}={}", k, v)).collect::<Vec<_>>().join(",")
}
}
impl Default for Counter {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct Gauge {
values: Arc<RwLock<HashMap<String, f64>>>,
labels: Arc<RwLock<HashMap<String, Vec<(String, String)>>>>,
}
impl Gauge {
pub fn new() -> Self {
Self {
values: Arc::new(RwLock::new(HashMap::new())),
labels: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn set(&self, value: f64, labels: &[(&str, &str)]) {
let key = Self::labels_to_key(labels);
let mut values = self.values.write().await;
values.insert(key.clone(), value);
let mut all_labels = self.labels.write().await;
all_labels.insert(key, labels.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect());
}
pub async fn inc(&self, labels: &[(&str, &str)]) {
let key = Self::labels_to_key(labels);
let mut values = self.values.write().await;
*values.entry(key.clone()).or_insert(0.0) += 1.0;
}
pub async fn dec(&self, labels: &[(&str, &str)]) {
let key = Self::labels_to_key(labels);
let mut values = self.values.write().await;
*values.entry(key.clone()).or_insert(0.0) -= 1.0;
}
pub async fn get_values(&self) -> Vec<(String, f64)> {
let values = self.values.read().await;
let labels = self.labels.read().await;
values
.iter()
.map(|(k, v)| {
let label_str = labels.get(k)
.map(|l| l.iter().map(|(k, v)| format!("{}=\"{}\"", k, v)).collect::<Vec<_>>().join(","))
.unwrap_or_default();
(format!("{{{}}}", label_str), *v)
})
.collect()
}
fn labels_to_key(labels: &[(&str, &str)]) -> String {
labels.iter().map(|(k, v)| format!("{}={}", k, v)).collect::<Vec<_>>().join(",")
}
}
impl Default for Gauge {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct Histogram {
buckets: Arc<RwLock<HashMap<String, Vec<BucketData>>>>,
sum: Arc<RwLock<HashMap<String, f64>>>,
count: Arc<RwLock<HashMap<String, u64>>>,
}
#[derive(Debug, Clone)]
pub struct BucketData {
le: f64,
count: u64,
}
impl Histogram {
pub fn new() -> Self {
Self {
buckets: Arc::new(RwLock::new(HashMap::new())),
sum: Arc::new(RwLock::new(HashMap::new())),
count: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn observe(&self, value: f64, labels: &[(&str, &str)]) {
let key = Self::labels_to_key(labels);
let mut sum = self.sum.write().await;
*sum.entry(key.clone()).or_insert(0.0) += value;
let mut count = self.count.write().await;
*count.entry(key.clone()).or_insert(0) += 1;
let bucket_bounds = vec![0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0];
let mut buckets = self.buckets.write().await;
for bound in bucket_bounds {
if value <= bound {
if let Some(bucket) = buckets.entry(key.clone()).or_insert_with(|| {
vec![BucketData { le: 0.005, count: 0 }; 11]
}).iter_mut().find(|b| b.le == bound) {
bucket.count += 1;
}
}
}
}
pub async fn get_buckets(&self) -> Vec<(String, Vec<BucketData>)> {
let buckets = self.buckets.read().await;
buckets.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
}
pub async fn get_sum(&self) -> f64 {
let sum = self.sum.read().await;
sum.values().sum()
}
pub async fn get_count(&self) -> u64 {
let count = self.count.read().await;
count.values().sum()
}
fn labels_to_key(labels: &[(&str, &str)]) -> String {
labels.iter().map(|(k, v)| format!("{}={}", k, v)).collect::<Vec<_>>().join(",")
}
}
impl Default for Histogram {
fn default() -> Self {
Self::new()
}
}