use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
#[derive(Debug, Clone)]
pub enum MetricValue {
Counter(u64),
Gauge(f64),
Histogram(Vec<f64>),
}
impl MetricValue {
pub fn as_counter(&self) -> u64 {
match self {
MetricValue::Counter(v) => *v,
_ => 0,
}
}
pub fn as_gauge(&self) -> f64 {
match self {
MetricValue::Gauge(v) => *v,
_ => 0.0,
}
}
pub fn as_histogram(&self) -> &[f64] {
match self {
MetricValue::Histogram(v) => v,
_ => &[],
}
}
}
#[derive(Debug, Default)]
pub struct MetricsStorage {
metrics: RwLock<HashMap<String, MetricValue>>,
}
impl MetricsStorage {
pub fn new() -> Self {
Self::default()
}
pub fn increment_counter(&self, name: &str, delta: u64) {
let mut metrics = self.metrics.write();
let entry = metrics
.entry(name.to_string())
.or_insert(MetricValue::Counter(0));
if let MetricValue::Counter(v) = entry {
*v += delta;
}
}
pub fn set_gauge(&self, name: &str, value: f64) {
let mut metrics = self.metrics.write();
metrics.insert(name.to_string(), MetricValue::Gauge(value));
}
pub fn record_histogram(&self, name: &str, value: f64) {
let mut metrics = self.metrics.write();
let entry = metrics
.entry(name.to_string())
.or_insert(MetricValue::Histogram(Vec::new()));
if let MetricValue::Histogram(v) = entry {
v.push(value);
}
}
pub fn get(&self, name: &str) -> Option<MetricValue> {
self.metrics.read().get(name).cloned()
}
pub fn list_names(&self) -> Vec<String> {
self.metrics.read().keys().cloned().collect()
}
}
#[derive(Debug, Clone)]
pub struct Counter {
name: String,
storage: Arc<MetricsStorage>,
}
impl Counter {
pub fn new(name: impl Into<String>, storage: Arc<MetricsStorage>) -> Self {
Self {
name: name.into(),
storage,
}
}
#[allow(clippy::let_unit_value)]
pub fn increment(&self, delta: u64) {
#[cfg(feature = "prometheus")]
{
let _ = metrics::counter!(self.name.clone()).increment(delta);
}
#[cfg(not(feature = "prometheus"))]
{
self.storage.increment_counter(&self.name, delta);
}
}
pub fn get(&self) -> u64 {
self.storage
.get(&self.name)
.map(|v| v.as_counter())
.unwrap_or(0)
}
}
#[derive(Debug, Clone)]
pub struct Histogram {
name: String,
storage: Arc<MetricsStorage>,
}
impl Histogram {
pub fn new(name: impl Into<String>, storage: Arc<MetricsStorage>) -> Self {
Self {
name: name.into(),
storage,
}
}
#[allow(clippy::let_unit_value)]
pub fn record(&self, value: f64) {
#[cfg(feature = "prometheus")]
{
let _ = metrics::histogram!(self.name.clone()).record(value);
}
#[cfg(not(feature = "prometheus"))]
{
self.storage.record_histogram(&self.name, value);
}
}
pub fn get_values(&self) -> Vec<f64> {
self.storage
.get(&self.name)
.map(|v| v.as_histogram().to_vec())
.unwrap_or_default()
}
}
#[derive(Debug, Clone)]
pub struct Gauge {
name: String,
storage: Arc<MetricsStorage>,
}
impl Gauge {
pub fn new(name: impl Into<String>, storage: Arc<MetricsStorage>) -> Self {
Self {
name: name.into(),
storage,
}
}
#[allow(clippy::let_unit_value)]
pub fn set(&self, value: f64) {
#[cfg(feature = "prometheus")]
{
let _ = metrics::gauge!(self.name.clone()).set(value);
}
#[cfg(not(feature = "prometheus"))]
{
self.storage.set_gauge(&self.name, value);
}
}
pub fn get(&self) -> f64 {
self.storage
.get(&self.name)
.map(|v| v.as_gauge())
.unwrap_or(0.0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_counter_increment() {
let storage = Arc::new(MetricsStorage::new());
let counter = Counter::new("test_counter", storage);
counter.increment(5);
assert_eq!(counter.get(), 5);
counter.increment(3);
assert_eq!(counter.get(), 8);
}
#[test]
fn test_gauge_set() {
let storage = Arc::new(MetricsStorage::new());
let gauge = Gauge::new("test_gauge", storage);
gauge.set(42.0);
assert_eq!(gauge.get(), 42.0);
gauge.set(100.0);
assert_eq!(gauge.get(), 100.0);
}
#[test]
fn test_histogram_record() {
let storage = Arc::new(MetricsStorage::new());
let histogram = Histogram::new("test_histogram", storage);
histogram.record(1.0);
histogram.record(2.0);
histogram.record(3.0);
let values = histogram.get_values();
assert_eq!(values, vec![1.0, 2.0, 3.0]);
}
#[test]
fn test_metrics_storage_list_names() {
let storage = Arc::new(MetricsStorage::new());
let counter = Counter::new("counter1", Arc::clone(&storage));
let gauge = Gauge::new("gauge1", Arc::clone(&storage));
counter.increment(1);
gauge.set(1.0);
let names = storage.list_names();
assert_eq!(names.len(), 2);
assert!(names.contains(&"counter1".to_string()));
assert!(names.contains(&"gauge1".to_string()));
}
}