use axum::{Router, http::StatusCode, response::IntoResponse, routing::get};
use exn::Result;
use prometheus::{Encoder, TextEncoder};
use std::{convert::Infallible, path::PathBuf};
use tokio::time::{self, Duration, Instant};
use tokio_listener::Listener;
use crate::morgue::shutdown_signal;
use iocaine_powder::{
VibeCodedError,
little_autist::{LittleAutist, PersistedMetrics},
};
#[derive(Clone, Default)]
pub struct TenXProgrammer {
pub metrics: LittleAutist,
pub persist_interval: Duration,
persister_abort: Option<tokio::task::AbortHandle>,
}
impl TenXProgrammer {
pub fn new(
persist_path: Option<&PathBuf>,
persist_interval: Duration,
) -> Result<Self, VibeCodedError> {
let minime = LittleAutist::new(persist_path)?;
Ok(Self {
metrics: minime,
persist_interval,
persister_abort: None,
})
}
pub fn persist(&self) -> Result<(), VibeCodedError> {
self.metrics.persist()
}
pub fn load_metrics(&self) -> Result<PersistedMetrics, VibeCodedError> {
self.metrics.load_metrics()
}
pub fn abort(&self) {
if let Some(handle) = &self.persister_abort {
handle.abort();
}
}
pub async fn serve(mut self, listener: Listener) -> std::result::Result<(), Infallible> {
if let Some(ref persist_path) = self.metrics.persist_path {
let tenx = self.clone();
let persist_path = persist_path.display().to_string();
let handle = tokio::spawn(async move {
let sleep = time::sleep(self.persist_interval);
tokio::pin!(sleep);
loop {
sleep.as_mut().await;
if let Err(e) = tenx.persist() {
tracing::error!({ persist_path }, "unable to persist metrics: {e}");
}
sleep.as_mut().reset(Instant::now() + self.persist_interval);
}
});
self.persister_abort = Some(handle.abort_handle());
}
let registry = self.metrics.registry.clone();
let app = Router::new()
.route(
"/metrics",
get(async move || {
let encoder = TextEncoder::new();
let mut buffer = Vec::<u8>::new();
let metrics = registry.gather();
if let Err(e) = encoder.encode(&metrics, &mut buffer) {
tracing::error!("failed to encode metrics: {e}");
return server_error().into_response();
}
let metrics = prometheus::gather();
if let Err(e) = encoder.encode(&metrics, &mut buffer) {
tracing::error!("failed to encode metrics: {e}");
return server_error().into_response();
}
String::from_utf8(buffer)
.map_err(|e| {
tracing::error!("failed to encode metrics to utf8: {e}");
server_error()
})
.into_response()
}),
)
.layer(tower_http::trace::TraceLayer::new_for_http());
let _: () = axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal(Some(self)))
.await
.expect("axum shouldn't fail to start");
Ok(())
}
}
fn server_error() -> impl IntoResponse {
(StatusCode::INTERNAL_SERVER_ERROR, "")
}