lavina/src/http.rs

use crate::prelude::*;

use std::convert::Infallible;
use std::sync::Arc;

use http_body_util::{BodyExt, Full};
use hyper::server::conn::http1;
use hyper::{body::Bytes, service::service_fn, Request, Response};
use hyper::{Method, StatusCode};

use prometheus::{Encoder, IntGauge, Opts, Registry, TextEncoder};
use tokio::net::TcpListener;
use tokio::sync::oneshot::Sender;
use tokio::task::JoinHandle;

mod ws;

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

async fn hello(
    _: Request<hyper::body::Incoming>,
) -> std::result::Result<Response<Full<Bytes>>, Infallible> {
    Ok(Response::new(Full::new(Bytes::from("Hello World!"))))
}

fn not_found() -> std::result::Result<Response<Full<Bytes>>, Infallible> {
    let mut response = Response::new(Full::new(Bytes::from("404")));
    *response.status_mut() = StatusCode::NOT_FOUND;
    Ok(response)
}

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

    Ok(Response::new(Full::new(Bytes::from(buffer))))
}

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

pub struct HttpServerActor {
    terminator: Sender<()>,
    fiber: JoinHandle<Result<()>>,
}
impl HttpServerActor {
    pub async fn launch(listener: TcpListener, metrics: Arc<Registry>) -> Result<HttpServerActor> {
        let (terminator, receiver) = tokio::sync::oneshot::channel::<()>();
        let fiber = tokio::task::spawn(Self::main_loop(listener, receiver, metrics));
        Ok(HttpServerActor { terminator, fiber })
    }

    pub async fn terminate(self) -> Result<()> {
        match self.terminator.send(()) {
            Ok(_) => {}
            Err(_) => failure("wat")?,
        }
        self.fiber.await??;
        Ok(())
    }

    async fn main_loop(
        listener: TcpListener,
        termination: impl Future,
        registry: Arc<Registry>,
    ) -> Result<()> {
        log::info!("Starting the http server");
        pin!(termination);
        let reqs = IntGauge::with_opts(Opts::new("sockets", "Number of open sockets"))?;
        registry.register(Box::new(reqs.clone()))?;

        loop {
            select! {
                _ = &mut termination => {
                    log::info!("Terminating the http server");
                    return Ok(())
                },
                result = listener.accept() => {
                    let (stream, _) = result?;
                    let c = registry.clone();
                    let reqs = reqs.clone();
                    tokio::task::spawn(async move {
                        reqs.inc();
                        let c = c.clone();
                        if let Err(err) = http1::Builder::new()
                            .serve_connection(stream, service_fn(move |r| route(c.clone(), r)))
                            .with_upgrades()
                            .await
                        {
                            tracing::error!("Error serving connection: {:?}", err);
                        }
                        reqs.dec();
                    });
                },
            }
        }
    }
}