lavina/src/util/telemetry.rs

100 lines
3.3 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;
2023-02-15 17:10:54 +00:00
use crate::core::room::RoomRegistry;
use crate::prelude::*;
use crate::util::http::*;
use crate::util::Terminator;
type BoxBody = http_body_util::combinators::BoxBody<Bytes, Infallible>;
2023-02-15 17:10:54 +00:00
type HttpResult<T> = std::result::Result<T, Infallible>;
#[derive(Deserialize, Debug)]
pub struct ServerConfig {
pub listen_on: SocketAddr,
}
2023-02-15 17:10:54 +00:00
pub async fn launch(
config: ServerConfig,
metrics: MetricsRegistry,
rooms: RoomRegistry,
) -> Result<Terminator> {
log::info!("Starting the telemetry service");
let listener = TcpListener::bind(config.listen_on).await?;
log::debug!("Listener started");
let (signal, rx) = channel();
2023-02-15 17:10:54 +00:00
let handle = tokio::task::spawn(main_loop(listener, metrics, rooms, rx.map(|_| ())));
let terminator = Terminator::from_raw(signal, handle);
Ok(terminator)
}
async fn main_loop(
listener: TcpListener,
metrics: MetricsRegistry,
2023-02-15 17:10:54 +00:00
rooms: RoomRegistry,
termination: impl Future<Output = ()>,
) -> Result<()> {
pin!(termination);
loop {
select! {
biased;
_ = &mut termination => break,
result = listener.accept() => {
let (stream, _) = result?;
let metrics = metrics.clone();
2023-02-15 17:10:54 +00:00
let rooms = rooms.clone();
tokio::task::spawn(async move {
let registry = metrics.clone();
2023-02-15 17:10:54 +00:00
let rooms = rooms.clone();
let server = http1::Builder::new().serve_connection(stream, service_fn(move |r| route(registry.clone(), rooms.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,
2023-02-15 17:10:54 +00:00
rooms: RoomRegistry,
request: Request<hyper::body::Incoming>,
) -> std::result::Result<Response<BoxBody>, Infallible> {
match (request.method(), request.uri().path()) {
2023-02-15 17:10:54 +00:00
(&Method::GET, "/metrics") => Ok(endpoint_metrics(registry)?.map(BodyExt::boxed)),
(&Method::GET, "/rooms") => Ok(endpoint_rooms(rooms).await?.map(BodyExt::boxed)),
_ => Ok(not_found()?.map(BodyExt::boxed)),
}
}
2023-02-15 17:10:54 +00:00
fn endpoint_metrics(registry: MetricsRegistry) -> HttpResult<Response<Full<Bytes>>> {
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))))
}
2023-02-15 17:10:54 +00:00
async fn endpoint_rooms(rooms: RoomRegistry) -> HttpResult<Response<Full<Bytes>>> {
let room_list = rooms.get_all_rooms().await;
let mut buffer = vec![];
serde_json::to_writer(&mut buffer, &room_list).expect("unexpected fail when writing to vec");
Ok(Response::new(Full::new(Bytes::from(buffer))))
}