mise_server/

mise.rs

use crate::server::Server;
use http::StatusCode;
use metrics::{Counter, Histogram, counter, histogram};
use serde_json::Value;
use std::{
    collections::HashMap,
    net::SocketAddr,
    panic::{AssertUnwindSafe, catch_unwind},
    time::Instant,
};
use tracing::{trace, warn};

pub struct Request(pub http::Request<Value>);
pub struct Response(pub http::Response<Value>);

/// Final routes are always in JSON.
pub trait Route: Fn(Request) -> Response + 'static {}
impl<T> Route for T where T: Fn(Request) -> Response + 'static {}

// In some cases it's possible to install a route that only returns text.
// this is used for certain side behavior such as returning prometheus scrape
// renders. These are meant to be static and have no parameter and respond only
// on absolute paths: no request object is available and the response is always
// a string.
pub trait TextRoute: Fn() -> String + 'static {}
impl<T> TextRoute for T where T: Fn() -> String + 'static {}

/// Server resource and entry point.
#[derive(Default)]
pub struct Mise {
    routes: HashMap<RouteMethod, HashMap<String, RouteContext>>,
    text_routes: HashMap<String, Box<dyn TextRoute>>,
}

impl Mise {
    /// Create a new default server.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a get route.
    pub fn get<F: Route>(mut self, path: &str, f: F) -> Self {
        let get_routes = self.routes.entry(RouteMethod::Get).or_default();
        let d: Box<dyn Route> = Box::new(f);
        get_routes.insert(path.to_string(), (path, d).into());
        self
    }

    /// Register a text result only.
    /// This is used for cases such as prometheus scrape renders.
    /// Text routes always take precedence even if the route is already defined
    /// for any other methods.
    pub fn text<F: TextRoute>(mut self, path: &str, f: F) -> Self {
        self.text_routes.insert(path.to_string(), Box::new(f));
        self
    }

    /// Registers a post route. Body is obtained from [Request::body] and it is
    /// always a json [Value].
    pub fn post<F: Route>(mut self, path: &str, f: F) -> Self {
        let get_routes = self.routes.entry(RouteMethod::Post).or_default();
        let d: Box<dyn Route> = Box::new(f);
        get_routes.insert(path.to_string(), (path, d).into());
        self
    }

    /// Starts the server. Blocks until the server quits.
    /// Can panic if cannot bind the server.
    pub fn serve(self, addr: SocketAddr) {
        Server::serve(
            RequestProcessor {
                routes: self.routes,
                text_routes: self.text_routes,
            },
            addr,
        )
        .run();
    }
}

impl From<Request> for Value {
    fn from(value: Request) -> Self {
        value.0.body().to_owned()
    }
}

impl From<Value> for Response {
    fn from(value: Value) -> Self {
        Response(http::Response::new(value))
    }
}

impl From<http::StatusCode> for Response {
    fn from(value: http::StatusCode) -> Self {
        Response(
            http::Response::builder()
                .status(value)
                .body(Value::Null)
                .expect("Statically built body should not fail"),
        )
    }
}

struct RouteMetrics {
    // Histograms also count the events so no need to also have a counter for
    // each request. The histogram will do both.
    call_seconds: Histogram,
    // call_seconds only measure successful requests, otherwise use this for
    // marking when a route returns a 500/panic
    error_count: Counter,
}

impl RouteMetrics {
    fn new(name: &str) -> Self {
        Self {
            call_seconds: histogram!("http_request_seconds", "uri" => name.to_string()),
            error_count: counter!("http_error", "uri" => name.to_string()),
        }
    }
}

struct RouteContext {
    route: Box<dyn Route>,
    metrics: RouteMetrics,
}

impl From<(&str, Box<dyn Route>)> for RouteContext {
    fn from(value: (&str, Box<dyn Route>)) -> Self {
        Self {
            route: value.1,
            metrics: RouteMetrics::new(value.0),
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) enum RouteMethod {
    Get,
    Post,
}

impl TryFrom<&str> for RouteMethod {
    type Error = ();

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        if value.eq_ignore_ascii_case("GET") {
            return Ok(RouteMethod::Get);
        }
        if value.eq_ignore_ascii_case("POST") {
            return Ok(RouteMethod::Post);
        }
        Err(())
    }
}

pub(crate) struct RequestProcessor {
    routes: HashMap<RouteMethod, HashMap<String, RouteContext>>,
    text_routes: HashMap<String, Box<dyn TextRoute>>,
}

pub(crate) enum ProcessorResponse {
    Json(Response),
    Text(String),
}

impl RequestProcessor {
    pub(crate) fn process(&self, req: Request) -> ProcessorResponse {
        if let Some(txt) = self.text_routes.get(req.0.uri().path()) {
            return ProcessorResponse::Text(txt());
        }

        let start = Instant::now();

        let Ok(method) = req.0.method().as_str().try_into() else {
            return ProcessorResponse::Json(StatusCode::METHOD_NOT_ALLOWED.into());
        };
        let Some(routes) = self.routes.get(&method) else {
            return ProcessorResponse::Json(StatusCode::NOT_FOUND.into());
        };
        let Some(route) = routes.get(req.0.uri().path()) else {
            return ProcessorResponse::Json(StatusCode::NOT_FOUND.into());
        };

        trace!("Received {:?}", req.0);
        let res = catch_unwind(AssertUnwindSafe(|| (route.route)(req)));
        let resp = match res {
            Ok(resp) => resp,
            Err(ref e) => {
                route.metrics.error_count.increment(1);
                warn!("Internal server error: {e:?}");
                return ProcessorResponse::Json(StatusCode::INTERNAL_SERVER_ERROR.into());
            }
        };
        trace!("Sending {:?}", resp.0);

        route
            .metrics
            .call_seconds
            .record(Instant::now().duration_since(start).as_secs_f64());

        ProcessorResponse::Json(resp)
    }
}

impl Request {
    /// Returns the query param by name.
    pub fn query_param(&self, name: &str) -> Option<String> {
        let q = self.0.uri().query()?;
        let f = format!("{name}=");
        let idx = q.find(&f)?;
        let end = q[idx..].find('&').unwrap_or(q.len());
        Some(q[idx + f.len()..end].to_string())
    }

    /// If there is a body in the request, then this will be the json of that
    pub fn body(&self) -> &Value {
        self.0.body()
    }
}