From 05e090789aa7c2f9c623f1ad823875e4deb78f36 Mon Sep 17 00:00:00 2001 From: Kristopher Wuollett Date: Fri, 13 Dec 2024 11:38:14 -0600 Subject: [PATCH 1/4] Refactor Server run to support customizing shutdown signals and behavior. This change supports building a server that can run Pingora without it calling std::process::exit during shutdown. --- pingora-core/src/server/mod.rs | 93 ++++++++++++++++++++++++++++------ 1 file changed, 78 insertions(+), 15 deletions(-) diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs index c056ed273..5c2ed639a 100644 --- a/pingora-core/src/server/mod.rs +++ b/pingora-core/src/server/mod.rs @@ -20,6 +20,7 @@ mod daemon; #[cfg(unix)] pub(crate) mod transfer_fd; +use async_trait::async_trait; #[cfg(unix)] use daemon::daemonize; use log::{debug, error, info, warn}; @@ -59,6 +60,61 @@ pub type ShutdownWatch = watch::Receiver; #[cfg(unix)] pub type ListenFds = Arc>; +pub enum ShutdownSignal { + GracefulUpgrade, + GracefulTerminate, + FastShutdown, +} + +#[async_trait] +pub trait ShutdownSignalWatch { + async fn recv(&self) -> ShutdownSignal; +} + +#[cfg(unix)] +pub struct UnixShutdownSignalWatch; + +#[cfg(unix)] +#[async_trait] +impl ShutdownSignalWatch for UnixShutdownSignalWatch { + async fn recv(&self) -> ShutdownSignal { + let mut graceful_upgrade_signal = unix::signal(unix::SignalKind::quit()).unwrap(); + let mut graceful_terminate_signal = unix::signal(unix::SignalKind::terminate()).unwrap(); + let mut fast_shutdown_signal = unix::signal(unix::SignalKind::interrupt()).unwrap(); + + tokio::select! { + _ = graceful_upgrade_signal.recv() => { + ShutdownSignal::GracefulUpgrade + }, + _ = graceful_terminate_signal.recv() => { + ShutdownSignal::GracefulTerminate + }, + _ = fast_shutdown_signal.recv() => { + ShutdownSignal::FastShutdown + }, + } + } +} + +pub struct RunArgs { + #[cfg(unix)] + pub shutdown_signal: Box, +} + +impl Default for RunArgs { + #[cfg(unix)] + fn default() -> Self { + Self { + shutdown_signal: Box::new(UnixShutdownSignalWatch), + } + } + + #[cfg(windows)] + fn default() -> Self { + Self {} + } +} + /// The server object /// /// This object represents an entire pingora server process which may have multiple independent @@ -87,18 +143,14 @@ pub struct Server { impl Server { #[cfg(unix)] - async fn main_loop(&self) -> ShutdownType { + async fn main_loop(&self, run_args: RunArgs) -> ShutdownType { // waiting for exit signal - // TODO: there should be a signal handling function - let mut graceful_upgrade_signal = unix::signal(unix::SignalKind::quit()).unwrap(); - let mut graceful_terminate_signal = unix::signal(unix::SignalKind::terminate()).unwrap(); - let mut fast_shutdown_signal = unix::signal(unix::SignalKind::interrupt()).unwrap(); - tokio::select! { - _ = fast_shutdown_signal.recv() => { + match run_args.shutdown_signal.recv().await { + ShutdownSignal::FastShutdown => { info!("SIGINT received, exiting"); ShutdownType::Quick - }, - _ = graceful_terminate_signal.recv() => { + } + ShutdownSignal::GracefulTerminate => { // we receive a graceful terminate, all instances are instructed to stop info!("SIGTERM received, gracefully exiting"); // graceful shutdown if there are listening sockets @@ -112,7 +164,7 @@ impl Server { info!("Broadcast graceful shutdown complete"); ShutdownType::Graceful } - _ = graceful_upgrade_signal.recv() => { + ShutdownSignal::GracefulUpgrade => { // TODO: still need to select! on signals in case a fast shutdown is needed // aka: move below to another task and only kick it off here info!("SIGQUIT received, sending socks and gracefully exiting"); @@ -148,7 +200,7 @@ impl Server { info!("No socks to send, shutting down."); ShutdownType::Graceful } - }, + } } } @@ -306,6 +358,17 @@ impl Server { } } + /// Start the server using [Self::run] and default [RunArgs]. + #[allow(unused_mut)] // TODO: May not need to keep mut self in interface + pub fn run_forever(mut self) -> ! { + info!("Server starting"); + + self.run(RunArgs::default()); + + info!("All runtimes exited, exiting now"); + std::process::exit(0) + } + /// Start the server /// /// This function will block forever until the server needs to quit. So this would be the last @@ -313,7 +376,7 @@ impl Server { /// /// Note: this function may fork the process for daemonization, so any additional threads created /// before this function will be lost to any service logic once this function is called. - pub fn run_forever(mut self) -> ! { + pub fn run(mut self, run_args: RunArgs) { info!("Server starting"); let conf = self.configuration.as_ref(); @@ -354,7 +417,9 @@ impl Server { // Only work steal runtime can use block_on() let server_runtime = Server::create_runtime("Server", 1, true); #[cfg(unix)] - let shutdown_type = server_runtime.get_handle().block_on(self.main_loop()); + let shutdown_type = server_runtime + .get_handle() + .block_on(self.main_loop(run_args)); #[cfg(windows)] let shutdown_type = ShutdownType::Graceful; @@ -394,8 +459,6 @@ impl Server { error!("Failed to shutdown runtime: {:?}", e); } } - info!("All runtimes exited, exiting now"); - std::process::exit(0) } fn create_runtime(name: &str, threads: usize, work_steal: bool) -> Runtime { From 11881c52c1e1ac7ae88ca6c9d60cee7779f32b42 Mon Sep 17 00:00:00 2001 From: Kristopher Wuollett Date: Fri, 13 Dec 2024 11:44:33 -0600 Subject: [PATCH 2/4] Fix pingora-core Server formatting to pass ci check. --- pingora-core/src/server/mod.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs index 5c2ed639a..a9b608fc8 100644 --- a/pingora-core/src/server/mod.rs +++ b/pingora-core/src/server/mod.rs @@ -156,7 +156,9 @@ impl Server { // graceful shutdown if there are listening sockets info!("Broadcasting graceful shutdown"); match self.shutdown_watch.send(true) { - Ok(_) => { info!("Graceful shutdown started!"); } + Ok(_) => { + info!("Graceful shutdown started!"); + } Err(e) => { error!("Graceful shutdown broadcast failed: {e}"); } @@ -172,10 +174,10 @@ impl Server { let fds = fds.lock().await; info!("Trying to send socks"); // XXX: this is blocking IO - match fds.send_to_sock( - self.configuration.as_ref().upgrade_sock.as_str()) - { - Ok(_) => {info!("listener sockets sent");}, + match fds.send_to_sock(self.configuration.as_ref().upgrade_sock.as_str()) { + Ok(_) => { + info!("listener sockets sent"); + } Err(e) => { error!("Unable to send listener sockets to new process: {e}"); // sentry log error on fd send failure @@ -187,7 +189,9 @@ impl Server { info!("Broadcasting graceful shutdown"); // gracefully exiting match self.shutdown_watch.send(true) { - Ok(_) => { info!("Graceful shutdown started!"); } + Ok(_) => { + info!("Graceful shutdown started!"); + } Err(e) => { error!("Graceful shutdown broadcast failed: {e}"); // switch to fast shutdown From 02c1ea1288ada6c9f5c0b1380c99da44593e1055 Mon Sep 17 00:00:00 2001 From: Kristopher Wuollett Date: Thu, 2 Jan 2025 13:45:23 -0600 Subject: [PATCH 3/4] Add doc comments for ShutdownSignal --- pingora-core/src/server/mod.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs index a9b608fc8..90a22df2e 100644 --- a/pingora-core/src/server/mod.rs +++ b/pingora-core/src/server/mod.rs @@ -60,9 +60,15 @@ pub type ShutdownWatch = watch::Receiver; #[cfg(unix)] pub type ListenFds = Arc>; +/// The type of shutdown process that has been requested. pub enum ShutdownSignal { + /// Send file descriptors to the new process before starting runtime shutdown with + /// [ServerConf::graceful_shutdown_timeout_seconds] timeout. GracefulUpgrade, + /// Wait for [ServerConf::grace_period_seconds] before starting runtime shutdown with + /// [ServerConf::graceful_shutdown_timeout_seconds] timeout. GracefulTerminate, + /// Shutdown with no timeout for runtime shutdown. FastShutdown, } From 0096b01b5378c52de6ab035de465f01d49cfa3e5 Mon Sep 17 00:00:00 2001 From: Kristopher Wuollett Date: Thu, 2 Jan 2025 13:46:10 -0600 Subject: [PATCH 4/4] Remove unused mut self from run_forever --- pingora-core/src/server/mod.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs index 90a22df2e..fdcd08620 100644 --- a/pingora-core/src/server/mod.rs +++ b/pingora-core/src/server/mod.rs @@ -369,8 +369,7 @@ impl Server { } /// Start the server using [Self::run] and default [RunArgs]. - #[allow(unused_mut)] // TODO: May not need to keep mut self in interface - pub fn run_forever(mut self) -> ! { + pub fn run_forever(self) -> ! { info!("Server starting"); self.run(RunArgs::default());