diff --git a/crates/claudear-engine/src/watcher.rs b/crates/claudear-engine/src/watcher.rs index 2f4961a..a4af2cd 100644 --- a/crates/claudear-engine/src/watcher.rs +++ b/crates/claudear-engine/src/watcher.rs @@ -44,6 +44,11 @@ type QueuedIssue = (Issue, MatchResult, Option); /// on (marked handled) instead of re-triggering the fix agent every cycle. const MAX_REVIEW_COMMENT_ATTEMPTS: i64 = 5; +/// How long graceful shutdown waits for in-flight issue-processing tasks to +/// finish before aborting the stragglers. Bounds shutdown so a task wedged in an +/// external git/LLM operation with no internal timeout cannot stall a redeploy. +const GRACEFUL_DRAIN_BUDGET: std::time::Duration = std::time::Duration::from_secs(30); + /// Extracts the source name from a processing key of the form "source:issue_id". fn source_from_processing_key(key: &str) -> &str { key.split_once(':').map_or(key, |(source, _)| source) @@ -251,7 +256,10 @@ pub struct Watcher { /// agent-based (Claude Code) by default, local-LLM-based when `qa.use_llm` is /// set. intent_classifier: Option>, - /// Join handles for spawned issue-processing tasks (used by tests to drain). + /// Join handles for in-flight issue-processing tasks. + /// + /// Finished handles are reaped by [`Self::reap_finished_spawn_handles`] so the + /// list stays bounded by concurrency rather than by issues-processed-ever. spawn_handles: tokio::sync::Mutex>>, } @@ -347,10 +355,26 @@ impl Watcher { } } + /// Drop the join handles of issue-processing tasks that have already finished. + /// + /// Tokio keeps a task's allocation alive until both the scheduler and the last + /// `JoinHandle` are dropped. Retaining handles for completed tasks therefore + /// pins one whole `process_issue` state machine per issue for the lifetime of + /// the daemon, which grows RSS without bound until the container is OOM-killed. + /// Called once per poll cycle and on every dispatch. + async fn reap_finished_spawn_handles(&self) { + let mut handles = self.spawn_handles.lock().await; + handles.retain(|handle| !handle.is_finished()); + // Also release the backing capacity of an unusually large burst. + if handles.capacity() > handles.len().saturating_mul(4) + 16 { + handles.shrink_to_fit(); + } + } + /// Wait for all spawned issue-processing tasks to complete. /// - /// Primarily useful in tests that need to assert on processing outcomes - /// after a non-blocking `poll_source` call. + /// Used on graceful shutdown, and by tests that need to assert on processing + /// outcomes after a non-blocking `poll_source` call. pub async fn drain_spawned_tasks(&self) { let handles: Vec<_> = { let mut guard = self.spawn_handles.lock().await; @@ -1149,32 +1173,61 @@ impl Watcher { /// all in-progress work completes before the application exits. pub async fn stop_and_drain(&self) { self.stop(); + self.drain_or_abort(GRACEFUL_DRAIN_BUDGET).await; + } - // Wait for any active processing to complete (up to 30 seconds). - // Uses slot_available to wake immediately when a task finishes rather - // than polling on a fixed interval. - let max_wait = std::time::Duration::from_secs(30); - let start = std::time::Instant::now(); + /// Drain in-flight issue-processing tasks, giving them `budget` to finish + /// gracefully and then aborting any stragglers. + /// + /// `stop()` has already cleared is_running, so no new tasks are recorded. + /// The two failure modes this balances, both flagged in review: + /// - Returning early while a task still runs let runtime teardown abort it + /// mid-operation and misreport a graceful stop. So we wait for real + /// completion within the budget, and any straggler is aborted *explicitly* + /// and logged here rather than left to implicit teardown. + /// - Waiting with no deadline stalled a non-interactive shutdown (e.g. a + /// redeploy's SIGTERM) when a task was wedged in an external git/LLM + /// operation with no internal timeout. So the wait is bounded. + /// + /// Cancel-safe: handles are joined via `&mut`, so the timeout drops only the + /// borrow — the owned `handles` survive for the abort pass, never detached. + async fn drain_or_abort(&self, budget: std::time::Duration) { + // Safe to take: is_running is already false and dispatch re-checks it + // under this same lock, so nothing new is pushed after this take. + let mut handles: Vec<_> = { + let mut guard = self.spawn_handles.lock().await; + std::mem::take(&mut *guard) + }; - while self.active_processing.load(Ordering::SeqCst) > 0 { - if start.elapsed() > max_wait { - tracing::warn!( - remaining = self.active_processing.load(Ordering::SeqCst), - "Graceful shutdown timeout reached, some tasks may not have completed" - ); - break; + let joined = tokio::time::timeout(budget, async { + for handle in &mut handles { + let _ = handle.await; } - tracing::info!( - active_count = self.active_processing.load(Ordering::SeqCst), - "Waiting for active tasks to complete..." - ); - // Wait for a task to finish (notifies via slot_available) or fall back - // to a periodic check in case the notification was missed. - let remaining = max_wait.saturating_sub(start.elapsed()); - let _ = tokio::time::timeout(remaining, self.slot_available.notified()).await; + }) + .await; + + if joined.is_ok() { + tracing::info!("Claude Watcher stopped gracefully"); + return; } - tracing::info!("Claude Watcher stopped gracefully"); + // Budget exhausted: abort the still-running stragglers deterministically + // instead of stalling shutdown or leaving them for runtime teardown. + let mut aborted = 0usize; + for handle in &mut handles { + if !handle.is_finished() { + handle.abort(); + aborted += 1; + } + } + for handle in handles { + let _ = handle.await; + } + tracing::warn!( + aborted, + budget_secs = budget.as_secs(), + "Graceful shutdown budget exhausted; aborted in-flight tasks to avoid stalling shutdown" + ); } /// Check if the watcher is currently running. @@ -3046,6 +3099,10 @@ Create a PR with your changes.{custom_instructions}"#, /// Poll a single source. async fn poll_source(self: &Arc, source: &Arc) -> Result<()> { + // Release tasks that finished since the last cycle before doing anything else, + // so an idle or rate-limit-paused watcher still frees their allocations. + self.reap_finished_spawn_handles().await; + if self.is_rate_limit_paused().await { return Ok(()); } @@ -3546,12 +3603,26 @@ Create a PR with your changes.{custom_instructions}"#, // the housekeeping loop (review checks, auto-close, retries) is not starved. let watcher = Arc::clone(self); let source_clone = Arc::clone(source); - let handle = tokio::spawn(async move { - watcher - .process_issue(source_clone, issue, match_result, None, None, intent) - .await; - }); - self.spawn_handles.lock().await.push(handle); + { + let mut handles = self.spawn_handles.lock().await; + // Re-check is_running under the same lock that drain_spawned_tasks + // takes, then spawn+record atomically. This closes the shutdown + // race: once stop() has set is_running=false (before the drain), + // a dispatch either records its handle before the drain's take (so + // shutdown joins it) or sees the stop here and never spawns. The + // top-of-loop check is not enough on its own — it runs before the + // spawn, so a concurrent drain could take the vector between it and + // the push. + if !self.is_running.load(Ordering::SeqCst) { + break; + } + handles.retain(|h| !h.is_finished()); + handles.push(tokio::spawn(async move { + watcher + .process_issue(source_clone, issue, match_result, None, None, intent) + .await; + })); + } // Add delay between starting new issues (skip trailing delay after the last item). if i + 1 < total && self.config.processing_delay_ms > 0 { @@ -3957,7 +4028,9 @@ Create a PR with your changes.{custom_instructions}"#, }; let context_provider = crate::processing::SourceContext(source.as_ref()); - let outcome = processor.run(input, &context_provider).await; + // Box the pipeline future: `run` inlines the whole processing state machine, + // so awaiting it directly would make every spawned task allocation carry it. + let outcome = Box::pin(processor.run(input, &context_provider)).await; // Watcher-specific: check for rate limit errors and pause if needed if let ProcessingOutcome::Failed { ref error } = outcome { @@ -5110,6 +5183,160 @@ mod tests { })) } + /// Completed issue-processing tasks must not stay pinned in `spawn_handles`. + /// + /// Tokio frees a task's allocation only once the scheduler *and* the last + /// `JoinHandle` are gone, so a push-only handle list leaks one whole + /// `process_issue` state machine per issue, growing the daemon's RSS until + /// the container is OOM-killed. + #[tokio::test] + async fn test_watcher_reaps_finished_spawn_handles() { + let notifier = Arc::new(MockNotifier::new(true)); + let tracker = Arc::new(SqliteTracker::in_memory().unwrap()); + // Source with no issues: this test is about reaping, not dispatching. + let source = Arc::new(MockSource::new("reap")) as Arc; + + let watcher = create_test_watcher(notifier, tracker, vec![source.clone()], false); + watcher.is_running.store(true, Ordering::SeqCst); + + // Stand in for the handles left behind by 64 already-processed issues. + { + let mut handles = watcher.spawn_handles.lock().await; + for _ in 0..64 { + handles.push(tokio::spawn(async {})); + } + } + + // Let every simulated task run to completion. + for _ in 0..1000 { + if watcher + .spawn_handles + .lock() + .await + .iter() + .all(|handle| handle.is_finished()) + { + break; + } + tokio::task::yield_now().await; + } + + // A normal poll cycle must release them; nothing calls drain in production. + watcher.poll_source(&source).await.unwrap(); + + let retained = watcher.spawn_handles.lock().await.len(); + assert_eq!( + retained, 0, + "watcher retained {retained} handles for completed issue-processing tasks; \ + each one pins a full process_issue allocation for the daemon's lifetime" + ); + } + + #[tokio::test] + async fn test_stop_and_drain_joins_recorded_task() { + use std::sync::atomic::AtomicBool; + let notifier = Arc::new(MockNotifier::new(true)); + let tracker = Arc::new(SqliteTracker::in_memory().unwrap()); + let source = Arc::new(MockSource::new("drain")) as Arc; + let watcher = create_test_watcher(notifier, tracker, vec![source], false); + watcher.is_running.store(true, Ordering::SeqCst); + + // A recorded task whose teardown completes shortly after it is spawned. + let done = Arc::new(AtomicBool::new(false)); + let done_clone = Arc::clone(&done); + { + let mut handles = watcher.spawn_handles.lock().await; + handles.push(tokio::spawn(async move { + tokio::task::yield_now().await; + done_clone.store(true, Ordering::SeqCst); + })); + } + + // Graceful shutdown must wait for the recorded task's teardown rather than + // reporting a clean stop while it is still mutating shared state. + watcher.stop_and_drain().await; + + assert!( + done.load(Ordering::SeqCst), + "stop_and_drain returned before the recorded task finished; a \ + concurrently processing issue could still be mutating tracker/notifier/agent state" + ); + assert!( + watcher.spawn_handles.lock().await.is_empty(), + "drain must release recorded handles" + ); + assert!(!watcher.is_running.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_stop_and_drain_joins_task_to_completion() { + use std::sync::atomic::AtomicBool; + use std::time::Duration; + let notifier = Arc::new(MockNotifier::new(true)); + let tracker = Arc::new(SqliteTracker::in_memory().unwrap()); + let source = Arc::new(MockSource::new("drain")) as Arc; + let watcher = create_test_watcher(notifier, tracker, vec![source], false); + watcher.is_running.store(true, Ordering::SeqCst); + + // An in-flight task that only finishes after a delay. Shutdown must wait + // for it to complete rather than returning and letting it be aborted. + let done = Arc::new(AtomicBool::new(false)); + let done_clone = Arc::clone(&done); + { + let mut handles = watcher.spawn_handles.lock().await; + handles.push(tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(150)).await; + done_clone.store(true, Ordering::SeqCst); + })); + } + + watcher.stop_and_drain().await; + + assert!( + done.load(Ordering::SeqCst), + "stop_and_drain returned before the in-flight task completed; runtime \ + teardown would then abort it mid-operation" + ); + assert!(watcher.spawn_handles.lock().await.is_empty()); + assert!(!watcher.is_running.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_drain_or_abort_bounds_wedged_task() { + use std::sync::atomic::AtomicBool; + use std::time::{Duration, Instant}; + let notifier = Arc::new(MockNotifier::new(true)); + let tracker = Arc::new(SqliteTracker::in_memory().unwrap()); + let source = Arc::new(MockSource::new("drain")) as Arc; + let watcher = create_test_watcher(notifier, tracker, vec![source], false); + + // A task wedged far longer than the budget (stands in for an external + // git/LLM op with no internal timeout). + let done = Arc::new(AtomicBool::new(false)); + let done_clone = Arc::clone(&done); + { + let mut handles = watcher.spawn_handles.lock().await; + handles.push(tokio::spawn(async move { + tokio::time::sleep(Duration::from_secs(30)).await; + done_clone.store(true, Ordering::SeqCst); + })); + } + + // Shutdown must not stall: it returns within the (tiny) budget, aborting + // the straggler rather than waiting for it or leaving it detached. + let start = Instant::now(); + watcher.drain_or_abort(Duration::from_millis(50)).await; + assert!( + start.elapsed() < Duration::from_secs(2), + "drain_or_abort stalled past its budget" + ); + assert!( + !done.load(Ordering::SeqCst), + "the wedged task was aborted, not awaited to completion" + ); + assert!(watcher.spawn_handles.lock().await.is_empty()); + } + #[test] fn test_watcher_new() { let notifier = Arc::new(MockNotifier::new(true)); @@ -7484,21 +7711,30 @@ mod tests { let watcher = Arc::new(create_test_watcher(notifier, tracker, sources, false)); watcher.is_running.store(true, Ordering::SeqCst); - watcher.active_processing.fetch_add(1, Ordering::SeqCst); - // Simulate task finishing after a short delay + // A handle-tracked in-flight task, mirroring process_issue: it holds an + // active_processing count and clears it only when it finishes. Shutdown + // must wait for the handle, so the count is 0 by the time it returns. + watcher.active_processing.fetch_add(1, Ordering::SeqCst); let release = Arc::clone(&watcher); - tokio::spawn(async move { - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - release.active_processing.fetch_sub(1, Ordering::SeqCst); - release.slot_available.notify_waiters(); - }); + { + let mut handles = watcher.spawn_handles.lock().await; + handles.push(tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + release.active_processing.fetch_sub(1, Ordering::SeqCst); + release.slot_available.notify_waiters(); + })); + } let result = tokio::time::timeout(std::time::Duration::from_secs(5), watcher.stop_and_drain()).await; assert!(result.is_ok(), "stop_and_drain timed out"); assert!(!watcher.is_running()); - assert_eq!(watcher.active_count(), 0); + assert_eq!( + watcher.active_count(), + 0, + "shutdown waited for the in-flight task to finish" + ); } #[test] @@ -9896,23 +10132,17 @@ mod tests { let watcher = Arc::new(create_test_watcher(notifier, tracker, vec![], false)); watcher.is_running.store(true, Ordering::SeqCst); - // Simulate a task that never completes (active count stays > 0) - watcher.active_processing.store(1, Ordering::SeqCst); - - // stop_and_drain has a 5-minute internal timeout, but we use an outer timeout - // We just verify it eventually returns (the internal max_wait breaks the loop) + // A stray active_processing count with no recorded handle must not wedge + // shutdown: stop_and_drain waits on the spawn_handles, which are empty + // here, so it returns promptly and clears is_running. let result = tokio::time::timeout(std::time::Duration::from_secs(10), watcher.stop_and_drain()) .await; - // In test the internal max_wait is 300s which we can't wait for, - // so this test verifies the method was called correctly and stop was set - // The timeout will trigger because 300s > 10s, but that's fine - if result.is_err() { - // Timed out externally - that's expected since internal timeout is 300s - assert!(!watcher.is_running()); - } else { - assert!(!watcher.is_running()); - } + assert!( + result.is_ok(), + "stop_and_drain should return promptly with no handles" + ); + assert!(!watcher.is_running()); } #[tokio::test]