use std::collections::HashMap;
use std::hash::Hasher;
use std::sync::Arc;
use fnv::FnvHasher;
use parking_lot::RwLock;
use crate::desc::{Desc, Describer};
use crate::errors::{Error, Result};
use crate::metrics::{Collector, Metric};
use crate::proto::{MetricFamily, MetricType};
pub trait MetricVecBuilder: Send + Sync + Clone {
type M: Metric;
type P: Describer + Sync + Send + Clone;
fn build(&self, _: &Self::P, _: &[&str]) -> Result<Self::M>;
}
#[derive(Debug)]
pub(crate) struct MetricVecCore<T: MetricVecBuilder> {
pub children: RwLock<HashMap<u64, T::M>>,
pub desc: Desc,
pub metric_type: MetricType,
pub new_metric: T,
pub opts: T::P,
}
impl<T: MetricVecBuilder> MetricVecCore<T> {
pub fn collect(&self) -> MetricFamily {
let mut m = MetricFamily::default();
m.set_name(self.desc.fq_name.clone());
m.set_help(self.desc.help.clone());
m.set_field_type(self.metric_type);
let children = self.children.read();
let mut metrics = Vec::with_capacity(children.len());
for child in children.values() {
metrics.push(child.metric());
}
m.set_metric(from_vec!(metrics));
m
}
pub fn get_metric_with_label_values(&self, vals: &[&str]) -> Result<T::M> {
let h = self.hash_label_values(vals)?;
if let Some(metric) = self.children.read().get(&h).cloned() {
return Ok(metric);
}
self.get_or_create_metric(h, vals)
}
pub fn get_metric_with(&self, labels: &HashMap<&str, &str>) -> Result<T::M> {
let h = self.hash_labels(labels)?;
if let Some(metric) = self.children.read().get(&h).cloned() {
return Ok(metric);
}
let vals = self.get_label_values(labels)?;
self.get_or_create_metric(h, &vals)
}
pub fn delete_label_values(&self, vals: &[&str]) -> Result<()> {
let h = self.hash_label_values(vals)?;
let mut children = self.children.write();
if children.remove(&h).is_none() {
return Err(Error::Msg(format!("missing label values {:?}", vals)));
}
Ok(())
}
pub fn delete(&self, labels: &HashMap<&str, &str>) -> Result<()> {
let h = self.hash_labels(labels)?;
let mut children = self.children.write();
if children.remove(&h).is_none() {
return Err(Error::Msg(format!("missing labels {:?}", labels)));
}
Ok(())
}
pub fn reset(&self) {
self.children.write().clear();
}
pub(crate) fn hash_label_values(&self, vals: &[&str]) -> Result<u64> {
if vals.len() != self.desc.variable_labels.len() {
return Err(Error::InconsistentCardinality {
expect: self.desc.variable_labels.len(),
got: vals.len(),
});
}
let mut h = FnvHasher::default();
for val in vals {
h.write(val.as_bytes());
}
Ok(h.finish())
}
fn hash_labels(&self, labels: &HashMap<&str, &str>) -> Result<u64> {
if labels.len() != self.desc.variable_labels.len() {
return Err(Error::InconsistentCardinality {
expect: self.desc.variable_labels.len(),
got: labels.len(),
});
}
let mut h = FnvHasher::default();
for name in &self.desc.variable_labels {
match labels.get(&name.as_ref()) {
Some(val) => h.write(val.as_bytes()),
None => {
return Err(Error::Msg(format!(
"label name {} missing in label map",
name
)));
}
}
}
Ok(h.finish())
}
fn get_label_values<'a>(&self, labels: &'a HashMap<&str, &str>) -> Result<Vec<&'a str>> {
let mut values = Vec::new();
for name in &self.desc.variable_labels {
match labels.get(&name.as_ref()) {
Some(val) => values.push(*val),
None => {
return Err(Error::Msg(format!(
"label name {} missing in label map",
name
)));
}
}
}
Ok(values)
}
fn get_or_create_metric(&self, hash: u64, label_values: &[&str]) -> Result<T::M> {
let mut children = self.children.write();
if let Some(metric) = children.get(&hash).cloned() {
return Ok(metric);
}
let metric = self.new_metric.build(&self.opts, label_values)?;
children.insert(hash, metric.clone());
Ok(metric)
}
}
#[derive(Clone)]
pub struct MetricVec<T: MetricVecBuilder> {
pub(crate) v: Arc<MetricVecCore<T>>,
}
impl<T: MetricVecBuilder> std::fmt::Debug for MetricVec<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MetricVec")
}
}
impl<T: MetricVecBuilder> MetricVec<T> {
pub fn create(metric_type: MetricType, new_metric: T, opts: T::P) -> Result<MetricVec<T>> {
let desc = opts.describe()?;
let v = MetricVecCore {
children: RwLock::new(HashMap::new()),
desc,
metric_type,
new_metric,
opts,
};
Ok(MetricVec { v: Arc::new(v) })
}
pub fn get_metric_with_label_values(&self, vals: &[&str]) -> Result<T::M> {
self.v.get_metric_with_label_values(vals)
}
pub fn get_metric_with(&self, labels: &HashMap<&str, &str>) -> Result<T::M> {
self.v.get_metric_with(labels)
}
pub fn with_label_values(&self, vals: &[&str]) -> T::M {
self.get_metric_with_label_values(vals).unwrap()
}
pub fn with(&self, labels: &HashMap<&str, &str>) -> T::M {
self.get_metric_with(labels).unwrap()
}
pub fn remove_label_values(&self, vals: &[&str]) -> Result<()> {
self.v.delete_label_values(vals)
}
pub fn remove(&self, labels: &HashMap<&str, &str>) -> Result<()> {
self.v.delete(labels)
}
pub fn reset(&self) {
self.v.reset()
}
}
impl<T: MetricVecBuilder> Collector for MetricVec<T> {
fn desc(&self) -> Vec<&Desc> {
vec![&self.v.desc]
}
fn collect(&self) -> Vec<MetricFamily> {
vec![self.v.collect()]
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use crate::counter::CounterVec;
use crate::gauge::GaugeVec;
use crate::metrics::{Metric, Opts};
#[test]
fn test_counter_vec_with_labels() {
let vec = CounterVec::new(
Opts::new("test_couter_vec", "test counter vec help"),
&["l1", "l2"],
)
.unwrap();
let mut labels = HashMap::new();
labels.insert("l1", "v1");
labels.insert("l2", "v2");
assert!(vec.remove(&labels).is_err());
vec.with(&labels).inc();
assert!(vec.remove(&labels).is_ok());
assert!(vec.remove(&labels).is_err());
let mut labels2 = HashMap::new();
labels2.insert("l1", "v2");
labels2.insert("l2", "v1");
vec.with(&labels).inc();
assert!(vec.remove(&labels2).is_err());
vec.with(&labels).inc();
let mut labels3 = HashMap::new();
labels3.insert("l1", "v1");
assert!(vec.remove(&labels3).is_err());
}
#[test]
fn test_counter_vec_with_label_values() {
let vec = CounterVec::new(
Opts::new("test_vec", "test counter vec help"),
&["l1", "l2"],
)
.unwrap();
assert!(vec.remove_label_values(&["v1", "v2"]).is_err());
vec.with_label_values(&["v1", "v2"]).inc();
assert!(vec.remove_label_values(&["v1", "v2"]).is_ok());
vec.with_label_values(&["v1", "v2"]).inc();
assert!(vec.remove_label_values(&["v1"]).is_err());
assert!(vec.remove_label_values(&["v1", "v3"]).is_err());
}
#[test]
fn test_gauge_vec_with_labels() {
let vec = GaugeVec::new(
Opts::new("test_gauge_vec", "test gauge vec help"),
&["l1", "l2"],
)
.unwrap();
let mut labels = HashMap::new();
labels.insert("l1", "v1");
labels.insert("l2", "v2");
assert!(vec.remove(&labels).is_err());
vec.with(&labels).inc();
vec.with(&labels).dec();
vec.with(&labels).add(42.0);
vec.with(&labels).sub(42.0);
vec.with(&labels).set(42.0);
assert!(vec.remove(&labels).is_ok());
assert!(vec.remove(&labels).is_err());
}
#[test]
fn test_gauge_vec_with_label_values() {
let vec = GaugeVec::new(
Opts::new("test_gauge_vec", "test gauge vec help"),
&["l1", "l2"],
)
.unwrap();
assert!(vec.remove_label_values(&["v1", "v2"]).is_err());
vec.with_label_values(&["v1", "v2"]).inc();
assert!(vec.remove_label_values(&["v1", "v2"]).is_ok());
vec.with_label_values(&["v1", "v2"]).inc();
vec.with_label_values(&["v1", "v2"]).dec();
vec.with_label_values(&["v1", "v2"]).add(42.0);
vec.with_label_values(&["v1", "v2"]).sub(42.0);
vec.with_label_values(&["v1", "v2"]).set(42.0);
assert!(vec.remove_label_values(&["v1"]).is_err());
assert!(vec.remove_label_values(&["v1", "v3"]).is_err());
}
#[test]
fn test_vec_get_metric_with() {
let vec = CounterVec::new(
Opts::new("test_vec", "test counter vec help"),
&["b", "c", "a"],
)
.unwrap();
let mut labels = HashMap::new();
labels.insert("a", "b");
labels.insert("b", "c");
labels.insert("c", "a");
let c = vec.get_metric_with(&labels).unwrap();
let m = c.metric();
let label_pairs = m.get_label();
assert_eq!(label_pairs.len(), labels.len());
for lp in label_pairs.iter() {
assert_eq!(lp.get_value(), labels[lp.get_name()]);
}
}
}