use anyhow::Result;
use axum::{Router, routing::get};
use prometheus::{Encoder, IntCounterVec, IntGaugeVec, Opts, Registry, TextEncoder};
use serde::Deserialize;
use std::collections::HashMap;
use std::fs::File;
use crate::{
app::VERSION,
config::{MetricsConfig, MetricsLabel},
queer::HRT,
};
#[derive(Debug, Clone)]
pub struct TenXProgrammer {
pub registry: Registry,
pub counters: TenXProgrammerCounters,
pub persist_path: Option<String>,
pub labels: Vec<&'static str>,
}
#[derive(Debug, Clone)]
pub struct TenXProgrammerCounters {
pub request_counter: IntCounterVec,
pub garbage_served_counter: IntCounterVec,
pub maze_depth: IntCounterVec,
pub challenge_counter: IntCounterVec,
pub handler_hits: 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");
}
if config.labels.contains(&MetricsLabel::Verdict) {
labels.push("verdict");
}
labels
}
#[must_use]
pub fn build_label_values<'a>(
config: &'a MetricsConfig,
headers: &'a axum::http::HeaderMap,
verdict: &'a str,
) -> Vec<String> {
let mut values = Vec::new();
if config.labels.contains(&MetricsLabel::Host) {
if let Some(host) = headers.get("host") {
let host = String::from_utf8_lossy(host.as_bytes()).to_string();
values.push(host);
} else {
values.push("<unknown>".into());
}
}
if config.labels.contains(&MetricsLabel::UserAgent)
|| config.labels.contains(&MetricsLabel::UserAgentGroup)
{
let ua = headers.get("user-agent").map_or_else(
|| "<unknown>".into(),
|agent| String::from_utf8_lossy(agent.as_bytes()).to_string(),
);
if config.labels.contains(&MetricsLabel::UserAgent) {
values.push(ua.clone());
}
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 = group.map_or("<unknown>", |cfg| &cfg.group);
values.push(group.to_owned());
}
}
if config.labels.contains(&MetricsLabel::Verdict) {
values.push(verdict.to_owned());
}
values
}
pub fn new(config: &MetricsConfig) -> Result<Option<Self>> {
if !config.enable {
return Ok(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)?;
registry.register(Box::new(request_counter.clone()))?;
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)?;
registry.register(Box::new(garbage_served_counter.clone()))?;
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)?;
registry.register(Box::new(maze_depth.clone()))?;
let challenge_counter_opts =
Opts::new("iocaine_challenges", "Number of challenges presented");
let challenge_counter = IntCounterVec::new(challenge_counter_opts, &labels)?;
registry.register(Box::new(challenge_counter.clone()))?;
let version_opts = Opts::new(
"iocaine_version",
"Version of the running iocaine (in the 'version' label)",
);
let version = IntGaugeVec::new(version_opts, &["version"])?;
registry.register(Box::new(version.clone()))?;
version.with_label_values(&[VERSION]).set(1);
let handler_hits_opts = Opts::new(
"iocaine_request_handler_hits",
"Total number of times the request handler was invoked",
);
let handler_hits = IntCounterVec::new(handler_hits_opts, &["id"])?;
registry.register(Box::new(handler_hits.clone()))?;
let tenx = Self {
registry,
counters: TenXProgrammerCounters {
request_counter,
garbage_served_counter,
maze_depth,
challenge_counter,
handler_hits,
},
persist_path: config.persist_path.clone(),
labels,
};
tenx.load_metrics()?;
Ok(Some(tenx))
}
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()
}),
)
}
pub fn persist(self) -> Result<()> {
let Some(persist_path) = self.persist_path else {
return Ok(());
};
tracing::info!({ persist_path }, "persisting metrics");
let mut f = File::create(persist_path)?;
let encoder = HRT::new();
let metrics = self.registry.gather();
Ok(encoder.encode(&metrics, &mut f)?)
}
fn build_persistent_label_values(
&self,
metric: &str,
labels: &HashMap<String, String>,
) -> Result<Vec<String>> {
if self.labels.len() != labels.keys().count() {
tracing::warn!(
{
persisted_labels = labels.keys().cloned().collect::<Vec<String>>().join(", "),
expected_labels = self.labels.join(", "),
metric
},
"incompatible labelset"
);
return Err(anyhow::anyhow!("incompatible metric labels"));
}
let mut values = Vec::new();
for label in &self.labels {
if let Some(value) = labels.get(*label) {
values.push(value.to_owned());
} else {
tracing::warn!(
{
persisted_labels = labels.keys().cloned().collect::<Vec<String>>().join(", "),
expected_labels = self.labels.join(", "),
metric
},
"incompatible labelset"
);
return Err(anyhow::anyhow!("incompatible metric labels"));
}
}
Ok(values)
}
#[allow(clippy::cognitive_complexity)]
fn load_metrics(&self) -> Result<()> {
#[derive(Deserialize, Debug)]
struct PersistedMetrics {
#[serde(flatten)]
metrics: HashMap<String, Vec<PersistedMetric>>,
}
#[derive(Deserialize, Debug)]
struct PersistedMetric {
labels: HashMap<String, String>,
value: f64,
}
let Some(ref persist_path) = self.persist_path else {
return Ok(());
};
tracing::info!({ persist_path }, "loading persisted metrics");
let Ok(data) = std::fs::read_to_string(persist_path) else {
return Ok(());
};
let persisted: PersistedMetrics = serde_json::from_str(&data)?;
for (name, metrics) in &persisted.metrics {
for metric in metrics {
let counter;
let label_values;
match name.as_str() {
"iocaine_requests_total" => {
if let Ok(labels) =
self.build_persistent_label_values(name.as_str(), &metric.labels)
{
counter = &self.counters.request_counter;
label_values = labels;
} else {
continue;
}
}
"iocaine_garbage_served" => {
if let Ok(labels) =
self.build_persistent_label_values(name.as_str(), &metric.labels)
{
counter = &self.counters.garbage_served_counter;
label_values = labels;
} else {
continue;
}
}
"iocaine_maze_depth" => {
if let Ok(labels) =
self.build_persistent_label_values(name.as_str(), &metric.labels)
{
counter = &self.counters.maze_depth;
label_values = labels;
} else {
continue;
}
}
"iocaine_challenges" => {
if let Ok(labels) =
self.build_persistent_label_values(name.as_str(), &metric.labels)
{
counter = &self.counters.challenge_counter;
label_values = labels;
} else {
continue;
}
}
"iocaine_request_handler_hits" => {
if metric.labels.len() != 1 || !metric.labels.contains_key("id") {
tracing::warn!(
{
persisted_labels = metric
.labels
.keys()
.cloned()
.collect::<Vec<String>>()
.join(", "),
expected_labels = "id",
metric = name
},
"incompatible labelset"
);
continue;
}
counter = &self.counters.handler_hits;
label_values = Vec::from(&[metric.labels.get("id").unwrap().to_owned()]);
}
_ => {
tracing::warn!({ metric = name }, "unknown metric");
continue;
}
}
#[allow(clippy::cast_possible_truncation)]
#[allow(clippy::cast_sign_loss)]
counter
.with_label_values(&label_values)
.inc_by(metric.value as u64);
}
}
Ok(())
}
}