use std::sync::Arc;
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use crate::Error;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MetricType {
Counter,
Gauge,
Timing,
}
pub trait Metric: Send + Sync {
fn increment(&self, value: u64);
fn decrement(&self, value: u64);
fn set(&self, value: u64);
fn timing(&self, duration: Duration);
fn metric_type(&self) -> MetricType;
fn name(&self) -> &str;
}
pub trait MetricCollector: Send + Sync {
fn counter(&self, name: &str) -> Arc<dyn Metric>;
fn gauge(&self, name: &str) -> Arc<dyn Metric>;
fn timer(&self, name: &str) -> Arc<dyn Metric>;
fn with_tags(
&self,
tags: std::collections::HashMap<String, String>,
) -> Box<dyn MetricCollector>;
fn close(&self) -> Result<(), Error>;
}
pub struct NoopMetric {
name: String,
metric_type: MetricType,
}
impl NoopMetric {
pub fn new(name: &str, metric_type: MetricType) -> Self {
Self {
name: name.to_string(),
metric_type,
}
}
}
impl Metric for NoopMetric {
fn increment(&self, _value: u64) {}
fn decrement(&self, _value: u64) {}
fn set(&self, _value: u64) {}
fn timing(&self, _duration: Duration) {}
fn metric_type(&self) -> MetricType {
self.metric_type.clone()
}
fn name(&self) -> &str {
&self.name
}
}
pub struct NoopCollector {}
impl NoopCollector {
pub fn new() -> Self {
Self {}
}
}
impl MetricCollector for NoopCollector {
fn counter(&self, name: &str) -> Arc<dyn Metric> {
Arc::new(NoopMetric::new(name, MetricType::Counter))
}
fn gauge(&self, name: &str) -> Arc<dyn Metric> {
Arc::new(NoopMetric::new(name, MetricType::Gauge))
}
fn timer(&self, name: &str) -> Arc<dyn Metric> {
Arc::new(NoopMetric::new(name, MetricType::Timing))
}
fn with_tags(
&self,
_tags: std::collections::HashMap<String, String>,
) -> Box<dyn MetricCollector> {
Box::new(Self::new())
}
fn close(&self) -> Result<(), Error> {
Ok(())
}
}
pub struct PrometheusCollector {
prefix: String,
}
impl PrometheusCollector {
pub fn new(prefix: &str) -> Result<Self, Error> {
Ok(Self {
prefix: prefix.to_string(),
})
}
}
impl MetricCollector for PrometheusCollector {
fn counter(&self, name: &str) -> Arc<dyn Metric> {
Arc::new(NoopMetric::new(name, MetricType::Counter))
}
fn gauge(&self, name: &str) -> Arc<dyn Metric> {
Arc::new(NoopMetric::new(name, MetricType::Gauge))
}
fn timer(&self, name: &str) -> Arc<dyn Metric> {
Arc::new(NoopMetric::new(name, MetricType::Timing))
}
fn with_tags(
&self,
_tags: std::collections::HashMap<String, String>,
) -> Box<dyn MetricCollector> {
Box::new(Self {
prefix: self.prefix.clone(),
})
}
fn close(&self) -> Result<(), Error> {
Ok(())
}
}
pub struct TimingHelper {
metric: Arc<dyn Metric>,
start: Instant,
}
impl TimingHelper {
pub fn new(metric: Arc<dyn Metric>) -> Self {
Self {
metric,
start: Instant::now(),
}
}
pub fn done(self) {
let duration = self.start.elapsed();
self.metric.timing(duration);
}
}
pub fn create_metrics(
config: &crate::config::MetricsConfig,
) -> Result<Box<dyn MetricCollector>, Error> {
if !config.enabled {
return Ok(Box::new(NoopCollector::new()));
}
let prefix = config.prefix.as_deref().unwrap_or("").to_string();
match config.type_name.as_str() {
"prometheus" => {
let collector = PrometheusCollector::new(&prefix)?;
if let Some(tags) = &config.tags {
Ok(collector.with_tags(tags.clone()))
} else {
Ok(Box::new(collector))
}
}
"statsd" => {
Ok(Box::new(NoopCollector::new()))
}
"none" | "noop" => Ok(Box::new(NoopCollector::new())),
_ => Err(Error::Config(format!(
"不支持的指标类型: {}",
config.type_name
))),
}
}