From a9e523deb02fc408257576967a916d5032f1d923 Mon Sep 17 00:00:00 2001 From: OpenAI Codex Date: Sun, 30 Aug 2026 21:16:58 +0900 Subject: [PATCH] fix(runtime): arm shutdown before readiness --- src/lib.rs | 38 +++++++++++++++++++ src/main.rs | 106 +++++++++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 134 insertions(+), 10 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 8f54751..b112e96 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4509,6 +4509,44 @@ mod tests { assert_eq!(support.threat_feed_freshness.len(), 1); assert!(!support.threat_feed_freshness[0].stale); assert!(support.event_count >= 1); + assert_eq!(support.route_count, support.kpis.route_count); + assert_eq!( + support.threat_indicator_count, + support.kpis.threat_indicator_count + ); + assert_eq!(support.dnsbl_entry_count, support.kpis.dnsbl_entry_count); + assert_eq!(support.threat_feed_count, support.kpis.threat_feed_count); + assert_eq!(support.event_count, support.kpis.event_count); + assert_eq!(support.audit_log_count, support.kpis.audit_log_count); + assert_eq!( + support.route_count, + support.evidence_manifest.runtime_counts.route_count + ); + assert_eq!( + support.threat_indicator_count, + support + .evidence_manifest + .runtime_counts + .threat_indicator_count + ); + assert_eq!( + support.dnsbl_entry_count, + support.evidence_manifest.runtime_counts.dnsbl_entry_count + ); + assert_eq!( + support.threat_feed_count, + support.evidence_manifest.runtime_counts.threat_feed_count + ); + assert_eq!( + support.event_count, + support.evidence_manifest.runtime_counts.event_count + ); + assert_eq!( + support.audit_log_count, + support.evidence_manifest.runtime_counts.audit_log_count + ); + let support_json = serde_json::to_string(&support).unwrap(); + assert!(!support_json.contains("\"secret\"")); let persisted: AppData = serde_json::from_str(&fs::read_to_string(&path).await.unwrap()).unwrap(); diff --git a/src/main.rs b/src/main.rs index 2c8fc82..fa65bc4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,24 +1,110 @@ +use std::{future::Future, pin::Pin}; +#[cfg(any(test, not(unix)))] +use std::{future::poll_fn, task::Poll}; + // The gateway entrypoint is intentionally a thin shim: all configuration // parsing, binding, and serving live in `waf_ids_ai_soc::run_from_env` so they // are unit-testable, while this file is covered end-to-end by `tests/binary.rs`. #[cfg(not(test))] #[tokio::main] async fn main() -> Result<(), Box> { - waf_ids_ai_soc::run_from_env(Box::pin(shutdown_signal())).await + waf_ids_ai_soc::run_from_env(shutdown_signal().await?).await } #[cfg(all(not(test), unix))] -async fn shutdown_signal() { - // Shut down gracefully on SIGTERM (what container runtimes and the e2e test - // harness send) so in-flight requests drain and the process exits cleanly. - let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) - .expect("install SIGTERM handler"); - term.recv().await; +async fn shutdown_signal() +-> Result + Send>>, Box> { + // Install SIGTERM handling before readiness can be reported, so a fast + // supervisor or test harness cannot kill the process before graceful + // shutdown is armed. + let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; + Ok(Box::pin(async move { + term.recv().await; + })) +} + +/// Poll a shutdown future once up front so listeners that install on first +/// poll, such as `tokio::signal::ctrl_c()`, are armed before startup runs. +#[cfg(any(test, not(unix)))] +async fn arm_shutdown_future(future: F) -> Result + Send>>, E> +where + F: Future> + Send + 'static, +{ + let mut future = Box::pin(future); + let ready = poll_fn(|cx| match future.as_mut().poll(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(true)), + Poll::Ready(Err(error)) => Poll::Ready(Err(error)), + Poll::Pending => Poll::Ready(Ok(false)), + }) + .await?; + if ready { + return Ok(Box::pin(async {})); + } + Ok(Box::pin(async move { + let _ = future.await; + })) } #[cfg(all(not(test), not(unix)))] -async fn shutdown_signal() { - tokio::signal::ctrl_c() +async fn shutdown_signal() +-> Result + Send>>, Box> { + Ok(arm_shutdown_future(tokio::signal::ctrl_c()).await?) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{ + io, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; + + struct PendingThenReady { + polls: Arc, + } + + impl Future for PendingThenReady { + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, _cx: &mut std::task::Context<'_>) -> Poll { + let polls = self.polls.fetch_add(1, Ordering::SeqCst); + if polls == 0 { + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + } + } + + #[tokio::test] + async fn arm_shutdown_future_registers_listener_before_await() { + let polls = Arc::new(AtomicUsize::new(0)); + let shutdown = arm_shutdown_future(PendingThenReady { + polls: polls.clone(), + }) + .await + .unwrap(); + + assert_eq!(polls.load(Ordering::SeqCst), 1); + shutdown.await; + assert_eq!(polls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn arm_shutdown_future_propagates_registration_error() { + let err = match arm_shutdown_future(async { + Err::<(), io::Error>(io::Error::other("listener failed")) + }) .await - .expect("install Ctrl-C handler"); + { + Ok(_) => panic!("listener registration should fail"), + Err(err) => err, + }; + + assert_eq!(err.kind(), io::ErrorKind::Other); + assert_eq!(err.to_string(), "listener failed"); + } }