From 777be1f723ae3fdb9bee5f225bc5e10296a38581 Mon Sep 17 00:00:00 2001
From: Paul-Nicolas Madelaine <pnm@pnm.tf>
Date: Sat, 30 Sep 2023 14:26:59 +0200
Subject: [PATCH] server: graceful shutdown

---
 typhon/src/lib.rs   | 12 ++++++++++++
 typhon/src/main.rs  | 13 ++++++++++---
 typhon/src/tasks.rs | 15 ++++++++++++++-
 3 files changed, 36 insertions(+), 4 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 239224e6..856e302f 100644
--- a/typhon/src/tasks.rs
+++ b/typhon/src/tasks.rs
@@ -15,7 +15,7 @@ pub struct Tasks<Id> {
     tasks: Mutex<HashMap<Id, TaskHandle>>,
 }
 
-impl<Id: std::cmp::Eq + std::hash::Hash + std::clone::Clone + Send> Tasks<Id> {
+impl<Id: std::cmp::Eq + std::hash::Hash + std::clone::Clone + Send + Sync> Tasks<Id> {
     pub fn new() -> Self {
         Tasks {
             tasks: Mutex::new(HashMap::new()),
@@ -83,4 +83,17 @@ impl<Id: std::cmp::Eq + std::hash::Hash + std::clone::Clone + Send> Tasks<Id> {
             .flatten()
             .is_some()
     }
+
+    pub async fn shutdown(&'static self) {
+        let tasks = self.tasks.lock().await;
+        let mut set = tokio::task::JoinSet::new();
+        for id in tasks.keys() {
+            set.spawn({
+                let id = id.clone();
+                async move { self.wait(&id).await }
+            });
+            self.cancel(id).await;
+        }
+        while let Some(_) = set.join_next().await {}
+    }
 }