1 files changed, 28 insertions, 7 deletions
diff --git a/atuin-server/src/lib.rs b/atuin-server/src/lib.rs
index fc4d9e0c..810b7db9 100644
--- a/atuin-server/src/lib.rs
+++ b/atuin-server/src/lib.rs
@@ -1,6 +1,9 @@
 #![forbid(unsafe_code)]
 
-use std::net::{IpAddr, SocketAddr};
+use std::{
+    future::Future,
+    net::{IpAddr, SocketAddr, TcpListener},
+};
 
 use atuin_server_database::Database;
 use axum::Server;
@@ -16,10 +19,15 @@ use tokio::signal;
 
 #[cfg(target_family = "unix")]
 async fn shutdown_signal() {
-    signal::unix::signal(signal::unix::SignalKind::terminate())
-        .expect("failed to register signal handler")
-        .recv()
-        .await;
+    let mut term = signal::unix::signal(signal::unix::SignalKind::terminate())
+        .expect("failed to register signal handler");
+    let mut interrupt = signal::unix::signal(signal::unix::SignalKind::interrupt())
+        .expect("failed to register signal handler");
+
+    tokio::select! {
+        _ = term.recv() => {},
+        _ = interrupt.recv() => {},
+    };
     eprintln!("Shutting down gracefully...");
 }
 
@@ -38,16 +46,29 @@ pub async fn launch<Db: Database>(
     port: u16,
 ) -> Result<()> {
     let host = host.parse::<IpAddr>()?;
+    launch_with_listener::<Db>(
+        settings,
+        TcpListener::bind(SocketAddr::new(host, port)).context("could not connect to socket")?,
+        shutdown_signal(),
+    )
+    .await
+}
 
+pub async fn launch_with_listener<Db: Database>(
+    settings: Settings<Db::Settings>,
+    listener: TcpListener,
+    shutdown: impl Future<Output = ()>,
+) -> Result<()> {
     let db = Db::new(&settings.db_settings)
         .await
         .wrap_err_with(|| format!("failed to connect to db: {:?}", settings.db_settings))?;
 
     let r = router::router(db, settings);
 
-    Server::bind(&SocketAddr::new(host, port))
+    Server::from_tcp(listener)
+        .context("could not launch server")?
         .serve(r.into_make_service())
-        .with_graceful_shutdown(shutdown_signal())
+        .with_graceful_shutdown(shutdown)
         .await?;
 
     Ok(())