o402 0.1.2

OpenAI-compatible gateway, paid with x402.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

//! Optional Prometheus scrape on `observability.metrics_bind`. Never on `:8080`.

use std::net::SocketAddr;
use std::time::Duration;

use axum::Router;
use axum::extract::State;
use axum::http::header::CONTENT_TYPE;
use axum::http::{HeaderValue, StatusCode};
use axum::response::IntoResponse;
use axum::routing::get;
use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
use tokio::net::TcpListener;

use crate::error::AppError;

/// Prometheus text exposition format.
const PROMETHEUS_CONTENT_TYPE: &str = "text/plain; version=0.0.4; charset=utf-8";

/// Listener plus scrape handle for the metrics bind.
pub(crate) struct MetricsBind {
    /// Bound metrics socket.
    listener: TcpListener,
    /// Prometheus scrape handle.
    handle: PrometheusHandle,
}

impl std::fmt::Debug for MetricsBind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MetricsBind")
            .field("listener", &self.listener)
            .finish_non_exhaustive()
    }
}

/// Install the recorder and bind `addr` when `Some`.
///
/// # Errors
///
/// Returns [`AppError::Metrics`] if the recorder cannot be installed, or
/// [`AppError::Bind`] if the listen socket fails.
pub(crate) async fn bind(addr: Option<SocketAddr>) -> Result<Option<MetricsBind>, AppError> {
    let Some(addr) = addr else {
        return Ok(None);
    };
    let handle = install()?;
    let listener = TcpListener::bind(addr)
        .await
        .map_err(|source| AppError::Bind { addr, source })?;
    let bound = listener.local_addr().unwrap_or(addr);
    tracing::info!(%bound, "metrics listening");
    Ok(Some(MetricsBind { listener, handle }))
}

/// Run `/metrics` in the background. Payment drain is independent of this task.
pub(crate) fn spawn(bound: Option<MetricsBind>) -> Option<tokio::task::JoinHandle<()>> {
    let MetricsBind { listener, handle } = bound?;
    Some(tokio::spawn(async move {
        if let Err(error) = axum::serve(listener, router(handle))
            .with_graceful_shutdown(crate::shutdown::signal())
            .await
        {
            tracing::error!(%error, "metrics server");
        }
    }))
}

fn install() -> Result<PrometheusHandle, AppError> {
    let handle = PrometheusBuilder::new()
        .install_recorder()
        .map_err(|error| AppError::Metrics(error.to_string()))?;
    metrics::gauge!("o402_up").set(1.0);
    let upkeep = handle.clone();
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_secs(5));
        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
        loop {
            interval.tick().await;
            upkeep.run_upkeep();
        }
    });
    Ok(handle)
}

fn router(handle: PrometheusHandle) -> Router {
    Router::new()
        .route("/metrics", get(scrape))
        .with_state(handle)
}

#[allow(clippy::unused_async, reason = "axum handler")]
async fn scrape(State(handle): State<PrometheusHandle>) -> impl IntoResponse {
    (
        StatusCode::OK,
        [(
            CONTENT_TYPE,
            HeaderValue::from_static(PROMETHEUS_CONTENT_TYPE),
        )],
        handle.render(),
    )
}

#[cfg(test)]
mod tests {
    use axum::body::Body;
    use axum::http::{Request, StatusCode};
    use metrics_exporter_prometheus::PrometheusBuilder;
    use tower::ServiceExt;

    use super::router;

    #[tokio::test]
    async fn metrics_route_ok() {
        let handle = PrometheusBuilder::new().build_recorder().handle();
        let response = router(handle)
            .oneshot(
                Request::get("/metrics")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "metrics status");
    }
}