use prometheus::{histogram_opts, opts, TextEncoder};
use rand::Rng;
use std::time::Duration;
use tokio::sync::OnceCell;
use trait_net::metrics::{
prometheus::{
ActiveCount, Elapsed, EndedCount, IntConstant, Latency, LatencyByStatus, StartedCount,
StatusCount,
},
AsStatusLabel, ObservedFutureExt,
};
struct Metrics {
constant: IntConstant,
elapsed: Elapsed,
started_count: StartedCount,
active_count: ActiveCount,
ended_count: EndedCount,
status: StatusCount,
latency: Latency,
latency_by_status: LatencyByStatus,
}
impl Metrics {
async fn global() -> &'static Self {
static METRICS: OnceCell<Metrics> = OnceCell::const_new();
METRICS
.get_or_init(|| async {
let constant = IntConstant::new(opts!("some_constant", "just 42"), 42).unwrap();
let elapsed = Elapsed::new(opts!("elapsed", "example runtime")).unwrap();
let started_count =
StartedCount::new(opts!("started_count", "started computations count"), &[])
.unwrap();
let active_count =
ActiveCount::new(opts!("active_count", "active computations count"), &[])
.unwrap();
let ended_count =
EndedCount::new(opts!("ended_count", "ended computations count"), &[]).unwrap();
let status =
StatusCount::new(opts!("status", "computation status count"), &[]).unwrap();
let latency =
Latency::new(histogram_opts!("latency", "computation latency"), &[]).unwrap();
let latency_by_status = LatencyByStatus::new(
histogram_opts!("latency_by_status", "computation latency by status"),
&[],
)
.unwrap();
Metrics {
constant,
elapsed,
started_count,
active_count,
ended_count,
latency,
status,
latency_by_status,
}
})
.await
}
fn register(&self) {
let registry = prometheus::default_registry();
registry.register(Box::new(self.constant.clone())).unwrap();
registry.register(Box::new(self.elapsed.clone())).unwrap();
registry
.register(Box::new(self.started_count.clone()))
.unwrap();
registry
.register(Box::new(self.active_count.clone()))
.unwrap();
registry
.register(Box::new(self.ended_count.clone()))
.unwrap();
registry.register(Box::new(self.latency.clone())).unwrap();
registry.register(Box::new(self.status.clone())).unwrap();
registry
.register(Box::new(self.latency_by_status.clone()))
.unwrap();
}
}
enum Error {
Ooops,
BigOoops,
}
impl AsStatusLabel for Error {
fn as_status_label(&self) -> String {
match self {
Error::Ooops => "ooops",
Error::BigOoops => "big_ooops",
}
.to_owned()
}
}
async fn computation() -> Result<(), Error> {
let millis = rand::thread_rng().gen_range(0..2000);
tokio::time::sleep(Duration::from_millis(millis)).await;
match rand::thread_rng().gen_range(0..3) {
0 => Ok(()),
1 => Err(Error::Ooops),
2 => Err(Error::BigOoops),
_ => unreachable!(),
}
}
async fn computation_loop() {
loop {
let metrics = Metrics::global().await;
let _ = computation()
.observe((
metrics.started_count.make_observer(&[]),
metrics.active_count.make_observer(&[]),
metrics.ended_count.make_observer(&[]),
metrics.latency.make_observer(&[]),
metrics.status.make_observer(&[]),
metrics.latency_by_status.make_observer(&[]),
))
.await;
}
}
async fn collect_loop() {
loop {
let text = TextEncoder::new()
.encode_to_string(&prometheus::default_registry().gather())
.unwrap();
println!("{text}");
println!("--------------------------------------------------");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
#[tokio::main]
async fn main() {
Metrics::global().await.register();
tokio::select! {
_ = tokio::spawn(computation_loop()) => {},
_ = tokio::spawn(computation_loop()) => {},
_ = tokio::spawn(collect_loop()) => {},
}
}