Skip to main content

slave_framework/
server.rs

1//! HTTP server lifecycle.
2
3use axum::Router;
4use tokio::net::TcpListener;
5
6use crate::{Error, Result};
7
8/// Binds to `host:port` and serves `router` until Ctrl+C or SIGTERM is received.
9///
10/// In-flight requests are allowed to finish before this function returns.
11pub async fn serve(router: Router, host: &str, port: u16) -> Result<()> {
12    let listener = TcpListener::bind((host, port))
13        .await
14        .map_err(|source| Error::Bind {
15            addr: format!("{host}:{port}"),
16            source,
17        })?;
18    let addr = listener.local_addr().map_err(Error::Server)?;
19    tracing::info!("Server running on http://{addr}");
20
21    axum::serve(listener, router)
22        .with_graceful_shutdown(shutdown_signal())
23        .await
24        .map_err(Error::Server)?;
25
26    tracing::info!("Server stopped");
27    Ok(())
28}
29
30/// Resolves when the process receives Ctrl+C, or SIGTERM on Unix.
31pub async fn shutdown_signal() {
32    let ctrl_c = async {
33        if tokio::signal::ctrl_c().await.is_err() {
34            std::future::pending::<()>().await;
35        }
36    };
37
38    #[cfg(unix)]
39    let terminate = async {
40        use tokio::signal::unix::{SignalKind, signal};
41        match signal(SignalKind::terminate()) {
42            Ok(mut sigterm) => {
43                sigterm.recv().await;
44            }
45            Err(_) => std::future::pending::<()>().await,
46        }
47    };
48
49    #[cfg(not(unix))]
50    let terminate = std::future::pending::<()>();
51
52    tokio::select! {
53        () = ctrl_c => {},
54        () = terminate => {},
55    }
56    tracing::info!("Shutting down gracefully...");
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62
63    #[tokio::test]
64    async fn reports_bind_errors_with_address() {
65        let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
66        let port = taken.local_addr().unwrap().port();
67
68        let err = serve(Router::new(), "127.0.0.1", port).await.unwrap_err();
69        assert!(matches!(err, Error::Bind { .. }), "{err}");
70        assert!(err.to_string().contains(&port.to_string()));
71    }
72}