use std::{
collections::HashMap,
sync::{Arc, RwLock},
};
use sealed::sealed;
use crate::{metric, Metric};
use super::KeyName;
pub type Map<K, V> = Arc<RwLock<HashMap<K, V>>>;
pub type Collection<M> = Map<KeyName, metric::Describable<Option<M>>>;
#[derive(Clone, Debug)]
pub struct Storage {
pub(crate) prometheus: prometheus::Registry,
pub(super) counters: Collection<metric::PrometheusIntCounter>,
pub(super) gauges: Collection<metric::PrometheusGauge>,
pub(super) histograms: Collection<metric::PrometheusHistogram>,
}
#[sealed]
impl super::Get<Collection<metric::PrometheusIntCounter>> for Storage {
fn collection(&self) -> &Collection<metric::PrometheusIntCounter> {
&self.counters
}
}
#[sealed]
impl super::Get<Collection<metric::PrometheusGauge>> for Storage {
fn collection(&self) -> &Collection<metric::PrometheusGauge> {
&self.gauges
}
}
#[sealed]
impl super::Get<Collection<metric::PrometheusHistogram>> for Storage {
fn collection(&self) -> &Collection<metric::PrometheusHistogram> {
&self.histograms
}
}
impl Default for Storage {
fn default() -> Self {
Self {
prometheus: prometheus::default_registry().clone(),
counters: Collection::default(),
gauges: Collection::default(),
histograms: Collection::default(),
}
}
}
impl Storage {
pub fn describe<M>(&self, name: &str, description: String)
where
M: metric::Bundled,
<M as metric::Bundled>::Bundle: Clone,
Self: super::Get<Collection<<M as metric::Bundled>::Bundle>>,
{
#![allow(
clippy::missing_panics_doc,
clippy::unwrap_in_result,
clippy::unwrap_used
)]
use super::Get as _;
let read_storage = self.collection().read().unwrap();
if let Some(metric) = read_storage.get(name) {
metric.description.store(Arc::new(description));
} else {
drop(read_storage);
let mut write_storage = self.collection().write().unwrap();
if let Some(metric) = write_storage.get(name) {
metric.description.store(Arc::new(description));
} else {
drop(write_storage.insert(
name.into(),
metric::Describable::only_description(description),
));
}
}
}
fn register<'k, M>(
&self,
key: &'k metrics::Key,
) -> prometheus::Result<Arc<Metric<M>>>
where
M: metric::Bundled + prometheus::core::Metric + Clone,
<M as metric::Bundled>::Bundle: metric::Bundle<Single = M>
+ prometheus::core::Collector
+ Clone
+ TryFrom<&'k metrics::Key, Error = prometheus::Error>
+ 'static,
Self: super::Get<Collection<<M as metric::Bundled>::Bundle>>,
{
#![allow(
clippy::missing_panics_doc,
clippy::unwrap_in_result,
clippy::unwrap_used
)]
use super::Get as _;
use metric::Bundle as _;
let name = key.name();
let mut bundle_opt = self
.collection()
.read()
.unwrap()
.get(name)
.and_then(|m| m.metric.clone());
let bundle = if let Some(bundle) = bundle_opt {
bundle
} else {
let mut storage = self.collection().write().unwrap();
bundle_opt = storage.get(name).and_then(|m| m.metric.clone());
if let Some(bundle) = bundle_opt {
bundle
} else {
let bundle: <M as metric::Bundled>::Bundle = key.try_into()?;
let mut entry = storage.entry(name.into()).or_default();
self.prometheus.register(Box::new(
entry.clone().map(|_| bundle.clone()),
))?;
entry.metric = Some(bundle.clone());
bundle
}
};
Ok(Arc::new(Metric::wrap(bundle.get_single_metric(key)?)))
}
pub fn register_external<M>(&self, metric: M) -> prometheus::Result<()>
where
M: metric::Bundled + prometheus::core::Collector,
<M as metric::Bundled>::Bundle:
prometheus::core::Collector + Clone + 'static,
Self: super::Get<Collection<<M as metric::Bundled>::Bundle>>,
{
#![allow(
clippy::missing_panics_doc,
clippy::unwrap_in_result,
clippy::unwrap_used
)]
use super::Get as _;
let name = metric
.desc()
.first()
.map(|d| d.fq_name.clone())
.unwrap_or_default();
let entry = metric::Describable::wrap(Some(metric.into_bundle()));
let mut storage = self.collection().write().unwrap();
self.prometheus
.register(Box::new(entry.clone().map(Option::unwrap)))?;
drop(storage.insert(name, entry));
Ok(())
}
}
impl metrics_util::registry::Storage<metrics::Key> for Storage {
type Counter = metric::Fallible<prometheus::IntCounter>;
type Gauge = metric::Fallible<prometheus::Gauge>;
type Histogram = metric::Fallible<prometheus::Histogram>;
fn counter(&self, key: &metrics::Key) -> Self::Counter {
self.register::<prometheus::IntCounter>(key).into()
}
fn gauge(&self, key: &metrics::Key) -> Self::Gauge {
self.register::<prometheus::Gauge>(key).into()
}
fn histogram(&self, key: &metrics::Key) -> Self::Histogram {
self.register::<prometheus::Histogram>(key).into()
}
}