From 0e0f830ac6cad404d1e0736c06ad74a1a9ea9b88 Mon Sep 17 00:00:00 2001
From: danielle-tfh <danielle.nagar@toolsforhumanity.com>
Date: Fri, 10 Jan 2025 15:51:20 +0100
Subject: [PATCH] Trigger graceful shutdown between nodes

---
 .../src/helpers/shutdown_handler.rs           |  4 ++++
 iris-mpc/src/bin/server.rs                    | 21 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/iris-mpc-common/src/helpers/shutdown_handler.rs b/iris-mpc-common/src/helpers/shutdown_handler.rs
index e596335d7..24513bac6 100644
--- a/iris-mpc-common/src/helpers/shutdown_handler.rs
+++ b/iris-mpc-common/src/helpers/shutdown_handler.rs
@@ -29,6 +29,10 @@ impl ShutdownHandler {
         self.shutdown_received.load(Ordering::Relaxed)
     }
 
+    pub fn trigger_manual_shutdown(&self) {
+        self.shutdown_received.store(true, Ordering::Relaxed);
+    }
+
     pub async fn wait_for_shutdown_signal(&self) {
         let shutdown_flag = self.shutdown_received.clone();
         tokio::spawn(async move {
diff --git a/iris-mpc/src/bin/server.rs b/iris-mpc/src/bin/server.rs
index 2b1ed4f54..943dd3b63 100644
--- a/iris-mpc/src/bin/server.rs
+++ b/iris-mpc/src/bin/server.rs
@@ -781,12 +781,14 @@ async fn server_main(config: Config) -> eyre::Result<()> {
     struct ReadyProbeResponse {
         image_name: String,
         uuid:       String,
+        shutdown:   bool,
     }
 
     let _health_check_abort = background_tasks.spawn({
         let uuid = uuid::Uuid::new_v4().to_string();
         let ready_probe_response = ReadyProbeResponse {
             image_name: config.image_name.clone(),
+            shutdown: false,
             uuid,
         };
         let serialized_response = serde_json::to_string(&ready_probe_response)
@@ -797,7 +799,16 @@ async fn server_main(config: Config) -> eyre::Result<()> {
             let app = Router::new()
                 .route(
                     "/health",
-                    get(move || async move { serialized_response.clone() }),
+                    get(move || {
+                        if shutdown_handler.is_shutting_down() {
+                            ready_probe_response.shutdown = true;
+                            serialized_response = serde_json::to_string(&ready_probe_response)
+                                .expect("Serialization to JSON to probe response failed");
+                        }
+                        async move {
+                            serialized_response.clone();
+                        }
+                    }),
                 )
                 .route(
                     "/ready",
@@ -863,6 +874,14 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                     .json::<ReadyProbeResponse>()
                     .await
                     .expect("Deserialization of probe response failed");
+                if probe_response.shutdown {
+                    tracing::error!(
+                        "Node {} has starting graceful shutdown. Therefore starting graceful \
+                         shutdown",
+                        host
+                    );
+                    shutdown_handler.trigger_manual_shutdown();
+                }
                 if probe_response.image_name != image_name {
                     // Do not create a panic as we still can continue to process before its
                     // updated