use axum::{routing::get, Router};
use prometheus::{Encoder, IntCounterVec, IntGaugeVec, Opts, Registry, TextEncoder};
use crate::{
app::AppError,
config::{MetricsConfig, MetricsLabel},
};
const VERSION: &str = env!("CARGO_PKG_VERSION");
#[derive(Debug, Clone)]
pub struct TenXProgrammer {
pub registry: Registry,
pub counters: TenXProgrammerCounters,
}
#[derive(Debug, Clone)]
pub struct TenXProgrammerCounters {
pub request_counter: IntCounterVec,
pub garbage_served_counter: IntCounterVec,
pub maze_depth: IntCounterVec,
}
impl TenXProgrammer {
fn build_label_names(config: &MetricsConfig) -> Vec<&'static str> {
let mut labels = Vec::new();
if config.labels.contains(&MetricsLabel::Host) {
labels.push("host");
}
if config.labels.contains(&MetricsLabel::UserAgent) {
labels.push("user_agent");
}
if config.labels.contains(&MetricsLabel::UserAgentGroup) {
labels.push("user_agent_group");
}
labels
}
pub fn build_label_values<'a>(
config: &'a MetricsConfig,
headers: &'a axum::http::HeaderMap,
) -> Result<Vec<&'a str>, AppError> {
let mut values = Vec::new();
if config.labels.contains(&MetricsLabel::Host) {
if let Some(host) = headers.get("host") {
let host = host.to_str()?;
values.push(host);
} else {
values.push("<unknown>");
}
}
if config.labels.contains(&MetricsLabel::UserAgent)
|| config.labels.contains(&MetricsLabel::UserAgentGroup)
{
let ua = match headers.get("user-agent") {
Some(agent) => agent.to_str()?,
_ => "<unknown>",
};
if config.labels.contains(&MetricsLabel::UserAgent) {
values.push(ua);
}
if config.labels.contains(&MetricsLabel::UserAgentGroup) {
let group = config
.agent_group
.iter()
.find(|agent_group_config| agent_group_config.agent.is_match(ua));
let group = match group {
Some(cfg) => &cfg.group,
_ => "<unknown>",
};
values.push(group);
}
}
Ok(values)
}
pub fn new(config: &MetricsConfig) -> Option<Self> {
if !config.enable {
return None;
}
let labels = Self::build_label_names(config);
let registry = Registry::new();
let request_counter_opts =
Opts::new("iocaine_requests_total", "Total number of requests served");
let request_counter = IntCounterVec::new(request_counter_opts, &labels).unwrap();
registry
.register(Box::new(request_counter.clone()))
.unwrap();
let garbage_served_counter_opts = Opts::new(
"iocaine_garbage_served",
"Total amount of garbage served (in bytes)",
);
let garbage_served_counter =
IntCounterVec::new(garbage_served_counter_opts, &labels).unwrap();
registry
.register(Box::new(garbage_served_counter.clone()))
.unwrap();
let maze_depth_opts = Opts::new(
"iocaine_maze_depth",
"Maximum explored depth of the maze (in path parts)",
);
let maze_depth = IntCounterVec::new(maze_depth_opts, &labels).unwrap();
registry.register(Box::new(maze_depth.clone())).unwrap();
let version_opts = Opts::new(
"iocaine_version",
"Version of the running iocaine (in the 'version' label)",
);
let version = IntGaugeVec::new(version_opts, &["version"]).unwrap();
registry.register(Box::new(version.clone())).unwrap();
version.with_label_values(&[VERSION]).set(1);
Some(Self {
registry,
counters: TenXProgrammerCounters {
request_counter,
garbage_served_counter,
maze_depth,
},
})
}
pub fn app(self) -> Router {
Router::new().route(
"/metrics",
get(|| async move {
let encoder = TextEncoder::new();
let mut buffer = Vec::<u8>::new();
let metrics = self.registry.gather();
encoder.encode(&metrics, &mut buffer).unwrap();
let metrics = prometheus::gather();
encoder.encode(&metrics, &mut buffer).unwrap();
String::from_utf8(buffer).unwrap()
}),
)
}
}