feat: handle SIGTERM in +watch and +subscribe for clean shutdown

.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.

Cargo.toml

@@ -84,4 +84,5 @@ inherits = "release"
 lto = "thin"
 
 [dev-dependencies]
+libc = "0.2"
 serial_test = "3.4.0"

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::<Option<()>>(),
+        );
+
         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::<Option<()>>(),
+        );
+
         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(())");
+    }
 }

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::<Option<()>>(),
+        );
+
         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::<Option<()>>(),
+        );
+
         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(())");
+    }
 }

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<Box<dyn Helper>> {
     match service {
         "gmail" => Some(Box::new(gmail::GmailHelper)),