use std::net::SocketAddr;
use alog::{MessageLevel, alog_channel, use_channel};
use axum::Router;
use tokio::sync::oneshot;
use_channel!("SBSRV");
pub struct SubServer {
pub local_addr: SocketAddr,
shutdown_tx: Option<oneshot::Sender<()>>,
join_handle: tokio::task::JoinHandle<()>,
}
impl SubServer {
pub fn spawn(router: Router, label: &str) -> anyhow::Result<Self> {
let std_listener = std::net::TcpListener::bind("127.0.0.1:0")?;
std_listener.set_nonblocking(true)?;
let local_addr = std_listener.local_addr()?;
let listener = tokio::net::TcpListener::from_std(std_listener)?;
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let label = label.to_string();
let join_handle = tokio::spawn(async move {
let server = axum::serve(listener, router).with_graceful_shutdown(async {
let _ = shutdown_rx.await;
});
if let Err(e) = server.await {
alog_channel!(MessageLevel::Warning, "sub-server '{label}' error: {e}");
}
});
Ok(Self {
local_addr,
shutdown_tx: Some(shutdown_tx),
join_handle,
})
}
pub async fn shutdown(mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
let _ = self.join_handle.await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use axum::routing::get;
#[tokio::test]
async fn spawn_serves_the_given_router_and_shutdown_frees_the_port() {
let router = Router::new().route("/ping", get(|| async { "pong" }));
let server = SubServer::spawn(router, "test").unwrap();
let addr = server.local_addr;
let resp = reqwest::get(format!("http://{addr}/ping")).await.unwrap();
assert!(resp.status().is_success());
assert_eq!(resp.text().await.unwrap(), "pong");
server.shutdown().await;
std::net::TcpListener::bind(addr).unwrap();
}
}