use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;
use axum::Router;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use crate::error::NativeError;
const READY_TIMEOUT: Duration = Duration::from_secs(5);
pub const DEFAULT_GRACE: Duration = Duration::from_secs(5);
pub struct EmbeddedServer {
listener: TcpListener,
addr: SocketAddr,
}
impl EmbeddedServer {
pub async fn bind() -> Result<Self, NativeError> {
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
.await
.map_err(|e| {
NativeError::new(
"listener",
format!("cannot bind a loopback port for the embedded server: {e}"),
)
})?;
let addr = listener.local_addr().map_err(|e| {
NativeError::new("listener", format!("the listener has no address: {e}"))
})?;
Ok(EmbeddedServer { listener, addr })
}
pub fn addr(&self) -> SocketAddr {
self.addr
}
pub fn port(&self) -> u16 {
self.addr.port()
}
pub fn base_url(&self) -> String {
format!("http://127.0.0.1:{}", self.addr.port())
}
pub fn serve(self, router: Router) -> RunningServer {
let addr = self.addr;
let (stop, stopped) = oneshot::channel::<()>();
let task = tokio::spawn(async move {
axum::serve(self.listener, router)
.with_graceful_shutdown(async {
let _ = stopped.await;
})
.await
});
RunningServer {
addr,
stop: Some(stop),
task: Some(task),
}
}
}
pub struct RunningServer {
addr: SocketAddr,
stop: Option<oneshot::Sender<()>>,
task: Option<JoinHandle<std::io::Result<()>>>,
}
impl RunningServer {
pub fn addr(&self) -> SocketAddr {
self.addr
}
pub fn base_url(&self) -> String {
format!("http://127.0.0.1:{}", self.addr.port())
}
pub async fn wait_until_ready(&self) -> Result<(), NativeError> {
let deadline = tokio::time::Instant::now() + READY_TIMEOUT;
loop {
if TcpStream::connect(self.addr).await.is_ok() {
return Ok(());
}
if tokio::time::Instant::now() >= deadline {
return Err(NativeError::new(
"server",
format!(
"the embedded server did not answer on {} within {} seconds",
self.addr,
READY_TIMEOUT.as_secs()
),
));
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
pub async fn shutdown(mut self, grace: Duration) -> Result<(), NativeError> {
if let Some(stop) = self.stop.take() {
let _ = stop.send(());
}
let Some(task) = self.task.take() else {
return Ok(());
};
match tokio::time::timeout(grace, task).await {
Ok(Ok(Ok(()))) => Ok(()),
Ok(Ok(Err(e))) => Err(NativeError::new(
"server",
format!("the embedded server stopped with an error: {e}"),
)),
Ok(Err(e)) => Err(NativeError::new(
"server",
format!("the embedded server task did not finish: {e}"),
)),
Err(_) => Err(NativeError::new(
"server",
format!(
"the embedded server still had work open after {} seconds and was abandoned",
grace.as_secs()
),
)),
}
}
}
impl Drop for RunningServer {
fn drop(&mut self) {
if let Some(stop) = self.stop.take() {
let _ = stop.send(());
}
if let Some(task) = self.task.take() {
task.abort();
}
}
}