use std::sync::{atomic::AtomicI64, Arc};
use serde::{
ser::{SerializeSeq, SerializeStruct},
Serialize, Serializer,
};
use super::prometheus::Encoder;
#[derive(Debug, Clone)]
pub struct Gauge {
pub metric: Arc<AtomicI64>,
pub labels: Vec<(&'static str, &'static str)>,
}
#[derive(Debug, Clone)]
pub struct GaugeVec {
pub metrics: Vec<Gauge>,
}
const GAUGE_MODIFIER: i64 = 100_000;
impl GaugeVec {
pub fn new(labels: &[&[(&'static str, &'static str)]]) -> Self {
Self {
metrics: labels
.iter()
.map(|label_list| Gauge::new_with_labels(label_list))
.collect(),
}
}
pub fn with_labels(&self, labels: &[(&str, &str)]) -> Option<&Gauge> {
'cnt: for gauge in &self.metrics {
if gauge.labels.len() != labels.len() {
continue;
}
for ((name1, value1), (name2, value2)) in gauge.labels.iter().zip(labels.iter()) {
if name1 != name2 || value1 != value2 {
continue 'cnt;
}
}
return Some(gauge);
}
None
}
pub fn reset(&self) {
for gauge in &self.metrics {
gauge.reset();
}
}
}
impl Default for Gauge {
fn default() -> Self {
Self::new()
}
}
impl Gauge {
pub fn new() -> Self {
Self {
metric: Arc::new(AtomicI64::new(0)),
labels: Vec::new(),
}
}
pub fn new_with_labels(labels: &[(&'static str, &'static str)]) -> Self {
let mut map = Vec::new();
for (name, value) in labels {
map.push((*name, *value));
}
Self {
metric: Arc::new(AtomicI64::new(0)),
labels: map,
}
}
pub fn inc(&self) {
self.inc_by(1.0);
}
pub fn dec(&self) {
self.dec_by(1.0);
}
pub fn inc_by(&self, v: f64) {
self.metric
.fetch_add(normalize_f64(v), std::sync::atomic::Ordering::SeqCst);
}
pub fn set(&self, v: f64) {
self.metric
.store(normalize_f64(v), std::sync::atomic::Ordering::SeqCst);
}
pub fn dec_by(&self, v: f64) {
self.metric
.fetch_sub(normalize_f64(v), std::sync::atomic::Ordering::SeqCst);
}
pub fn get(&self) -> f64 {
normalize_i64(self.metric.load(std::sync::atomic::Ordering::Relaxed))
}
pub fn reset(&self) {
self.metric.store(0, std::sync::atomic::Ordering::SeqCst);
}
}
fn normalize_f64(numb: f64) -> i64 {
(numb * (GAUGE_MODIFIER as f64)) as i64
}
fn normalize_i64(numb: i64) -> f64 {
numb as f64 / GAUGE_MODIFIER as f64
}
impl Encoder for GaugeVec {
fn encode<W: std::fmt::Write>(
&self,
f: &mut W,
name: &str,
description: &str,
help: bool,
) -> Result<(), std::fmt::Error> {
if help {
f.write_str("# HELP ")?;
f.write_str(name)?;
f.write_str(" ")?;
f.write_str(description)?;
f.write_str("\n")?;
}
f.write_str("# TYPE ")?;
f.write_str(name)?;
f.write_str(" gauge\n")?;
for counter in &self.metrics {
counter.encode(f, name, description, help)?;
}
Ok(())
}
}
impl Encoder for Gauge {
fn encode<W: std::fmt::Write>(
&self,
f: &mut W,
name: &str,
_description: &str,
_help: bool,
) -> Result<(), std::fmt::Error> {
f.write_str(name)?;
f.write_str("{")?;
let mut i = 0;
for (name, value) in &self.labels {
i += 1;
if value.is_empty() {
continue;
}
f.write_str(name)?;
f.write_str("=\"")?;
f.write_str(value)?;
f.write_str("\"")?;
if i != self.labels.len() {
f.write_str(",")?;
}
}
f.write_str("} ")?;
f.write_fmt(format_args!("{}", self.get()))?;
f.write_str("\n")?;
Ok(())
}
}
impl Serialize for GaugeVec {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut state = serializer.serialize_seq(Some(self.metrics.len()))?;
for metric in &self.metrics {
state.serialize_element(metric)?;
}
state.end()
}
}
impl Serialize for Gauge {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut state = serializer.serialize_struct("Gauge", 2)?;
state.serialize_field("labels", &self.labels)?;
state.serialize_field("value", &self.get())?;
state.end()
}
}
#[test]
fn gauge_metric_should_work() {
let metric = GaugeVec::new(&[
&[("name", "Paco"), ("v1", "1")],
&[("name", "Pepe"), ("v1", "2")],
]);
let gauge = metric
.with_labels(&[("name", "Paco"), ("v1", "1")])
.unwrap();
gauge.inc_by(222.0);
assert_eq!(222.0, gauge.get());
gauge.reset();
assert_eq!(0.0, gauge.get());
gauge.inc_by(222.0);
metric.reset();
assert_eq!(0.0, gauge.get());
let gauge = metric
.with_labels(&[("name", "Pepe"), ("v1", "2")])
.unwrap();
gauge.inc_by(222.0);
assert_eq!(222.0, gauge.get());
gauge.reset();
assert_eq!(0.0, gauge.get());
gauge.inc_by(222.0);
metric.reset();
assert_eq!(0.0, gauge.get());
assert!(metric
.with_labels(&[("name", "Paco"), ("v1", "2")])
.is_none());
gauge.set(123.0);
assert_eq!(123.0, gauge.get());
}
#[test]
fn gauge_should_be_encoded_in_prometheus() {
let counter = GaugeVec::new(&[
&[("name", "Nombre"), ("v1", "1")],
&[("name", "Nombre"), ("v1", "2")],
]);
let mut st = String::with_capacity(1_000_000);
counter
.encode(&mut st, "simple_gauge", "Simple Gauge metric", true)
.unwrap();
assert_eq!(
r#"# HELP simple_gauge Simple Gauge metric
# TYPE simple_gauge gauge
simple_gauge{name="Nombre",v1="1"} 0
simple_gauge{name="Nombre",v1="2"} 0
"#,
st
);
}