diff --git a/src/app.rs b/src/app.rs index a028b2c..89e3771 100644 --- a/src/app.rs +++ b/src/app.rs @@ -5,24 +5,20 @@ use anyhow::Result; use axum::{ - extract::{Path, Request, State}, - middleware::{self, Next}, - response::{Html, IntoResponse}, + extract::{Path, State}, + middleware, + response::Html, routing::get, Router, }; -use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle}; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; use crate::{ - assembled_statistical_sequences::AssembledStatisticalSequences, - config::{Config, MetricsLabel}, - garglebargle::GargleBargle, - wurstsalat_generator_pro::WurstsalatGeneratorPro, + assembled_statistical_sequences::AssembledStatisticalSequences, config::Config, + garglebargle::GargleBargle, tenx_programmer, wurstsalat_generator_pro::WurstsalatGeneratorPro, }; -type StatefulIocaine = Arc; +pub type StatefulIocaine = Arc; #[derive(Debug)] pub struct Iocaine { @@ -60,85 +56,13 @@ impl Iocaine { if state.config.metrics.enable { app = app.route_layer(middleware::from_fn_with_state( state.clone(), - Self::track_metrics, + tenx_programmer::track_metrics, )); } app.with_state(state) } - async fn metrics_app() -> Router { - let recorder_handle = Self::setup_metrics_recorder(); - Router::new().route("/metrics", get(|| async move { recorder_handle.render() })) - } - - async fn track_metrics( - State(iocaine): State, - req: Request, - next: Next, - ) -> impl IntoResponse { - let headers = req.headers().clone(); - let response = next.run(req).await; - - let mut labels = Vec::new(); - if iocaine.config.metrics.labels.contains(&MetricsLabel::Host) { - if let Some(host) = headers.get("host") { - let host = host.to_str().unwrap().to_string(); - labels.push(("host", host)); - } - } - - if iocaine - .config - .metrics - .labels - .contains(&MetricsLabel::UserAgent) - || iocaine - .config - .metrics - .labels - .contains(&MetricsLabel::UserAgentGroup) - { - if let Some(ua) = headers.get("user-agent") { - let user_agent = ua.to_str().unwrap().to_string(); - - if iocaine - .config - .metrics - .labels - .contains(&MetricsLabel::UserAgent) - { - labels.push(("user-agent", user_agent.clone())); - } - - if iocaine - .config - .metrics - .labels - .contains(&MetricsLabel::UserAgentGroup) - { - if let Some(group) = iocaine - .config - .metrics - .agent_group - .iter() - .find(|agent_group_config| agent_group_config.agent.is_match(&user_agent)) - { - labels.push(("user-agent-group", group.group.clone())); - } - } - } - } - - metrics::counter!("iocaine_requests_total", &labels).increment(1); - - response - } - - fn setup_metrics_recorder() -> PrometheusHandle { - PrometheusBuilder::new().install_recorder().unwrap() - } - async fn start_main_server(self) -> std::result::Result<(), std::io::Error> { let bind = &self.config.server.bind.clone(); @@ -149,25 +73,10 @@ impl Iocaine { .await } - async fn start_metrics_server(metrics_bind: String) -> std::result::Result<(), std::io::Error> { - let metrics_listener = tokio::net::TcpListener::bind(metrics_bind).await?; - let app = Self::metrics_app().await; - - let ts = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Time went backwards"); - let labels = [("service", "iocaine".to_string())]; - metrics::gauge!("process_start_time_seconds", &labels).set(ts); - - axum::serve(metrics_listener, app) - .with_graceful_shutdown(shutdown_signal()) - .await - } - pub async fn run(self) -> Result<()> { if self.config.metrics.enable { let metrics_bind = &self.config.metrics.bind.clone(); - let metrics_server = Self::start_metrics_server(metrics_bind.to_string()); + let metrics_server = tenx_programmer::start_metrics_server(metrics_bind.to_string()); let (main_server, _metrics_server) = tokio::join!(self.start_main_server(), metrics_server); Ok(main_server?) @@ -178,7 +87,7 @@ impl Iocaine { } } -async fn shutdown_signal() { +pub async fn shutdown_signal() { let ctrl_c = async { tokio::signal::ctrl_c() .await diff --git a/src/lib.rs b/src/lib.rs index 5601b8a..6f0e31b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,4 +10,5 @@ pub mod config; pub mod assembled_statistical_sequences; pub mod garglebargle; pub mod gobbledygook; +pub mod tenx_programmer; pub mod wurstsalat_generator_pro; diff --git a/src/tenx_programmer.rs b/src/tenx_programmer.rs new file mode 100644 index 0000000..4447043 --- /dev/null +++ b/src/tenx_programmer.rs @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: 2025 Gergely Nagy +// SPDX-FileContributor: Gergely Nagy +// +// SPDX-License-Identifier: MIT + +use axum::{ + extract::{Request, State}, + middleware::Next, + response::IntoResponse, + routing::get, + Router, +}; +use metrics_exporter_prometheus::PrometheusBuilder; +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::{ + app::{shutdown_signal, StatefulIocaine}, + config::MetricsLabel, +}; + +pub async fn track_metrics( + State(iocaine): State, + req: Request, + next: Next, +) -> impl IntoResponse { + let headers = req.headers().clone(); + let response = next.run(req).await; + let cfg = &iocaine.config.metrics; + + let mut labels = Vec::new(); + if cfg.labels.contains(&MetricsLabel::Host) { + if let Some(host) = headers.get("host") { + let host = host.to_str().unwrap().to_string(); + labels.push(("host", host)); + } + } + + if cfg.labels.contains(&MetricsLabel::UserAgent) + || cfg.labels.contains(&MetricsLabel::UserAgentGroup) + { + if let Some(ua) = headers.get("user-agent") { + let user_agent = ua.to_str().unwrap().to_string(); + + if cfg.labels.contains(&MetricsLabel::UserAgent) { + labels.push(("user-agent", user_agent.clone())); + } + + if cfg.labels.contains(&MetricsLabel::UserAgentGroup) { + if let Some(group) = cfg + .agent_group + .iter() + .find(|agent_group_config| agent_group_config.agent.is_match(&user_agent)) + { + labels.push(("user-agent-group", group.group.clone())); + } + } + } + } + + metrics::counter!("iocaine_requests_total", &labels).increment(1); + + response +} + +pub async fn start_metrics_server(metrics_bind: String) -> std::result::Result<(), std::io::Error> { + let metrics_listener = tokio::net::TcpListener::bind(metrics_bind).await?; + let recorder_handle = PrometheusBuilder::new().install_recorder().unwrap(); + let app = Router::new().route("/metrics", get(|| async move { recorder_handle.render() })); + + let ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Time went backwards"); + let labels = [("service", "iocaine".to_string())]; + metrics::gauge!("process_start_time_seconds", &labels).set(ts); + + axum::serve(metrics_listener, app) + .with_graceful_shutdown(shutdown_signal()) + .await +}