From f0dc59aeae3f66d50b6a46215abb9fa6424fffd5 Mon Sep 17 00:00:00 2001
From: Hiroaki Nakamura <hnakamur@gmail.com>
Date: Sun, 3 Mar 2024 09:56:19 +0900
Subject: [PATCH] Add config for grace period and graceful shutdown timeout

---
 pingora-core/src/server/configuration/mod.rs |  6 ++++++
 pingora-core/src/server/mod.rs               | 16 +++++++++++++---
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/pingora-core/src/server/configuration/mod.rs b/pingora-core/src/server/configuration/mod.rs
index 19db553..8f82566 100644
--- a/pingora-core/src/server/configuration/mod.rs
+++ b/pingora-core/src/server/configuration/mod.rs
@@ -60,6 +60,10 @@ pub struct ServerConf {
     /// The path to CA file the SSL library should use. If empty, the default trust store location
     /// defined by the SSL library will be used.
     pub ca_file: Option<String>,
+    /// Grace period in seconds before starting the final step of the graceful shutdown after signaling shutdown.
+    pub grace_period_seconds: Option<u64>,
+    /// Timeout in seconds of the final step for the graceful shutdown.
+    pub graceful_shutdown_timeout_seconds: Option<u64>,
     // These options don't belong here as they are specific to certain services
     pub(crate) client_bind_to_ipv4: Vec<String>,
     pub(crate) client_bind_to_ipv6: Vec<String>,
@@ -86,6 +90,8 @@ impl Default for ServerConf {
             upstream_keepalive_pool_size: 128,
             upstream_connect_offload_threadpools: None,
             upstream_connect_offload_thread_per_pool: None,
+            grace_period_seconds: None,
+            graceful_shutdown_timeout_seconds: None,
         }
     }
 }
diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs
index 4273550..c9721e3 100644
--- a/pingora-core/src/server/mod.rs
+++ b/pingora-core/src/server/mod.rs
@@ -302,15 +302,25 @@ impl Server {
         let shutdown_type = server_runtime.get_handle().block_on(self.main_loop());
 
         if matches!(shutdown_type, ShutdownType::Graceful) {
-            info!("Graceful shutdown: grace period {}s starts", EXIT_TIMEOUT);
-            thread::sleep(Duration::from_secs(EXIT_TIMEOUT));
+            let exit_timeout = self
+                .configuration
+                .as_ref()
+                .grace_period_seconds
+                .unwrap_or(EXIT_TIMEOUT);
+            info!("Graceful shutdown: grace period {}s starts", exit_timeout);
+            thread::sleep(Duration::from_secs(exit_timeout));
             info!("Graceful shutdown: grace period ends");
         }
 
         // Give tokio runtimes time to exit
         let shutdown_timeout = match shutdown_type {
             ShutdownType::Quick => Duration::from_secs(0),
-            ShutdownType::Graceful => Duration::from_secs(5),
+            ShutdownType::Graceful => Duration::from_secs(
+                self.configuration
+                    .as_ref()
+                    .graceful_shutdown_timeout_seconds
+                    .unwrap_or(5),
+            ),
         };
         let shutdowns: Vec<_> = runtimes
             .into_iter()