lavina/src/util/telemetry.rs

use std::convert::Infallible;
use std::net::SocketAddr;

use futures_util::FutureExt;
use http_body_util::{BodyExt, Full};
use hyper::body::Bytes;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response};
use prometheus::{Encoder, Registry as MetricsRegistry, TextEncoder};
use serde::Deserialize;
use tokio::net::TcpListener;
use tokio::sync::oneshot::channel;

use crate::prelude::*;

use crate::util::http::*;
use crate::util::Terminator;

type BoxBody = http_body_util::combinators::BoxBody<Bytes, Infallible>;

#[derive(Deserialize, Debug)]
pub struct ServerConfig {
    pub listen_on: SocketAddr,
}

pub async fn launch(config: ServerConfig, metrics: MetricsRegistry) -> Result<Terminator> {
    log::info!("Starting the telemetry service");
    let listener = TcpListener::bind(config.listen_on).await?;
    log::debug!("Listener started");
    let (signal, rx) = channel();
    let handle = tokio::task::spawn(main_loop(listener, metrics, rx.map(|_| ())));
    let terminator = Terminator::from_raw(signal, handle);
    Ok(terminator)
}

async fn main_loop(
    listener: TcpListener,
    metrics: MetricsRegistry,
    termination: impl Future<Output = ()>,
) -> Result<()> {
    pin!(termination);
    loop {
        select! {
            biased;
            _ = &mut termination => break,
            result = listener.accept() => {
                let (stream, _) = result?;
                let metrics = metrics.clone();
                tokio::task::spawn(async move {
                    let registry = metrics.clone();
                    let server = http1::Builder::new().serve_connection(stream, service_fn(move |r| route(registry.clone(), r)));
                    if let Err(err) = server.await {
                        tracing::error!("Error serving connection: {:?}", err);
                    }
                });
            },
        }
    }
    log::info!("Terminating the telemetry service");
    Ok(())
}

async fn route(
    registry: MetricsRegistry,
    request: Request<hyper::body::Incoming>,
) -> std::result::Result<Response<BoxBody>, Infallible> {
    match (request.method(), request.uri().path()) {
        (&Method::GET, "/metrics") => Ok(metrics(registry)?.map(BodyExt::boxed)),
        _ => Ok(not_found()?.map(BodyExt::boxed)),
    }
}

fn metrics(registry: MetricsRegistry) -> std::result::Result<Response<Full<Bytes>>, Infallible> {
    let mf = registry.gather();
    let mut buffer = vec![];
    TextEncoder
        .encode(&mf, &mut buffer)
        .expect("write to vec cannot fail");
    Ok(Response::new(Full::new(Bytes::from(buffer))))
}