prometheus_mcp/mcp/
exporter.rs

1use crate::mcp::metrics;
2use hyper::{
3    service::{make_service_fn, service_fn},
4    Body, Request, Response, Server,
5};
6use std::convert::Infallible;
7use std::net::SocketAddr;
8use tokio::sync::oneshot;
9
10/// Handle HTTP requests for metrics
11async fn metrics_handler(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
12    let metrics = metrics::get_metrics_as_string();
13    Ok(Response::new(Body::from(metrics)))
14}
15
16/// Start the metrics server
17pub async fn start_metrics_server(
18    addr: SocketAddr,
19    shutdown_signal: oneshot::Receiver<()>,
20) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
21    // Create a service to handle the metrics endpoint
22    let make_svc =
23        make_service_fn(|_conn| async { Ok::<_, Infallible>(service_fn(metrics_handler)) });
24
25    // Create the server
26    let server = Server::bind(&addr).serve(make_svc);
27
28    // Add a graceful shutdown
29    let server = server.with_graceful_shutdown(async {
30        shutdown_signal.await.ok();
31    });
32
33    println!("Metrics server listening on http://{}/metrics", addr);
34
35    // Start the server
36    server.await?;
37
38    Ok(())
39}
40
41/// Create a metrics server with the given port
42#[allow(clippy::type_complexity)]
43pub fn create_metrics_server(
44    port: u16,
45) -> (
46    tokio::task::JoinHandle<Result<(), Box<dyn std::error::Error + Send + Sync>>>,
47    oneshot::Sender<()>,
48) {
49    let addr = SocketAddr::from(([0, 0, 0, 0], port));
50    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
51
52    let handle = tokio::spawn(async move { start_metrics_server(addr, shutdown_rx).await });
53
54    (handle, shutdown_tx)
55}