From 8e9ea437dcc24c008f60ea506f32be7160ff12b1 Mon Sep 17 00:00:00 2001
From: Paul-Nicolas Madelaine <pnm@pnm.tf>
Date: Fri, 29 Sep 2023 23:49:08 +0200
Subject: [PATCH] server: graceful shutdown

---
 typhon/src/lib.rs   | 12 ++++++++++++
 typhon/src/main.rs  | 13 ++++++++++---
 typhon/src/tasks.rs | 14 ++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/typhon/src/lib.rs b/typhon/src/lib.rs
index c6aa36e2..d9bc9e8b 100644
--- a/typhon/src/lib.rs
+++ b/typhon/src/lib.rs
@@ -284,3 +284,15 @@ pub fn log_event(event: Event) {
         LISTENERS.lock().await.log(event);
     });
 }
+
+pub async fn shutdown() {
+    let res = tokio::signal::ctrl_c().await;
+    eprintln!("Typhon is shutting down...");
+    let _ = res.map_err(|e| log::error!("{}", e));
+    tokio::join!(
+        EVALUATIONS.shutdown(),
+        JOBS_BUILD.shutdown(),
+        JOBS_POST.shutdown(),
+        JOBS_PRE.shutdown(),
+    );
+}
diff --git a/typhon/src/main.rs b/typhon/src/main.rs
index 232b31fd..71975ce7 100644
--- a/typhon/src/main.rs
+++ b/typhon/src/main.rs
@@ -30,8 +30,15 @@ async fn main() -> std::io::Result<()> {
         .expect("failed to run migrations");
 
     // Run actix server
-    HttpServer::new(|| App::new().configure(typhon::api::config))
+    let actix = HttpServer::new(|| App::new().configure(typhon::api::config))
         .bind(("127.0.0.1", 8000))?
-        .run()
-        .await
+        .run();
+
+    // Graceful shutdown
+    tokio::select! {
+        _ = actix => panic!(),
+        _ = typhon::shutdown() => eprintln!("Good bye!"),
+    }
+
+    Ok(())
 }
diff --git a/typhon/src/tasks.rs b/typhon/src/tasks.rs
index 92b0b488..4de9ff82 100644
--- a/typhon/src/tasks.rs
+++ b/typhon/src/tasks.rs
@@ -91,4 +91,18 @@ impl<Id: std::cmp::Eq + std::hash::Hash + std::clone::Clone + Send> Tasks<Id> {
             })
             .unwrap_or(false)
     }
+
+    pub async fn keys(&self) -> Vec<Id> {
+        let tasks = self.tasks.lock().await;
+        tasks.keys().cloned().collect()
+    }
+
+    pub async fn shutdown(&'static self) {
+        let mut set = tokio::task::JoinSet::new();
+        let keys = self.keys().await;
+        for id in keys {
+            set.spawn(self.cancel(id));
+        }
+        while let Some(_) = set.join_next().await {}
+    }
 }