diff --git a/.changeset/sigterm-graceful-shutdown.md b/.changeset/sigterm-graceful-shutdown.md new file mode 100644 index 00000000..e98bce63 --- /dev/null +++ b/.changeset/sigterm-graceful-shutdown.md @@ -0,0 +1,8 @@ +--- +"googleworkspace-cli": patch +--- + +Handle SIGTERM in `gws gmail +watch` and `gws events +subscribe` for clean shutdown. + +Long-running pull loops now exit gracefully on SIGTERM (in addition to Ctrl+C), +enabling clean shutdown in containers and process supervisors. diff --git a/Cargo.lock b/Cargo.lock index 10e12097..7122d964 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -914,6 +914,7 @@ dependencies = [ "hostname", "iana-time-zone", "keyring", + "libc", "mail-builder", "mime_guess2", "percent-encoding", diff --git a/Cargo.toml b/Cargo.toml index 1f73b8f7..fbd98448 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,4 +84,5 @@ inherits = "release" lto = "thin" [dev-dependencies] +libc = "0.2" serial_test = "3.4.0" diff --git a/src/helpers/events/subscribe.rs b/src/helpers/events/subscribe.rs index 00474ac7..3b22eeaf 100644 --- a/src/helpers/events/subscribe.rs +++ b/src/helpers/events/subscribe.rs @@ -320,6 +320,14 @@ async fn pull_loop( pubsub_api_base: &str, ) -> Result<(), GwsError> { let mut file_counter: u64 = 0; + + #[cfg(unix)] + let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .context("failed to register SIGTERM handler")?; + #[cfg(unix)] + let mut sigint = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt()) + .context("failed to register SIGINT handler")?; + loop { let token = token_provider .access_token() @@ -337,6 +345,16 @@ async fn pull_loop( .timeout(std::time::Duration::from_secs(config.poll_interval.max(10))) .send(); + // Hoist signal futures outside select! — #[cfg] is not supported inside + // tokio::select! branches; cfg'd let bindings are the workaround. + #[cfg(unix)] + hoist_signals!(sigint, sigterm, sigint_fut, sigterm_fut); + #[cfg(not(unix))] + let (sigint_fut, sigterm_fut) = ( + tokio::signal::ctrl_c(), + std::future::pending::>(), + ); + let resp = tokio::select! { result = pull_future => { match result { @@ -345,10 +363,14 @@ async fn pull_loop( Err(e) => return Err(anyhow::anyhow!("Pub/Sub pull failed: {e}").into()), } } - _ = tokio::signal::ctrl_c() => { + _ = sigint_fut => { eprintln!("\nReceived interrupt, stopping..."); return Ok(()); } + _ = sigterm_fut => { + eprintln!("\nReceived SIGTERM, stopping..."); + return Ok(()); + } }; if !resp.status().is_success() { @@ -411,13 +433,25 @@ async fn pull_loop( break; } - // Check for SIGINT between polls + // Check for SIGINT/SIGTERM between polls + #[cfg(unix)] + hoist_signals!(sigint, sigterm, sigint_fut, sigterm_fut); + #[cfg(not(unix))] + let (sigint_fut, sigterm_fut) = ( + tokio::signal::ctrl_c(), + std::future::pending::>(), + ); + tokio::select! { _ = tokio::time::sleep(std::time::Duration::from_secs(config.poll_interval)) => {}, - _ = tokio::signal::ctrl_c() => { + _ = sigint_fut => { eprintln!("\nReceived interrupt, stopping..."); break; } + _ = sigterm_fut => { + eprintln!("\nReceived SIGTERM, stopping..."); + break; + } } } @@ -879,4 +913,46 @@ mod tests { ); assert_eq!(requests[1].1, "authorization: Bearer pubsub-token"); } + + #[cfg(unix)] + #[tokio::test] + #[serial_test::serial] + async fn test_pull_loop_exits_on_sigterm() { + use std::time::Duration; + + let client = reqwest::Client::new(); + let token_provider = FakeTokenProvider::new(["tok"; 16]); + let (base, server) = crate::helpers::test_helpers::spawn_empty_pull_server().await; + let config = SubscribeConfigBuilder::default() + .subscription(Some(SubscriptionName( + "projects/test/subscriptions/demo".to_string(), + ))) + .max_messages(1_u32) + .poll_interval(60_u64) // Long sleep so SIGTERM fires during the sleep select! + .once(false) + .build() + .unwrap(); + + // Send SIGTERM after the first pull completes and the loop enters the sleep select! + tokio::spawn(async { + tokio::time::sleep(Duration::from_millis(200)).await; + unsafe { libc::raise(libc::SIGTERM) }; + }); + + let result = tokio::time::timeout( + Duration::from_secs(2), + pull_loop( + &client, + &token_provider, + "projects/test/subscriptions/demo", + config, + &base, + ), + ) + .await; + + server.abort(); + let inner = result.expect("loop should exit on SIGTERM within 2 seconds"); + assert!(inner.is_ok(), "loop should return Ok(())"); + } } diff --git a/src/helpers/gmail/watch.rs b/src/helpers/gmail/watch.rs index 43707ee6..525c1a2e 100644 --- a/src/helpers/gmail/watch.rs +++ b/src/helpers/gmail/watch.rs @@ -264,6 +264,13 @@ async fn watch_pull_loop( last_history_id: &mut u64, config: WatchConfig, ) -> Result<(), GwsError> { + #[cfg(unix)] + let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .context("failed to register SIGTERM handler")?; + #[cfg(unix)] + let mut sigint = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt()) + .context("failed to register SIGINT handler")?; + loop { let pubsub_token = runtime .pubsub_token_provider @@ -280,6 +287,16 @@ async fn watch_pull_loop( .timeout(std::time::Duration::from_secs(config.poll_interval.max(10))) .send(); + // Hoist signal futures outside select! — #[cfg] is not supported inside + // tokio::select! branches; cfg'd let bindings are the workaround. + #[cfg(unix)] + hoist_signals!(sigint, sigterm, sigint_fut, sigterm_fut); + #[cfg(not(unix))] + let (sigint_fut, sigterm_fut) = ( + tokio::signal::ctrl_c(), + std::future::pending::>(), + ); + let resp = tokio::select! { result = pull_future => { match result { @@ -288,10 +305,14 @@ async fn watch_pull_loop( Err(e) => return Err(GwsError::Other(anyhow::anyhow!("Pub/Sub pull failed: {e}"))), } } - _ = tokio::signal::ctrl_c() => { + _ = sigint_fut => { eprintln!("\nReceived interrupt, stopping..."); return Ok(()); } + _ = sigterm_fut => { + eprintln!("\nReceived SIGTERM, stopping..."); + return Ok(()); + } }; if !resp.status().is_success() { @@ -346,12 +367,24 @@ async fn watch_pull_loop( break; } + #[cfg(unix)] + hoist_signals!(sigint, sigterm, sigint_fut, sigterm_fut); + #[cfg(not(unix))] + let (sigint_fut, sigterm_fut) = ( + tokio::signal::ctrl_c(), + std::future::pending::>(), + ); + tokio::select! { _ = tokio::time::sleep(std::time::Duration::from_secs(config.poll_interval)) => {}, - _ = tokio::signal::ctrl_c() => { + _ = sigint_fut => { eprintln!("\nReceived interrupt, stopping..."); break; } + _ = sigterm_fut => { + eprintln!("\nReceived SIGTERM, stopping..."); + break; + } } } @@ -990,4 +1023,63 @@ mod tests { assert_eq!(requests[3].1, "authorization: Bearer pubsub-token"); assert_eq!(last_history_id, 2); } + + #[cfg(unix)] + #[tokio::test] + #[serial_test::serial] + async fn test_watch_pull_loop_exits_on_sigterm() { + use std::time::Duration; + + let client = reqwest::Client::new(); + // Supply enough tokens for several pull iterations + let pubsub_provider = FakeTokenProvider::new(["tok"; 16]); + let gmail_provider = FakeTokenProvider::new(["tok"; 16]); + let (base, server) = crate::helpers::test_helpers::spawn_empty_pull_server().await; + let sanitize_config = crate::helpers::modelarmor::SanitizeConfig { + template: None, + mode: crate::helpers::modelarmor::SanitizeMode::Warn, + }; + let runtime = WatchRuntime { + client: &client, + pubsub_token_provider: &pubsub_provider, + gmail_token_provider: &gmail_provider, + sanitize_config: &sanitize_config, + pubsub_api_base: &base, + gmail_api_base: &base, + }; + let mut last_history_id = 0u64; + let config = WatchConfig { + project: None, + subscription: None, + topic: None, + label_ids: None, + max_messages: 1, + poll_interval: 60, // Long sleep so SIGTERM fires during the sleep select! + format: "full".to_string(), + once: false, + cleanup: false, + output_dir: None, + }; + + // Send SIGTERM after the first pull completes and the loop enters the sleep select! + tokio::spawn(async { + tokio::time::sleep(Duration::from_millis(200)).await; + unsafe { libc::raise(libc::SIGTERM) }; + }); + + let result = tokio::time::timeout( + Duration::from_secs(2), + watch_pull_loop( + &runtime, + "projects/test/subscriptions/demo", + &mut last_history_id, + config, + ), + ) + .await; + + server.abort(); + let inner = result.expect("loop should exit on SIGTERM within 2 seconds"); + assert!(inner.is_ok(), "loop should return Ok(())"); + } } diff --git a/src/helpers/mod.rs b/src/helpers/mod.rs index 72d31272..352957ee 100644 --- a/src/helpers/mod.rs +++ b/src/helpers/mod.rs @@ -16,6 +16,23 @@ use crate::error::GwsError; use clap::{ArgMatches, Command}; use std::future::Future; use std::pin::Pin; + +/// Hoists Unix signal futures into caller-named bindings for use inside +/// `tokio::select!`, which does not support `#[cfg]` inside its branches. +/// Output variable names are passed as arguments so they are visible after the +/// macro invocation (Rust macro hygiene: only call-site identifiers leak out). +/// +/// Only defined on Unix — call sites must pair this with a `#[cfg(not(unix))]` +/// fallback that creates the same two bindings via `ctrl_c()` / `pending()`. +/// Defined here so all helper submodules can use it without an explicit import. +#[cfg(unix)] +macro_rules! hoist_signals { + ($sigint:ident, $sigterm:ident, $sigint_fut:ident, $sigterm_fut:ident) => { + let $sigint_fut = $sigint.recv(); + let $sigterm_fut = $sigterm.recv(); + }; +} + pub mod calendar; pub mod chat; pub mod docs; @@ -54,6 +71,39 @@ pub trait Helper: Send + Sync { } } +#[cfg(test)] +pub(crate) mod test_helpers { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + /// Spawns a mock Pub/Sub server that returns empty pull responses + /// indefinitely. Shared by the SIGTERM tests in `gmail::watch` and + /// `events::subscribe` so the loop blocks in the sleep `select!`. + pub(crate) async fn spawn_empty_pull_server() -> (String, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let handle = tokio::spawn(async move { + loop { + let Ok((mut stream, _)) = listener.accept().await else { + break; + }; + tokio::spawn(async move { + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await; + let body = r#"{"receivedMessages":[]}"#; + let resp = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: application/json\r\n\r\n{}", + body.len(), + body + ); + let _ = stream.write_all(resp.as_bytes()).await; + }); + } + }); + (format!("http://{addr}"), handle) + } +} + pub fn get_helper(service: &str) -> Option> { match service { "gmail" => Some(Box::new(gmail::GmailHelper)),