From 2c579489a139dd920e7a4152d9cf1f8bce945cac Mon Sep 17 00:00:00 2001
From: Piotr Heilman <1212808+piohei@users.noreply.github.com>
Date: Wed, 27 Nov 2024 15:43:41 +0100
Subject: [PATCH] Fix shutdown. (#822)

---
 src/main.rs                |  2 ++
 src/shutdown.rs            |  1 -
 src/utils/mod.rs           | 10 ++++++++--
 tests/graceful_shutdown.rs | 11 +++++++++--
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index c53d5d97..922d85f3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -57,6 +57,8 @@ async fn sequencer_app(args: Args) -> anyhow::Result<()> {
     tracing::info!("Shutting down");
     shutdown.shutdown();
 
+    shutdown.await_shutdown_complete().await;
+
     Ok(())
 }
 
diff --git a/src/shutdown.rs b/src/shutdown.rs
index b7d6fa5c..4695a982 100644
--- a/src/shutdown.rs
+++ b/src/shutdown.rs
@@ -102,7 +102,6 @@ impl Shutdown {
                 _ = self.await_shutdown_complete() => {
                     let elapsed = start.elapsed();
                     info!("shutdown channel closed, gracefully shut down in {:?}", elapsed);
-                    std::process::exit(0);
                 },
                 _ = tokio::time::sleep(timeout) => {
                     error!("shutdown monitor timed out");
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index b6db4dbe..57dd408b 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -223,7 +223,13 @@ mod tests {
         tokio::time::sleep(Duration::from_millis(50)).await;
         assert!(handle.is_finished(), "Task should be finished");
 
-        tokio::time::sleep(Duration::from_millis(100)).await;
-        panic!("The process should have exited already");
+        select! {
+                _ = shutdown.await_shutdown_complete() => {
+                    Ok(())
+                },
+                _ = tokio::time::sleep(Duration::from_millis(100)) => {
+                    panic!("The process should have exited already");
+                }
+        }
     }
 }
diff --git a/tests/graceful_shutdown.rs b/tests/graceful_shutdown.rs
index 36a3b63e..36263fdc 100644
--- a/tests/graceful_shutdown.rs
+++ b/tests/graceful_shutdown.rs
@@ -3,6 +3,7 @@
 mod common;
 
 use common::prelude::*;
+use tokio::select;
 
 const IDLE_TIME: u64 = 5;
 
@@ -64,6 +65,12 @@ async fn graceful_shutdown(offchain_mode_enabled: bool) -> anyhow::Result<()> {
     tokio::time::sleep(Duration::from_secs(IDLE_TIME)).await;
     shutdown.shutdown();
 
-    tokio::time::sleep(Duration::from_secs(5)).await;
-    panic!("error: process took longer than 5 seconds to shutdown");
+    select! {
+            _ = shutdown.await_shutdown_complete() => {
+                Ok(())
+            },
+            _ = tokio::time::sleep(Duration::from_secs(5)) => {
+                panic!("error: process took longer than 5 seconds to shutdown");
+            }
+    }
 }