iocaine 3.3.0

The deadliest poison known to AI
// SPDX-FileCopyrightText: Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

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, "")
}