From 6e420f70f66443d3707f0c79e76881bc0b6809bd Mon Sep 17 00:00:00 2001 From: Cody Carlsen Date: Tue, 3 Nov 2026 09:19:09 -0800 Subject: [PATCH 1/2] Keep opted-in listening services accepting through graceful terminate - Readiness and metrics endpoints stop answering once SIGTERM closes every listener, so orchestrators see a dead process during drain (#505) - Add Service::keep_listening_on_terminate(&Server) behind the default-off keep_listening_on_terminate feature (pingora-core, forwarded by pingora-proxy and pingora); graceful upgrade still closes the listeners - Gate the method itself: a no-op when the feature is off would compile and still refuse readiness probes on SIGTERM - Document the SIGTERM/SIGQUIT behavior and add a proxy example with a draining readiness endpoint - Add unit and real-server tests, plus a CI step that lints, tests and builds the example with the feature enabled --- .github/workflows/build.yml | 9 + docs/user_guide/start_stop.md | 7 + pingora-core/Cargo.toml | 1 + pingora-core/src/services/listening.rs | 39 +- .../src/services/listening/keep_listening.rs | 87 ++++ .../listening/keep_listening_tests.rs | 293 +++++++++++++ .../tests/test_keep_listening_on_terminate.rs | 384 ++++++++++++++++++ pingora-proxy/Cargo.toml | 5 + .../examples/keep_listening_on_terminate.rs | 122 ++++++ pingora/Cargo.toml | 10 +- 10 files changed, 955 insertions(+), 2 deletions(-) create mode 100644 pingora-core/src/services/listening/keep_listening.rs create mode 100644 pingora-core/src/services/listening/keep_listening_tests.rs create mode 100644 pingora-core/tests/test_keep_listening_on_terminate.rs create mode 100644 pingora-proxy/examples/keep_listening_on_terminate.rs diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 69c41ce2e..a4a3ac6a6 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -90,6 +90,15 @@ jobs: cargo test --verbose --lib --bins --tests --no-fail-fast -- --skip connectors::http::v2::tests::test_h2_late_frames_on_cancelled_streams_keep_connection + - name: Test keep listening on terminate + if: matrix.toolchain == '1.97.1' + run: | + cargo clippy --package pingora-core --all-targets --features keep_listening_on_terminate -- --allow=unknown-lints --deny=warnings + cargo test --package pingora-core --features keep_listening_on_terminate \ + --lib --test test_keep_listening_on_terminate --no-fail-fast -- \ + --skip connectors::http::v2::tests::test_h2_late_frames_on_cancelled_streams_keep_connection + cargo build --package pingora-proxy --example keep_listening_on_terminate --features keep_listening_on_terminate + # Need to run doc tests separately. # (https://github.com/rust-lang/cargo/issues/6669) - name: Run cargo doc test diff --git a/docs/user_guide/start_stop.md b/docs/user_guide/start_stop.md index 6e26981a4..5b6f032b0 100644 --- a/docs/user_guide/start_stop.md +++ b/docs/user_guide/start_stop.md @@ -23,5 +23,12 @@ Upon receiving SIGINT (ctrl + c), the server will exit immediately with no delay ### SIGTERM: graceful shutdown Upon receiving SIGTERM, the server will notify all its services to shutdown, wait for some preconfigured time and then exit. This behavior gives requests a grace period to finish. +Listening services stop accepting new connections when shutdown starts. A service that must stay reachable during the grace period, such as a readiness or metrics endpoint, can opt in to keep accepting with `Service::keep_listening_on_terminate()`, behind the `keep_listening_on_terminate` cargo feature ([example](../../pingora-proxy/examples/keep_listening_on_terminate.rs)). The application still reports draining itself, for example from a background service's `ShutdownWatch`. + +| Signal | Default listener | `keep_listening_on_terminate` | +| ------ | ---------------- | ----------------------------- | +| SIGTERM | Stops accepting | Accepts until the grace period ends | +| SIGQUIT | Transfers its socket, then stops accepting | Same as default | + ### SIGQUIT: graceful upgrade Similar to SIGTERM, but the server will also transfer all its listening sockets to a new Pingora server so that there is no downtime during the upgrade. See the [graceful upgrade](graceful.md) section for more details. diff --git a/pingora-core/Cargo.toml b/pingora-core/Cargo.toml index a8af5b70a..3e867a04a 100644 --- a/pingora-core/Cargo.toml +++ b/pingora-core/Cargo.toml @@ -128,3 +128,4 @@ openssl_derived = ["any_tls"] any_tls = ["dep:tokio-util"] sentry = ["dep:sentry"] connection_filter = [] +keep_listening_on_terminate = [] diff --git a/pingora-core/src/services/listening.rs b/pingora-core/src/services/listening.rs index f041e303b..2c14daa0f 100644 --- a/pingora-core/src/services/listening.rs +++ b/pingora-core/src/services/listening.rs @@ -42,6 +42,13 @@ use std::marker::PhantomData; use std::sync::Arc; use std::time::Duration; +#[cfg(feature = "keep_listening_on_terminate")] +mod keep_listening; +#[cfg(all(test, feature = "keep_listening_on_terminate"))] +mod keep_listening_tests; +#[cfg(feature = "keep_listening_on_terminate")] +use keep_listening::TerminatePhases; + /// Override the runtime options used to run a listening service. pub type RuntimeOptsOverride = Arc Option + Send + Sync>; @@ -58,6 +65,9 @@ where pub threads: Option, runtime_opts_override: Option, runtime_group: Option, + /// Set by [`Self::keep_listening_on_terminate`]. + #[cfg(feature = "keep_listening_on_terminate")] + terminate_phases: Option>, #[cfg(feature = "connection_filter")] connection_filter: Arc, } @@ -76,6 +86,8 @@ where threads: None, runtime_opts_override: None, runtime_group: None, + #[cfg(feature = "keep_listening_on_terminate")] + terminate_phases: None, #[cfg(feature = "connection_filter")] connection_filter: Arc::new(AcceptAllFilter), } @@ -95,6 +107,8 @@ where threads: None, runtime_opts_override: None, runtime_group: None, + #[cfg(feature = "keep_listening_on_terminate")] + terminate_phases: None, #[cfg(feature = "connection_filter")] connection_filter: Arc::new(AcceptAllFilter), } @@ -245,6 +259,7 @@ where app_logic: Arc, mut stack: TransportStack, mut shutdown: ShutdownWatch, + #[cfg(feature = "keep_listening_on_terminate")] phases: Option>, ) { // the accept loop, until the system is shutting down loop { @@ -257,6 +272,17 @@ where // happen in the initial read continue; } + #[cfg(feature = "keep_listening_on_terminate")] + if phases + .as_ref() + .is_some_and(|phases| !phases.is_graceful_upgrade()) + { + info!( + "Graceful terminate: {} keeps listening until runtime shutdown", + stack.as_str() + ); + continue; + } info!("Shutting down {}", stack.as_str()); break; } @@ -343,6 +369,8 @@ where .take() .expect("can only start_service() once"); let app_logic = Arc::new(app_logic); + #[cfg(feature = "keep_listening_on_terminate")] + let terminate_phases = self.terminate_phases.take(); let mut handlers = Vec::new(); @@ -351,9 +379,18 @@ where let shutdown = shutdown.clone(); let my_app_logic = app_logic.clone(); let endpoint = endpoint.clone(); + #[cfg(feature = "keep_listening_on_terminate")] + let phases = terminate_phases.clone(); let jh = runtime.spawn(async move { - Self::run_endpoint(my_app_logic, endpoint, shutdown).await; + Self::run_endpoint( + my_app_logic, + endpoint, + shutdown, + #[cfg(feature = "keep_listening_on_terminate")] + phases, + ) + .await; }); handlers.push(jh); diff --git a/pingora-core/src/services/listening/keep_listening.rs b/pingora-core/src/services/listening/keep_listening.rs new file mode 100644 index 000000000..694959561 --- /dev/null +++ b/pingora-core/src/services/listening/keep_listening.rs @@ -0,0 +1,87 @@ +// Copyright 2026 Cloudflare, Inc. +// +// 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. + +use super::Service; +use crate::protocols::http::custom::server::Session as CustomServerSession; +use crate::server::{ExecutionPhase, Server}; +use parking_lot::Mutex; +use std::sync::Arc; +use tokio::sync::broadcast::{self, error::TryRecvError}; + +impl Service +where + DS: CustomServerSession, +{ + /// Keep accepting connections through a graceful terminate (SIGTERM) until `server` shuts its + /// runtimes down, instead of closing the listeners when shutdown starts. + /// + /// Meant for endpoints that must stay reachable while the server drains, such as readiness + /// or metrics. The listeners still close on a graceful upgrade (SIGQUIT), so the new process + /// takes over the transferred sockets. `server` must be the server this service is added to. + /// + /// Only accepting is kept. The application still owns its drain behavior: + /// - On terminate this bypasses the normal cleanup path: runtime shutdown drops the + /// listeners, so [`ServerApp::cleanup`](crate::apps::ServerApp::cleanup) is not called, and + /// an HTTP proxy's `Session::is_process_shutting_down()` stays `false`. Observe shutdown + /// separately, e.g. through a background service's + /// [`ShutdownWatch`](crate::server::ShutdownWatch). + /// - Keepalive, response status, and readiness reporting are unchanged. + /// + /// A service that serves regular traffic should not opt in: it would keep admitting work that + /// runtime shutdown then drops. + pub fn keep_listening_on_terminate(&mut self, server: &Server) { + self.keep_listening_with(server.watch_execution_phase()); + } + + pub(super) fn keep_listening_with(&mut self, phases: broadcast::Receiver) { + self.terminate_phases = Some(Arc::new(TerminatePhases(Mutex::new(Decision { + phases, + upgrade: None, + })))); + } +} + +/// Shared by a service's endpoints so they all make the same shutdown decision. +pub(super) struct TerminatePhases(Mutex); + +struct Decision { + phases: broadcast::Receiver, + upgrade: Option, +} + +impl TerminatePhases { + /// Whether the shutdown just broadcast is a graceful upgrade. + /// + /// The server sends the upgrade phases before it broadcasts shutdown, so they are already + /// queued. Lost phases or a closed channel count as an upgrade, the default close. + pub(super) fn is_graceful_upgrade(&self) -> bool { + let mut decision = self.0.lock(); + if let Some(upgrade) = decision.upgrade { + return upgrade; + } + let upgrade = loop { + match decision.phases.try_recv() { + Ok( + ExecutionPhase::GracefulUpgradeTransferringFds + | ExecutionPhase::GracefulUpgradeCloseTimeout, + ) + | Err(TryRecvError::Lagged(_) | TryRecvError::Closed) => break true, + Ok(_) => {} + Err(TryRecvError::Empty) => break false, + } + }; + decision.upgrade = Some(upgrade); + upgrade + } +} diff --git a/pingora-core/src/services/listening/keep_listening_tests.rs b/pingora-core/src/services/listening/keep_listening_tests.rs new file mode 100644 index 000000000..fe9661640 --- /dev/null +++ b/pingora-core/src/services/listening/keep_listening_tests.rs @@ -0,0 +1,293 @@ +// Copyright 2026 Cloudflare, Inc. +// +// 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. + +//! Drives listener shutdown directly, sending phases and the shutdown broadcast in the order +//! `Server` does. Sockets are bound up front and handed over through the fd table. + +#![cfg(unix)] + +use super::*; +use crate::server::{ExecutionPhase, Fds}; +use std::os::fd::IntoRawFd; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio::sync::{broadcast, watch}; +use tokio::task::JoinHandle; +use tokio::time::{sleep, timeout}; + +const REPLY: &[u8] = b"ok"; +const DEADLINE: Duration = Duration::from_secs(5); + +struct Reply; + +#[async_trait] +impl ServerApp for Reply { + async fn process_new( + self: &Arc, + mut session: Stream, + _shutdown: &ShutdownWatch, + ) -> Option { + session.write_all(REPLY).await.ok(); + session.shutdown().await.ok(); + None + } +} + +/// Binds `count` listeners and returns their addresses with an fd table holding them. +fn bind(count: usize) -> (Vec, ListenFds) { + let mut fds = Fds::new(); + let addrs = (0..count) + .map(|_| { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + fds.add(addr.clone(), listener.into_raw_fd()); + addr + }) + .collect(); + (addrs, Arc::new(parking_lot::Mutex::new(fds))) +} + +async fn accepts(addr: &str) -> bool { + let Ok(Ok(mut stream)) = timeout(DEADLINE, TcpStream::connect(addr)).await else { + return false; + }; + let mut reply = Vec::new(); + matches!( + timeout(DEADLINE, stream.read_to_end(&mut reply)).await, + Ok(Ok(_)) + ) && reply == REPLY +} + +/// Refused, not merely unanswered: a stalled accept loop keeps its socket open. +async fn refuses(addr: &str) -> bool { + matches!( + timeout(DEADLINE, TcpStream::connect(addr)).await, + Ok(Err(e)) if e.kind() == std::io::ErrorKind::ConnectionRefused + ) +} + +async fn eventually_refuses(addr: &str) { + timeout(DEADLINE, async { + while !refuses(addr).await { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{addr} was never refused")); +} + +struct Setup { + keep_listening: bool, + endpoints: usize, + listeners_per_fd: usize, + phase_capacity: usize, + phases_before_start: Vec, +} + +impl Default for Setup { + fn default() -> Self { + Setup { + keep_listening: true, + endpoints: 1, + listeners_per_fd: 1, + phase_capacity: 16, + phases_before_start: Vec::new(), + } + } +} + +struct Running { + addrs: Vec, + /// A listener that does not opt in, sharing the shutdown broadcast. + witness: String, + shutdown: watch::Sender, + /// `None` once the phase channel is closed. + phases: Option>, + tasks: Vec>, +} + +fn spawn( + mut service: Service, + fds: ListenFds, + shutdown: watch::Receiver, + listeners_per_fd: usize, +) -> JoinHandle<()> { + tokio::spawn(async move { + service + .start_service(Some(fds), shutdown, listeners_per_fd) + .await; + }) +} + +async fn start(setup: Setup) -> Running { + let (addrs, fds) = bind(setup.endpoints); + let (witness_addrs, witness_fds) = bind(1); + let (phases, phase_rx) = broadcast::channel(setup.phase_capacity); + let (shutdown, shutdown_rx) = watch::channel(false); + + let mut service = Service::new("kept".to_string(), Reply); + for addr in &addrs { + service.add_tcp(addr); + } + if setup.keep_listening { + service.keep_listening_with(phase_rx); + } + let mut witness = Service::new("witness".to_string(), Reply); + witness.add_tcp(&witness_addrs[0]); + for phase in setup.phases_before_start { + phases.send(phase).unwrap(); + } + + let tasks = vec![ + spawn(service, fds, shutdown_rx.clone(), setup.listeners_per_fd), + spawn(witness, witness_fds, shutdown_rx, 1), + ]; + for addr in addrs.iter().chain(&witness_addrs) { + assert!(accepts(addr).await, "{addr} never started accepting"); + } + Running { + addrs, + witness: witness_addrs[0].clone(), + shutdown, + phases: Some(phases), + tasks, + } +} + +impl Running { + fn send(&self, phase: ExecutionPhase) { + if let Some(phases) = &self.phases { + phases.send(phase).ok(); + } + } + + /// Mirrors `Server`: the upgrade phases precede the shutdown broadcast. + fn upgrade(&self) { + for phase in [ + ExecutionPhase::GracefulUpgradeTransferringFds, + ExecutionPhase::GracefulUpgradeCloseTimeout, + ] { + self.send(phase); + } + self.shutdown.send(true).unwrap(); + } + + /// Mirrors `Server`: the terminate phase follows the shutdown broadcast. + fn terminate(&self) { + self.shutdown.send(true).unwrap(); + self.send(ExecutionPhase::GracefulTerminate); + } + + async fn assert_keeps_accepting(self) { + eventually_refuses(&self.witness).await; + for addr in &self.addrs { + assert!(accepts(addr).await, "{addr} stopped accepting"); + } + assert!(!self.tasks[0].is_finished()); + self.tasks[0].abort(); + } + + async fn assert_stopped(self) { + for addr in &self.addrs { + eventually_refuses(addr).await; + } + for task in self.tasks { + timeout(DEADLINE, task) + .await + .expect("service kept running") + .unwrap(); + } + } +} + +#[tokio::test] +async fn graceful_terminate_keeps_accepting() { + let running = start(Setup::default()).await; + running.terminate(); + running.assert_keeps_accepting().await; +} + +#[tokio::test] +async fn graceful_upgrade_closes_listener() { + let running = start(Setup::default()).await; + running.upgrade(); + running.assert_stopped().await; +} + +#[tokio::test] +async fn graceful_terminate_closes_listener_without_opt_in() { + let running = start(Setup { + keep_listening: false, + ..Default::default() + }) + .await; + running.terminate(); + running.assert_stopped().await; +} + +/// A service that waits on a dependency can start after the upgrade began. +#[tokio::test] +async fn upgrade_reported_before_start_closes_listener() { + let running = start(Setup { + phases_before_start: vec![ExecutionPhase::GracefulUpgradeTransferringFds], + ..Default::default() + }) + .await; + running.shutdown.send(true).unwrap(); + running.assert_stopped().await; +} + +#[tokio::test] +async fn lagged_phases_close_listener() { + let running = start(Setup { + phase_capacity: 1, + ..Default::default() + }) + .await; + running.send(ExecutionPhase::Running); + running.send(ExecutionPhase::Running); + running.terminate(); + running.assert_stopped().await; +} + +#[tokio::test] +async fn closed_phases_close_listener() { + let mut running = start(Setup::default()).await; + running.phases = None; + running.shutdown.send(true).unwrap(); + running.assert_stopped().await; +} + +const CONCURRENT: Setup = Setup { + keep_listening: true, + endpoints: 3, + listeners_per_fd: 4, + phase_capacity: 16, + phases_before_start: Vec::new(), +}; + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_listeners_keep_accepting_on_terminate() { + let running = start(CONCURRENT).await; + running.terminate(); + running.assert_keeps_accepting().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_listeners_close_on_upgrade() { + let running = start(CONCURRENT).await; + running.upgrade(); + running.assert_stopped().await; +} diff --git a/pingora-core/tests/test_keep_listening_on_terminate.rs b/pingora-core/tests/test_keep_listening_on_terminate.rs new file mode 100644 index 000000000..7891f3703 --- /dev/null +++ b/pingora-core/tests/test_keep_listening_on_terminate.rs @@ -0,0 +1,384 @@ +// Copyright 2026 Cloudflare, Inc. +// +// 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. + +//! Runs a real `Server` with an injected shutdown signal, so its own phase and shutdown ordering +//! decides each listener's fate. Sockets are bound up front and handed to the server's fd table. + +#![cfg(all(unix, feature = "keep_listening_on_terminate"))] + +use async_trait::async_trait; +use http::{Response, StatusCode}; +use pingora_core::apps::http_app::{HttpServer, ServeHttp}; +use pingora_core::apps::ServerApp; +use pingora_core::protocols::http::ServerSession; +use pingora_core::protocols::Stream; +use pingora_core::server::configuration::ServerConf; +use pingora_core::server::{ + ExecutionPhase, ListenFds, RunArgs, Server, ShutdownSignal, ShutdownSignalWatch, ShutdownWatch, +}; +use pingora_core::services::background::{background_service, BackgroundService}; +use pingora_core::services::listening::Service as ListeningService; +use pingora_core::services::Service; +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::os::fd::{IntoRawFd, RawFd}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, Instant}; +use tokio::io::AsyncWriteExt; +use tokio::sync::{broadcast, mpsc, Mutex}; + +const DEADLINE: Duration = Duration::from_secs(15); +const IO_TIMEOUT: Duration = Duration::from_secs(1); + +/// Delivers shutdown signals sent by the test instead of process signals. +struct InjectedSignal(Mutex>); + +#[async_trait] +impl ShutdownSignalWatch for InjectedSignal { + async fn recv(&self) -> ShutdownSignal { + self.0 + .lock() + .await + .recv() + .await + .unwrap_or(ShutdownSignal::FastShutdown) + } +} + +/// Adds its pre-bound socket to the server's fd table, so the listener reuses it and a graceful +/// upgrade transfers it. +struct PreBound { + inner: S, + addr: String, + fd: Option, +} + +fn pre_bound(inner: impl FnOnce(&str) -> S) -> PreBound { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + PreBound { + inner: inner(&addr), + addr, + fd: Some(listener.into_raw_fd()), + } +} + +#[async_trait] +impl Service for PreBound { + async fn start_service( + &mut self, + fds: Option, + shutdown: ShutdownWatch, + listeners_per_fd: usize, + ) { + let fds = fds.expect("server provides an fd table"); + fds.lock().add(self.addr.clone(), self.fd.take().unwrap()); + self.inner + .start_service(Some(fds), shutdown, listeners_per_fd) + .await; + } + + fn name(&self) -> &str { + self.inner.name() + } + + fn threads(&self) -> Option { + Some(1) + } +} + +struct Reply; + +#[async_trait] +impl ServerApp for Reply { + async fn process_new( + self: &Arc, + mut session: Stream, + _shutdown: &ShutdownWatch, + ) -> Option { + session.write_all(b"ok").await.ok(); + session.shutdown().await.ok(); + None + } +} + +/// The supported readiness pattern: shutdown is observed separately from the listener. +struct Readiness { + draining: Arc, +} + +#[async_trait] +impl ServeHttp for Readiness { + async fn response(&self, _session: &mut ServerSession) -> Response> { + let (status, body) = if self.draining.load(Ordering::Relaxed) { + (StatusCode::SERVICE_UNAVAILABLE, "draining\n") + } else { + (StatusCode::OK, "ready\n") + }; + Response::builder() + .status(status) + .header(http::header::CONTENT_LENGTH, body.len()) + .body(body.into()) + .unwrap() + } +} + +struct DrainWatch { + draining: Arc, +} + +#[async_trait] +impl BackgroundService for DrainWatch { + async fn start(&self, mut shutdown: ShutdownWatch) { + shutdown.wait_for(|&shutting_down| shutting_down).await.ok(); + self.draining.store(true, Ordering::Relaxed); + } +} + +/// Returns the raw response to a fresh `GET /`, or `None` if nothing answers. +fn get(addr: &str) -> Option { + let mut stream = TcpStream::connect_timeout(&addr.parse().unwrap(), IO_TIMEOUT).ok()?; + stream.set_read_timeout(Some(IO_TIMEOUT)).unwrap(); + stream + .write_all(b"GET / HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n") + .ok()?; + let mut response = String::new(); + stream.read_to_string(&mut response).ok()?; + (!response.is_empty()).then_some(response) +} + +fn accepts(addr: &str) -> bool { + let Ok(mut stream) = TcpStream::connect_timeout(&addr.parse().unwrap(), IO_TIMEOUT) else { + return false; + }; + stream.set_read_timeout(Some(IO_TIMEOUT)).unwrap(); + let mut reply = Vec::new(); + stream.read_to_end(&mut reply).is_ok() && reply == b"ok" +} + +/// Refused, not merely unanswered: a stalled accept loop keeps its socket open. +fn refuses(addr: &str) -> bool { + matches!( + TcpStream::connect_timeout(&addr.parse().unwrap(), IO_TIMEOUT), + Err(e) if e.kind() == std::io::ErrorKind::ConnectionRefused + ) +} + +fn eventually(what: &str, mut done: impl FnMut() -> bool) { + let start = Instant::now(); + while !done() { + assert!(start.elapsed() < DEADLINE, "timed out waiting: {what}"); + thread::sleep(Duration::from_millis(10)); + } +} + +fn wait_for_phase( + phases: &mut broadcast::Receiver, + done: fn(&ExecutionPhase) -> bool, +) { + eventually("execution phase", || loop { + match phases.try_recv() { + Ok(phase) if done(&phase) => return true, + Ok(_) => {} + Err(broadcast::error::TryRecvError::Empty) => return false, + Err(e) => panic!("phase watch failed: {e}"), + } + }); +} + +/// Fails if runtime shutdown has begun, since it closes every listener regardless. +fn assert_before_runtime_shutdown(phases: &mut broadcast::Receiver) { + loop { + match phases.try_recv() { + Ok(ExecutionPhase::ShutdownRuntimes | ExecutionPhase::Terminated) => { + panic!("observed only after runtime shutdown began") + } + Ok(_) => {} + Err(broadcast::error::TryRecvError::Empty) => return, + Err(e) => panic!("phase watch failed: {e}"), + } + } +} + +/// Leaves room to observe the drain before runtime shutdown. +const GRACE_PERIOD: Duration = Duration::from_secs(5); + +struct Running { + /// Opted in. + kept: String, + /// Not opted in. + closed: String, + phases: broadcast::Receiver, + signal: mpsc::UnboundedSender, + server: JoinHandle<()>, +} + +/// Starts a server with the opted-in service `kept(&server, addr)` and a default listener. +fn run( + conf: ServerConf, + kept: impl FnOnce(&Server, &str) -> S, + extra: impl FnOnce(&mut Server), +) -> Running { + let mut server = Server::new_with_opt_and_conf( + None, + ServerConf { + grace_period_seconds: Some(GRACE_PERIOD.as_secs()), + graceful_shutdown_timeout_seconds: Some(1), + ..conf + }, + ); + let mut phases = server.watch_execution_phase(); + + let kept = pre_bound(|addr| kept(&server, addr)); + let closed = pre_bound(|addr| { + let mut service = ListeningService::new("closed".to_string(), Reply); + service.add_tcp(addr); + service + }); + let (kept_addr, closed_addr) = (kept.addr.clone(), closed.addr.clone()); + server.add_service(kept); + server.add_service(closed); + extra(&mut server); + + let (signal, signal_rx) = mpsc::unbounded_channel(); + let server = thread::spawn(move || { + server.bootstrap(); + server.run(RunArgs { + shutdown_signal: Box::new(InjectedSignal(Mutex::new(signal_rx))), + }); + }); + wait_for_phase(&mut phases, |p| matches!(p, ExecutionPhase::Running)); + eventually("default listener starts", || accepts(&closed_addr)); + Running { + kept: kept_addr, + closed: closed_addr, + phases, + signal, + server, + } +} + +fn join(server: JoinHandle<()>) { + eventually("server exits", || server.is_finished()); + server.join().unwrap(); +} + +#[test] +fn terminate_keeps_serving_readiness_until_runtime_shutdown() { + let draining = Arc::new(AtomicBool::new(false)); + let watch = DrainWatch { + draining: draining.clone(), + }; + let Running { + kept, + closed, + mut phases, + signal, + server, + } = run( + ServerConf::default(), + |server, addr| { + let mut service = ListeningService::new( + "readiness".to_string(), + HttpServer::new_app(Readiness { draining }), + ); + service.add_tcp(addr); + service.keep_listening_on_terminate(server); + service + }, + |server| { + server.add_service(background_service("drain watch", watch)); + }, + ); + eventually("readiness ready", || { + get(&kept).is_some_and(|r| r.starts_with("HTTP/1.1 200") && r.ends_with("ready\n")) + }); + + signal.send(ShutdownSignal::GracefulTerminate).unwrap(); + let draining_response = |r: &String| r.starts_with("HTTP/1.1 503") && r.ends_with("draining\n"); + eventually("readiness reports draining", || { + get(&kept).is_some_and(|r| draining_response(&r)) + }); + for _ in 0..3 { + let response = get(&kept).expect("readiness stopped answering during the drain"); + assert!(draining_response(&response), "{response}"); + } + eventually("default listener refuses", || refuses(&closed)); + assert_before_runtime_shutdown(&mut phases); + + join(server); + assert!(refuses(&kept), "readiness outlived runtime shutdown"); +} + +/// Socket transfer is Linux-only; elsewhere the server reports the upgrade without a taker. +#[cfg(target_os = "linux")] +fn take_over(upgrade_sock: String) -> JoinHandle> { + thread::spawn(move || { + let mut fds = pingora_core::server::Fds::new(); + fds.get_from_sock(upgrade_sock.as_str()) + .expect("socket transfer failed"); + fds.serialize().1 + }) +} + +#[test] +fn upgrade_hands_over_and_closes_kept_listener() { + // Short and on a local filesystem: socket paths are length-limited, and the socket is chmodded. + let upgrade_sock = std::env::temp_dir() + .join(format!("keep_listening_{}.sock", std::process::id())) + .to_string_lossy() + .into_owned(); + let _ = std::fs::remove_file(&upgrade_sock); + let mut running = run( + ServerConf { + upgrade_sock: upgrade_sock.clone(), + ..Default::default() + }, + |server, addr| { + let mut service = ListeningService::new("kept".to_string(), Reply); + service.add_tcp(addr); + service.keep_listening_on_terminate(server); + service + }, + |_| {}, + ); + eventually("kept listener starts", || accepts(&running.kept)); + + #[cfg(target_os = "linux")] + let new_process = take_over(upgrade_sock); + running + .signal + .send(ShutdownSignal::GracefulUpgrade) + .unwrap(); + #[cfg(target_os = "linux")] + { + let transferred = new_process.join().unwrap(); + assert_eq!(transferred.len(), 2, "both listeners are handed over"); + // Closed so refused connections below prove the old listeners stopped. + for fd in transferred { + unsafe { libc::close(fd) }; + } + } + + wait_for_phase(&mut running.phases, |p| { + matches!(p, ExecutionPhase::ShutdownGracePeriod) + }); + eventually("kept listener refuses", || refuses(&running.kept)); + eventually("default listener refuses", || refuses(&running.closed)); + assert_before_runtime_shutdown(&mut running.phases); + join(running.server); +} diff --git a/pingora-proxy/Cargo.toml b/pingora-proxy/Cargo.toml index e5cc0a236..ebca46c03 100644 --- a/pingora-proxy/Cargo.toml +++ b/pingora-proxy/Cargo.toml @@ -84,6 +84,7 @@ sentry = ["pingora-core/sentry"] patched_http1 = ["pingora-core/patched_http1"] upstream_modules = [] connection_filter = ["pingora-core/connection_filter"] +keep_listening_on_terminate = ["pingora-core/keep_listening_on_terminate"] trace = ["pingora-cache/trace"] early_body_buffer = ["dep:pingora-timeout"] @@ -95,6 +96,10 @@ required-features = ["connection_filter"] name = "body_routing" required-features = ["early_body_buffer"] +[[example]] +name = "keep_listening_on_terminate" +required-features = ["keep_listening_on_terminate"] + # or locally cargo doc --config "build.rustdocflags='--cfg doc_async_trait'" [package.metadata.docs.rs] rustdoc-args = ["--cfg", "doc_async_trait"] diff --git a/pingora-proxy/examples/keep_listening_on_terminate.rs b/pingora-proxy/examples/keep_listening_on_terminate.rs new file mode 100644 index 000000000..cbc907d65 --- /dev/null +++ b/pingora-proxy/examples/keep_listening_on_terminate.rs @@ -0,0 +1,122 @@ +// Copyright 2026 Cloudflare, Inc. +// +// 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. + +//! On SIGTERM the proxy stops accepting so the load balancer moves traffic away, while the +//! readiness and metrics endpoints stay reachable for the whole grace period. +//! +//! ```bash +//! cargo run --example keep_listening_on_terminate --features keep_listening_on_terminate +//! curl -i 127.0.0.1:6193 # readiness: 200 ready +//! pkill -TERM -f keep_listening_on_terminate +//! curl -i 127.0.0.1:6191 # proxy: connection refused +//! curl -i 127.0.0.1:6193 # readiness: 503 draining +//! curl 127.0.0.1:6192 # metrics: still served +//! ``` + +use async_trait::async_trait; +use http::{Response, StatusCode}; +use pingora_core::apps::http_app::{HttpServer, ServeHttp}; +use pingora_core::protocols::http::ServerSession; +use pingora_core::server::configuration::Opt; +use pingora_core::server::{Server, ShutdownWatch}; +use pingora_core::services::background::{background_service, BackgroundService}; +use pingora_core::services::listening::Service; +use pingora_core::upstreams::peer::HttpPeer; +use pingora_core::Result; +use pingora_proxy::{ProxyHttp, Session}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +struct Proxy; + +#[async_trait] +impl ProxyHttp for Proxy { + type CTX = (); + fn new_ctx(&self) -> Self::CTX {} + + async fn upstream_peer(&self, _session: &mut Session, _ctx: &mut ()) -> Result> { + Ok(Box::new(HttpPeer::new( + ("1.1.1.1", 80), + false, + String::new(), + ))) + } +} + +struct Readiness { + draining: Arc, +} + +#[async_trait] +impl ServeHttp for Readiness { + async fn response(&self, _session: &mut ServerSession) -> Response> { + let (status, body) = if self.draining.load(Ordering::Relaxed) { + (StatusCode::SERVICE_UNAVAILABLE, "draining\n") + } else { + (StatusCode::OK, "ready\n") + }; + Response::builder() + .status(status) + .header(http::header::CONTENT_LENGTH, body.len()) + .body(body.into()) + .unwrap() + } +} + +/// Marks the server as draining once shutdown starts. A kept listener's app is never cleaned up +/// on terminate, so draining has to be observed here. +struct DrainWatch { + draining: Arc, +} + +#[async_trait] +impl BackgroundService for DrainWatch { + async fn start(&self, mut shutdown: ShutdownWatch) { + shutdown.wait_for(|&shutting_down| shutting_down).await.ok(); + self.draining.store(true, Ordering::Relaxed); + } +} + +fn main() { + env_logger::init(); + + let mut server = Server::new(Some(Opt::parse_args())).unwrap(); + server.bootstrap(); + + let mut proxy = pingora_proxy::http_proxy_service(&server.configuration, Proxy); + proxy.add_tcp("127.0.0.1:6191"); + server.add_service(proxy); + + let mut metrics = pingora_prometheus::prometheus_http_service(); + metrics.add_tcp("127.0.0.1:6192"); + metrics.keep_listening_on_terminate(&server); + server.add_service(metrics); + + let draining = Arc::new(AtomicBool::new(false)); + server.add_service(background_service( + "drain watch", + DrainWatch { + draining: draining.clone(), + }, + )); + let mut readiness = Service::new( + "readiness".to_string(), + HttpServer::new_app(Readiness { draining }), + ); + readiness.add_tcp("127.0.0.1:6193"); + readiness.keep_listening_on_terminate(&server); + server.add_service(readiness); + + server.run_forever(); +} diff --git a/pingora/Cargo.toml b/pingora/Cargo.toml index 4caf24f55..5ee0910a2 100644 --- a/pingora/Cargo.toml +++ b/pingora/Cargo.toml @@ -163,6 +163,13 @@ connection_filter = [ ## `early_request_body_filter` trait hooks. early_body_buffer = ["pingora-proxy?/early_body_buffer"] +## Allow listening services to keep accepting through a graceful terminate, via +## `Service::keep_listening_on_terminate`. +keep_listening_on_terminate = [ + "pingora-core/keep_listening_on_terminate", + "pingora-proxy?/keep_listening_on_terminate", +] + # These features are intentionally not documented openssl_derived = ["any_tls"] @@ -179,6 +186,7 @@ document-features = [ "dial9-worker-s3", "dial9-cpu-profiling", "connection_filter", - "early_body_buffer" + "early_body_buffer", + "keep_listening_on_terminate" ] trace = ["pingora-cache?/trace", "pingora-proxy?/trace"] From 71037b8d4361f031d25d25507923179a6095f382 Mon Sep 17 00:00:00 2001 From: Cody Carlsen Date: Tue, 10 Nov 2026 11:57:49 -0800 Subject: [PATCH 2/2] Fix keep-listening example for fallible HttpPeer::new --- pingora-proxy/examples/keep_listening_on_terminate.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pingora-proxy/examples/keep_listening_on_terminate.rs b/pingora-proxy/examples/keep_listening_on_terminate.rs index cbc907d65..c5503daa3 100644 --- a/pingora-proxy/examples/keep_listening_on_terminate.rs +++ b/pingora-proxy/examples/keep_listening_on_terminate.rs @@ -50,7 +50,7 @@ impl ProxyHttp for Proxy { ("1.1.1.1", 80), false, String::new(), - ))) + )?)) } }