Graceful shutdown

Cargo.toml

@@ -4,7 +4,8 @@ members = [
   "api",
   "ingestion",
   "integration_tests",
-  "rove_connector", "util",
+  "rove_connector",
+  "util",
 ]
 resolver = "2"
 

api/Cargo.toml

@@ -8,6 +8,7 @@ axum.workspace = true
 bb8.workspace = true
 bb8-postgres.workspace = true
 chrono.workspace = true
+util = { path = "../util" }
 postgres-types.workspace = true
 serde.workspace = true
 tokio.workspace = true

api/src/lib.rs

@@ -135,5 +135,8 @@ pub async fn run(pool: PgConnectionPool) {
 
     // run it with hyper on localhost:3000
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
-    axum::serve(listener, app).await.unwrap();
+    axum::serve(listener, app)
+        .with_graceful_shutdown(::util::shutdown_signal())
+        .await
+        .unwrap();
 }

ingestion/Cargo.toml

@@ -16,6 +16,7 @@ bb8-postgres.workspace = true
 bytes.workspace = true
 chrono.workspace = true
 chronoutil.workspace = true
+util = { path = "../util" }
 csv.workspace = true
 futures.workspace = true
 kafka.workspace = true

ingestion/src/lib.rs

@@ -452,7 +452,9 @@ pub async fn run(
 
     // run our app with hyper, listening globally on port 3001
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3001").await?;
-    axum::serve(listener, app).await?;
+    axum::serve(listener, app)
+        .with_graceful_shutdown(util::shutdown_signal())
+        .await?;
 
     Ok(())
 }

util/Cargo.toml

@@ -8,3 +8,4 @@ name = "generate_partition_queries"
 
 [dependencies]
 chrono.workspace = true
+tokio.workspace = true

util/src/lib.rs

@@ -1 +1,27 @@
+use tokio::signal;
 
+/// Returns a Future that completes once a signal to shutdown the service is caught.
+pub async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c() // aka. SIGINT on Unix
+            .await
+            .expect("failed to install Ctrl+C (SIGINT) handler");
+    };
+
+    #[cfg(unix)]
+    let sigterm = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler for SIGTERM")
+            .recv()
+            .await;
+    };
+    #[cfg(not(unix))]
+    let sigterm = std::future::pending::<()>();
+
+    // TODO: add more signals that should result in a shutdown?
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = sigterm => {},
+    }
+}