feat(server): allow graceful shutdown to exit when idle

pingora-core/src/server/mod.rs

@@ -126,6 +126,29 @@ pub type ShutdownWatch = watch::Receiver<bool>;
 #[cfg(unix)]
 pub type ListenFds = Arc<Mutex<Fds>>;
 
+/// Optional process-owned idle wait used to end a graceful shutdown early.
+///
+/// The callback must return `true` when all accepted work has become idle and
+/// `false` when the supplied deadline expires. The owner is responsible for
+/// synchronizing accept quiescence with its activity count before returning
+/// `true`; a zero activity count alone is not sufficient.
+pub type GracefulShutdownIdle = Arc<dyn Fn(Duration) -> bool + Send + Sync>;
+
+/// Wait for graceful shutdown to become idle, retaining the fixed wait when
+/// no idle callback is installed.
+fn wait_for_graceful_shutdown_idle(
+    idle_wait: Option<&GracefulShutdownIdle>,
+    grace_period: Duration,
+) -> bool {
+    idle_wait.map_or_else(
+        || {
+            thread::sleep(grace_period);
+            false
+        },
+        |wait_for_idle| wait_for_idle(grace_period),
+    )
+}
+
 /// The type of shutdown process that has been requested.
 #[derive(Debug)]
 pub enum ShutdownSignal {
@@ -224,11 +247,31 @@ pub struct Server {
     pub configuration: Arc<ServerConf>,
     /// The parser command line options
     pub options: Option<Opt>,
+
+    /// Optional process-owned idle wait for graceful shutdown.
+    graceful_shutdown_idle: Option<GracefulShutdownIdle>,
 }
 
 // TODO: delete the pid when exit
 
 impl Server {
+    /// Install an optional idle wait for graceful termination.
+    ///
+    /// When no wait is installed, graceful shutdown retains the historical
+    /// fixed grace-period behavior. The callback is invoked immediately after
+    /// the shutdown broadcast, without waiting for listeners to stop accepting
+    /// new work, and receives the configured grace period as its maximum wait.
+    /// The callback owner must synchronize accept quiescence with its activity
+    /// count before returning `true`.
+    pub fn set_graceful_shutdown_idle(&mut self, idle_wait: GracefulShutdownIdle) {
+        self.graceful_shutdown_idle = Some(idle_wait);
+    }
+
+    /// Remove a previously installed idle wait and restore the fixed grace period.
+    pub fn clear_graceful_shutdown_idle(&mut self) {
+        self.graceful_shutdown_idle = None;
+    }
+
     /// Acquire a receiver for the server's execution phase.
     ///
     /// The receiver will produce values for each transition.
@@ -479,6 +522,7 @@ impl Server {
             options: opt,
             dependencies: Arc::new(Mutex::new(DependencyGraph::new())),
             bootstrap,
+            graceful_shutdown_idle: None,
         }
     }
 
@@ -527,6 +571,7 @@ impl Server {
             options: opt,
             dependencies: Arc::new(Mutex::new(DependencyGraph::new())),
             bootstrap,
+            graceful_shutdown_idle: None,
         })
     }
 
@@ -808,7 +853,12 @@ impl Server {
                 .grace_period_seconds
                 .unwrap_or(EXIT_TIMEOUT);
             info!("Graceful shutdown: grace period {}s starts", exit_timeout);
-            thread::sleep(Duration::from_secs(exit_timeout));
+            let grace_period = Duration::from_secs(exit_timeout);
+            if wait_for_graceful_shutdown_idle(self.graceful_shutdown_idle.as_ref(), grace_period) {
+                info!("Graceful shutdown: accepted work became idle");
+            } else {
+                info!("Graceful shutdown: grace period expired");
+            }
             info!("Graceful shutdown: grace period ends");
         }
 
@@ -894,3 +944,29 @@ mod listen_address_tests {
         assert_eq!(server.collect_listen_addresses(), None);
     }
 }
+
+#[cfg(test)]
+mod graceful_shutdown_tests {
+    use super::*;
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    #[test]
+    fn idle_callback_short_circuits_grace_period() {
+        let called = Arc::new(AtomicBool::new(false));
+        let called_by_waiter = Arc::clone(&called);
+        let idle_wait: GracefulShutdownIdle = Arc::new(move |grace_period| {
+            assert_eq!(grace_period, Duration::from_secs(30));
+            called_by_waiter.store(true, Ordering::Relaxed);
+            true
+        });
+
+        assert!(
+            wait_for_graceful_shutdown_idle(Some(&idle_wait), Duration::from_secs(30)),
+            "the idle callback should report that work is complete"
+        );
+        assert!(
+            called.load(Ordering::Relaxed),
+            "the idle callback was not invoked"
+        );
+    }
+}

pingora-core/tests/server_idle_gracefulshutdown.rs

@@ -0,0 +1,109 @@
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// NOTE: This test sends a shutdown signal to itself, so it needs to be an
+// isolated test to prevent concurrency.
+
+use pingora_core::server::{
+    configuration::ServerConf, ExecutionPhase, GracefulShutdownIdle, RunArgs, Server,
+};
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tokio::sync::broadcast;
+
+const GRACE_PERIOD: Duration = Duration::from_secs(10);
+const PHASE_TIMEOUT: Duration = Duration::from_secs(5);
+
+fn next_phase(phase: &mut broadcast::Receiver<ExecutionPhase>) -> ExecutionPhase {
+    let deadline = Instant::now() + PHASE_TIMEOUT;
+    loop {
+        match phase.try_recv() {
+            Ok(value) => return value,
+            Err(broadcast::error::TryRecvError::Closed) => panic!("execution phase channel closed"),
+            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
+                panic!("execution phase receiver lagged by {skipped} messages")
+            }
+            Err(broadcast::error::TryRecvError::Empty) if Instant::now() >= deadline => {
+                panic!("timed out waiting for execution phase")
+            }
+            Err(broadcast::error::TryRecvError::Empty) => {
+                std::thread::sleep(Duration::from_millis(10))
+            }
+        }
+    }
+}
+
+#[test]
+fn server_uses_idle_callback_during_graceful_shutdown() {
+    let callback_grace_period = Arc::new(AtomicU64::new(0));
+    let callback_grace_period_writer = Arc::clone(&callback_grace_period);
+    let idle_wait: GracefulShutdownIdle = Arc::new(move |grace_period| {
+        callback_grace_period_writer.store(grace_period.as_secs(), Ordering::SeqCst);
+        true
+    });
+    let conf = ServerConf {
+        grace_period_seconds: Some(GRACE_PERIOD.as_secs()),
+        ..Default::default()
+    };
+    let mut server = Server::new_with_opt_and_conf(None, conf);
+    server.set_graceful_shutdown_idle(idle_wait);
+    let mut phase = server.watch_execution_phase();
+
+    let join = std::thread::spawn(move || {
+        server.bootstrap();
+        server.run(RunArgs::default());
+    });
+
+    assert!(matches!(next_phase(&mut phase), ExecutionPhase::Bootstrap));
+    assert!(matches!(
+        next_phase(&mut phase),
+        ExecutionPhase::BootstrapComplete
+    ));
+    assert!(matches!(next_phase(&mut phase), ExecutionPhase::Running));
+    std::thread::sleep(Duration::from_millis(500));
+
+    unsafe {
+        libc::raise(libc::SIGTERM);
+    }
+
+    assert!(matches!(
+        next_phase(&mut phase),
+        ExecutionPhase::GracefulTerminate
+    ));
+    assert!(matches!(
+        next_phase(&mut phase),
+        ExecutionPhase::ShutdownStarted
+    ));
+    assert!(matches!(
+        next_phase(&mut phase),
+        ExecutionPhase::ShutdownGracePeriod
+    ));
+    let grace_period_started = Instant::now();
+    assert!(matches!(
+        phase.blocking_recv().unwrap(),
+        ExecutionPhase::ShutdownRuntimes
+    ));
+    let callback_wait = grace_period_started.elapsed();
+
+    join.join().unwrap();
+
+    assert_eq!(
+        callback_grace_period.load(Ordering::SeqCst),
+        GRACE_PERIOD.as_secs(),
+        "the idle callback should receive the configured grace period"
+    );
+    assert!(
+        callback_wait < GRACE_PERIOD,
+        "the idle callback should allow shutdown before the full grace period"
+    );
+}