From 2bde4bfa0127cc6fa9396f091e60fbebb2c210e4 Mon Sep 17 00:00:00 2001
From: segfault_magnet <ahmed.sagdati.ets@gmail.com>
Date: Thu, 2 May 2024 18:06:39 +0200
Subject: [PATCH] graceful close on shutdown

---
 src/adapters/storage/postgresql.rs | 6 ++++++
 src/main.rs                        | 3 ++-
 src/setup/helpers.rs               | 2 ++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/adapters/storage/postgresql.rs b/src/adapters/storage/postgresql.rs
index d385c013..fe0e2ed5 100644
--- a/src/adapters/storage/postgresql.rs
+++ b/src/adapters/storage/postgresql.rs
@@ -47,6 +47,12 @@ impl Postgres {
         Ok(Self { connection_pool })
     }
 
+    /// Close only when shutting down the application. Will close the connection pool even if it is
+    /// shared.
+    pub async fn close(self) {
+        self.connection_pool.close().await;
+    }
+
     pub async fn migrate(&self) -> Result<()> {
         sqlx::migrate!().run(&self.connection_pool).await?;
         Ok(())
diff --git a/src/main.rs b/src/main.rs
index b7781205..a7fbe6ce 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -62,7 +62,7 @@ async fn main() -> Result<()> {
     launch_api_server(
         &config,
         metrics_registry,
-        storage,
+        storage.clone(),
         fuel_health_check,
         eth_health_check,
     )
@@ -74,6 +74,7 @@ async fn main() -> Result<()> {
         wallet_balance_tracker_handle,
         committer_handle,
         listener_handle,
+        storage,
     )
     .await
 }
diff --git a/src/setup/helpers.rs b/src/setup/helpers.rs
index ea84f46d..38aeae74 100644
--- a/src/setup/helpers.rs
+++ b/src/setup/helpers.rs
@@ -201,6 +201,7 @@ pub async fn shut_down(
     wallet_balance_tracker_handle: JoinHandle<()>,
     committer_handle: JoinHandle<()>,
     listener_handle: JoinHandle<()>,
+    storage: Postgres,
 ) -> Result<()> {
     cancel_token.cancel();
 
@@ -213,5 +214,6 @@ pub async fn shut_down(
         handle.await?;
     }
 
+    storage.close().await;
     Ok(())
 }