forked from lavina/lavina
1
0
Fork 0
lavina/src/util/telemetry.rs

82 lines
2.6 KiB
Rust
Raw Normal View History

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