Trigger graceful shutdown between nodes

worldcoin · Jan 10, 2025 · 0e0f830 · 0e0f830
1 parent 7d6de6f
commit 0e0f830
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git 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
@@ -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