zeph_common/task_supervisor.rs
1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Supervised lifecycle task manager for long-running named services.
5//!
6//! [`TaskSupervisor`] manages named, long-lived background tasks (config watcher,
7//! scheduler loop, gateway, MCP connections, etc.) with restart policies, health
8//! snapshots, and graceful shutdown. Unlike `BackgroundSupervisor`
9//! (which is `&mut self`-only, lossy, and turn-scoped), `TaskSupervisor` is
10//! `Clone + Send + Sync` and designed for the full agent session lifetime.
11//!
12//! # Design rationale
13//!
14//! - **Shared handle**: `Arc<Inner>` interior allows passing the supervisor to bootstrap
15//! code, TUI status display, and shutdown orchestration without lifetime coupling.
16//! - **Event-driven reap**: An internal mpsc channel delivers completion events to a
17//! reap driver task; no polling interval required.
18//! - **No `JoinSet`**: Individual `JoinHandle`s per task enable per-name abort, status
19//! tracking, and restart policies — `JoinSet` is better for homogeneous work.
20//! - **Mutex held briefly**: `parking_lot::Mutex` guards only bookkeeping operations
21//! (insert/remove from `HashMap`). The lock is **never held across `.await`**.
22//!
23//! # Examples
24//!
25//! ```rust,no_run
26//! use std::time::Duration;
27//! use tokio_util::sync::CancellationToken;
28//! use zeph_common::task_supervisor::{RestartPolicy, TaskDescriptor, TaskSupervisor};
29//!
30//! # #[tokio::main]
31//! # async fn main() {
32//! let cancel = CancellationToken::new();
33//! let supervisor = TaskSupervisor::new(cancel.clone());
34//!
35//! supervisor.spawn(TaskDescriptor {
36//! name: "my-service",
37//! restart: RestartPolicy::Restart { max: 3, base_delay: Duration::from_secs(1) },
38//! factory: || async { /* service loop */ },
39//! });
40//!
41//! // Graceful shutdown — waits up to 5 s for all tasks to stop.
42//! supervisor.shutdown_all(Duration::from_secs(5)).await;
43//! # }
44//! ```
45
46use std::collections::HashMap;
47use std::future::Future;
48use std::pin::Pin;
49use std::sync::Arc;
50use std::time::{Duration, Instant};
51
52use tokio::sync::{mpsc, oneshot};
53use tokio::task::AbortHandle;
54use tokio_util::sync::CancellationToken;
55use tracing::Instrument as _;
56
57use crate::BlockingSpawner;
58
59// ── Public types ─────────────────────────────────────────────────────────────
60
61/// Policy governing what happens when a supervised task completes or panics.
62///
63/// Used in [`TaskDescriptor`] to configure restart behaviour for a task.
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65#[non_exhaustive]
66pub enum RestartPolicy {
67 /// Task runs once; normal completion removes it from the registry.
68 RunOnce,
69 /// Task is restarted **only on panic**, up to `max` times.
70 ///
71 /// Normal completion (the future returns `()`) does **not** trigger a restart.
72 /// The task is removed from the registry on normal exit.
73 ///
74 /// A `max` of `0` means the task is monitored but **never** restarted —
75 /// a panic leaves the entry as `Failed` in the registry for observability.
76 /// Use `RunOnce` when you want the entry removed on completion.
77 ///
78 /// Restart delays follow **exponential backoff**: the delay before attempt `n`
79 /// is `base_delay * 2^(n-1)`, capped at [`MAX_RESTART_DELAY`].
80 ///
81 /// # Examples
82 ///
83 /// ```
84 /// use std::time::Duration;
85 /// use zeph_common::task_supervisor::RestartPolicy;
86 ///
87 /// // Restart up to 3 times with exponential backoff starting at 1 s.
88 /// let policy = RestartPolicy::Restart { max: 3, base_delay: Duration::from_secs(1) };
89 /// ```
90 Restart { max: u32, base_delay: Duration },
91}
92
93/// Maximum delay between restart attempts (caps exponential backoff).
94pub const MAX_RESTART_DELAY: Duration = Duration::from_mins(1);
95
96/// Configuration passed to [`TaskSupervisor::spawn`] to describe a supervised task.
97///
98/// `F` must be `Fn` (not `FnOnce`) to support restarts: the factory is called once on
99/// initial spawn and once per restart attempt.
100pub struct TaskDescriptor<F> {
101 /// Unique name for this task (e.g., `"config-watcher"`, `"scheduler-loop"`).
102 ///
103 /// Names must be `'static` — they are typically compile-time string literals.
104 /// Spawning a task with a name that already exists aborts the prior instance.
105 pub name: &'static str,
106 /// Restart policy applied when the task exits unexpectedly.
107 pub restart: RestartPolicy,
108 /// Factory called to produce a new future. Must be `Fn` for restart support.
109 pub factory: F,
110}
111
112/// Opaque handle to a single supervised task.
113///
114/// Can be used to abort the task by name independently of the supervisor.
115#[derive(Debug, Clone)]
116pub struct TaskHandle {
117 name: &'static str,
118 abort: AbortHandle,
119}
120
121impl TaskHandle {
122 /// Abort the task immediately.
123 pub fn abort(&self) {
124 tracing::debug!(task.name = self.name, "task aborted via handle");
125 self.abort.abort();
126 }
127
128 /// Return the task's name.
129 #[must_use]
130 pub const fn name(&self) -> &'static str {
131 self.name
132 }
133}
134
135/// Error returned by [`BlockingHandle::join`].
136#[derive(Debug, PartialEq, Eq)]
137#[non_exhaustive]
138pub enum BlockingError {
139 /// The task panicked before producing a result.
140 Panicked,
141 /// The supervisor (or the task's abort handle) was dropped before the task completed.
142 SupervisorDropped,
143}
144
145impl std::fmt::Display for BlockingError {
146 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147 match self {
148 Self::Panicked => write!(f, "supervised blocking task panicked"),
149 Self::SupervisorDropped => write!(f, "supervisor dropped before task completed"),
150 }
151 }
152}
153
154impl std::error::Error for BlockingError {}
155
156/// Handle returned by [`TaskSupervisor::spawn_blocking`].
157///
158/// Awaiting [`BlockingHandle::join`] blocks until the OS-thread task produces a
159/// value. Dropping the handle without joining does **not** cancel the task — it
160/// continues to run on the blocking thread pool but the result is discarded.
161///
162/// A panic inside the closure is captured and returned as
163/// [`BlockingError::Panicked`] rather than propagating to the caller.
164pub struct BlockingHandle<R> {
165 rx: oneshot::Receiver<Result<R, BlockingError>>,
166 abort: AbortHandle,
167}
168
169impl<R> BlockingHandle<R> {
170 /// Await the task result.
171 ///
172 /// # Errors
173 ///
174 /// - [`BlockingError::Panicked`] — the task closure panicked.
175 /// - [`BlockingError::SupervisorDropped`] — the task was aborted or the
176 /// supervisor was dropped before a value was produced.
177 pub async fn join(self) -> Result<R, BlockingError> {
178 self.rx
179 .await
180 .unwrap_or(Err(BlockingError::SupervisorDropped))
181 }
182
183 /// Non-blocking poll: return the result if the task has already finished, or `None`
184 /// if it is still running.
185 ///
186 /// This is the `BlockingHandle` equivalent of `FutureExt::now_or_never` on a
187 /// [`tokio::task::JoinHandle`]. Call this inside a synchronous context (e.g., between
188 /// agent turns) to apply a completed background result without blocking.
189 ///
190 /// The handle is consumed on success. If the task is not yet done, the handle
191 /// is returned as `Err(self)` so the caller can re-store it.
192 ///
193 /// # Examples
194 ///
195 /// ```rust,no_run
196 /// # use zeph_common::task_supervisor::{BlockingHandle, BlockingError};
197 /// async fn example(mut handle: BlockingHandle<u32>) {
198 /// // Try to get the result without blocking.
199 /// match handle.try_join() {
200 /// Ok(result) => println!("done: {result:?}"),
201 /// Err(handle) => {
202 /// // Task still running — `handle` is returned for re-storage.
203 /// drop(handle);
204 /// }
205 /// }
206 /// }
207 /// ```
208 ///
209 /// # Errors
210 ///
211 /// Returns `Err(self)` when the task has not yet produced a result (still running).
212 /// The inner `Ok(Err(BlockingError::...))` variants are returned when the task
213 /// panicked or the supervisor was dropped before the task completed.
214 pub fn try_join(mut self) -> Result<Result<R, BlockingError>, Self> {
215 match self.rx.try_recv() {
216 Ok(result) => Ok(result),
217 Err(tokio::sync::oneshot::error::TryRecvError::Empty) => Err(self),
218 Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
219 Ok(Err(BlockingError::SupervisorDropped))
220 }
221 }
222 }
223
224 /// Abort the underlying task immediately.
225 pub fn abort(&self) {
226 self.abort.abort();
227 }
228
229 /// Non-blocking, non-consuming check for whether the task has finished.
230 ///
231 /// Unlike [`try_join`][Self::try_join], this does not consume `self` or retrieve the
232 /// result — it only reports completion, including a panicked or aborted task. Useful
233 /// for reap sweeps that must detect a task exit even when the task's own
234 /// application-level status channel never got a chance to publish a terminal value
235 /// before exiting (e.g. a panic partway through the task body).
236 #[must_use]
237 pub fn is_finished(&self) -> bool {
238 self.abort.is_finished()
239 }
240}
241
242/// Point-in-time state of a supervised task.
243#[derive(Debug, Clone, PartialEq, Eq)]
244#[non_exhaustive]
245pub enum TaskStatus {
246 /// Task is actively running.
247 Running,
248 /// Task is waiting for the restart delay before the next attempt.
249 Restarting { attempt: u32, max: u32 },
250 /// Task completed normally.
251 Completed,
252 /// Task was force-aborted during shutdown.
253 Aborted,
254 /// Task exhausted all restart attempts and is permanently failed.
255 Failed { reason: String },
256}
257
258/// Point-in-time snapshot of a supervised task, returned by [`TaskSupervisor::snapshot`].
259#[derive(Debug, Clone)]
260/// Observability surface per field:
261///
262/// | Field | tokio-console | Jaeger / OTLP | TUI | `metrics` histogram |
263/// |-------|--------------|--------------|-----|---------------------|
264/// | `name` | span name | span name | task list | label `"task"` |
265/// | `task.wall_time_ms` | — | span field | — | `zeph.task.wall_time_ms` |
266/// | `task.cpu_time_ms` | — | span field | — | `zeph.task.cpu_time_ms` |
267/// | `status` | — | — | task list | — |
268/// | `restart_count` | — | — | task list | — |
269pub struct TaskSnapshot {
270 /// Task name.
271 pub name: Arc<str>,
272 /// Current status.
273 pub status: TaskStatus,
274 /// Instant the task was first spawned.
275 pub started_at: Instant,
276 /// Number of times the task has been restarted.
277 pub restart_count: u32,
278}
279
280// ── Internal types ───────────────────────────────────────────────────────────
281
282type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
283type BoxFactory = Box<dyn Fn() -> BoxFuture + Send + Sync>;
284
285struct TaskEntry {
286 name: Arc<str>,
287 status: TaskStatus,
288 started_at: Instant,
289 restart_count: u32,
290 restart_policy: RestartPolicy,
291 abort_handle: AbortHandle,
292 /// `Some` only for `Restart` policy tasks.
293 factory: Option<BoxFactory>,
294}
295
296/// How a supervised task ended.
297#[derive(Debug, Clone, Copy, PartialEq, Eq)]
298enum CompletionKind {
299 /// Future returned normally.
300 Normal,
301 /// Future returned a value that the caller-supplied classifier (see
302 /// [`TaskSupervisor::spawn_oneshot_classified`]) identified as an inner failure —
303 /// distinct from `Normal` even though the outer `JoinHandle` itself resolved
304 /// successfully (no panic, no cancellation).
305 Failed,
306 /// Future panicked.
307 Panicked,
308 /// Future was cancelled via the cancellation token or abort handle.
309 Cancelled,
310}
311
312struct Completion {
313 name: Arc<str>,
314 kind: CompletionKind,
315}
316
317struct SupervisorState {
318 tasks: HashMap<Arc<str>, TaskEntry>,
319}
320
321struct Inner {
322 state: parking_lot::Mutex<SupervisorState>,
323 /// Completion events from spawned tasks → reap driver.
324 /// Lives in `Inner` (not `SupervisorState`) to avoid double mutex acquisition
325 /// — callers clone it once during spawn without re-locking state.
326 completion_tx: mpsc::UnboundedSender<Completion>,
327 cancel: CancellationToken,
328 /// Limits the number of concurrently running `spawn_blocking` tasks to prevent
329 /// runaway thread-pool growth under burst load.
330 blocking_semaphore: Arc<tokio::sync::Semaphore>,
331 /// Notified when `active_count()` reaches zero so `shutdown_all` wakes immediately.
332 shutdown_notify: Arc<tokio::sync::Notify>,
333}
334
335// ── Main type ────────────────────────────────────────────────────────────────
336
337/// Shared, cloneable handle to the supervised lifecycle task registry.
338///
339/// `TaskSupervisor` manages named, long-lived background tasks with restart
340/// policies, health snapshots, and graceful shutdown. It is `Clone + Send + Sync`
341/// so it can be distributed to bootstrap code, TUI, and shutdown orchestration
342/// without any additional synchronisation.
343///
344/// # Thread safety
345///
346/// Interior state is guarded by a `parking_lot::Mutex`. The lock is **never**
347/// held across `.await` points.
348///
349/// # Examples
350///
351/// ```rust,no_run
352/// use std::time::Duration;
353/// use tokio_util::sync::CancellationToken;
354/// use zeph_common::task_supervisor::{RestartPolicy, TaskDescriptor, TaskSupervisor};
355///
356/// # #[tokio::main]
357/// # async fn main() {
358/// let cancel = CancellationToken::new();
359/// let sup = TaskSupervisor::new(cancel.clone());
360///
361/// let _handle = sup.spawn(TaskDescriptor {
362/// name: "watcher",
363/// restart: RestartPolicy::RunOnce,
364/// factory: || async { tokio::time::sleep(std::time::Duration::from_secs(1)).await },
365/// });
366///
367/// sup.shutdown_all(Duration::from_secs(5)).await;
368/// # }
369/// ```
370#[derive(Clone)]
371pub struct TaskSupervisor {
372 inner: Arc<Inner>,
373}
374
375impl std::fmt::Debug for TaskSupervisor {
376 /// Prints a stable placeholder rather than task internals (factory closures
377 /// inside `TaskEntry` are not `Debug`); callers needing task details should
378 /// use [`TaskSupervisor::snapshot`] instead.
379 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
380 f.debug_struct("TaskSupervisor").finish_non_exhaustive()
381 }
382}
383
384impl TaskSupervisor {
385 /// Create a new supervisor and start its reap driver.
386 ///
387 /// The `cancel` token is propagated into every spawned task via `tokio::select!`.
388 /// When the token is cancelled, all tasks exit cooperatively on their next
389 /// cancellation check. Call [`shutdown_all`][Self::shutdown_all] to wait for
390 /// them to finish.
391 ///
392 /// When called outside a Tokio runtime context (e.g. in synchronous unit tests),
393 /// the reap driver is skipped. The supervisor still accepts task registrations but
394 /// completion callbacks are not processed — safe because no tasks can actually be
395 /// spawned without a runtime.
396 #[must_use]
397 pub fn new(cancel: CancellationToken) -> Self {
398 // NOTE: unbounded channel is acceptable here because supervised tasks are
399 // O(10–20) lifecycle services, not high-throughput work. Backpressure would
400 // complicate the spawn path without practical benefit.
401 let (completion_tx, completion_rx) = mpsc::unbounded_channel();
402 let inner = Arc::new(Inner {
403 state: parking_lot::Mutex::new(SupervisorState {
404 tasks: HashMap::new(),
405 }),
406 completion_tx,
407 cancel: cancel.clone(),
408 blocking_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
409 shutdown_notify: Arc::new(tokio::sync::Notify::new()),
410 });
411
412 // Only start the reap driver when a Tokio runtime is available. In synchronous
413 // unit tests that construct Agent/LifecycleState directly there is no reactor,
414 // so we skip the spawn. Without a runtime no tasks can be spawned either, so
415 // the driver is not needed.
416 if tokio::runtime::Handle::try_current().is_ok() {
417 Self::start_reap_driver(Arc::clone(&inner), completion_rx, cancel);
418 }
419
420 Self { inner }
421 }
422
423 /// Spawn a named, supervised async task.
424 ///
425 /// If a task with the same `name` already exists, it is aborted before the
426 /// new one is started.
427 ///
428 /// # Examples
429 ///
430 /// ```rust,no_run
431 /// use std::time::Duration;
432 /// use tokio_util::sync::CancellationToken;
433 /// use zeph_common::task_supervisor::{RestartPolicy, TaskDescriptor, TaskHandle, TaskSupervisor};
434 ///
435 /// # #[tokio::main]
436 /// # async fn main() {
437 /// let cancel = CancellationToken::new();
438 /// let sup = TaskSupervisor::new(cancel.clone());
439 ///
440 /// let handle: TaskHandle = sup.spawn(TaskDescriptor {
441 /// name: "config-watcher",
442 /// restart: RestartPolicy::Restart { max: 3, base_delay: Duration::from_secs(1) },
443 /// factory: || async { /* watch loop */ },
444 /// });
445 /// # }
446 /// ```
447 pub fn spawn<F, Fut>(&self, desc: TaskDescriptor<F>) -> TaskHandle
448 where
449 F: Fn() -> Fut + Send + Sync + 'static,
450 Fut: Future<Output = ()> + Send + 'static,
451 {
452 let factory: BoxFactory = Box::new(move || Box::pin((desc.factory)()));
453 let cancel = self.inner.cancel.clone();
454 let completion_tx = self.inner.completion_tx.clone();
455 let name: Arc<str> = Arc::from(desc.name);
456
457 let (abort_handle, jh) = Self::do_spawn(desc.name, &factory, cancel);
458 Self::wire_completion_reporter(Arc::clone(&name), jh, completion_tx);
459
460 let entry = TaskEntry {
461 name: Arc::clone(&name),
462 status: TaskStatus::Running,
463 started_at: Instant::now(),
464 restart_count: 0,
465 restart_policy: desc.restart,
466 abort_handle: abort_handle.clone(),
467 factory: match desc.restart {
468 RestartPolicy::RunOnce => None,
469 RestartPolicy::Restart { .. } => Some(factory),
470 },
471 };
472
473 {
474 let mut state = self.inner.state.lock();
475 if let Some(old) = state.tasks.remove(&name) {
476 old.abort_handle.abort();
477 }
478 state.tasks.insert(Arc::clone(&name), entry);
479 }
480
481 TaskHandle {
482 name: desc.name,
483 abort: abort_handle,
484 }
485 }
486
487 /// Spawn a CPU-bound closure on the OS blocking thread pool.
488 ///
489 /// The closure runs via [`tokio::task::spawn_blocking`] — it is never polled
490 /// on tokio worker threads and cannot block async I/O. The task is registered
491 /// in the supervisor registry and is visible to [`snapshot`][Self::snapshot]
492 /// and [`shutdown_all`][Self::shutdown_all].
493 ///
494 /// Dropping the returned [`BlockingHandle`] without calling `.join()` does
495 /// **not** cancel the task; it runs to completion but the result is discarded.
496 ///
497 /// A panic inside `f` is captured and returned as [`BlockingError::Panicked`]
498 /// rather than propagating to the caller.
499 ///
500 /// # Examples
501 ///
502 /// ```rust,no_run
503 /// use std::sync::Arc;
504 /// use tokio_util::sync::CancellationToken;
505 /// use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
506 ///
507 /// # #[tokio::main]
508 /// # async fn main() {
509 /// let cancel = CancellationToken::new();
510 /// let sup = TaskSupervisor::new(cancel);
511 ///
512 /// let handle: BlockingHandle<u32> = sup.spawn_blocking(Arc::from("compute"), || {
513 /// // CPU-bound work — safe to block here
514 /// 42_u32
515 /// });
516 /// let result = handle.join().await.unwrap();
517 /// assert_eq!(result, 42);
518 /// # }
519 /// ```
520 ///
521 /// # Capacity limit
522 ///
523 /// At most 8 `spawn_blocking` tasks run concurrently. Additional tasks wait for a
524 /// semaphore permit, bounding thread-pool growth under burst load.
525 ///
526 /// # Panics
527 ///
528 /// Panics inside `f` are captured and returned as [`BlockingError::Panicked`] — they
529 /// do not propagate to the caller.
530 #[allow(clippy::needless_pass_by_value)] // `name` is cloned into async task and registry
531 pub fn spawn_blocking<F, R>(&self, name: Arc<str>, f: F) -> BlockingHandle<R>
532 where
533 F: FnOnce() -> R + Send + 'static,
534 R: Send + 'static,
535 {
536 let (tx, rx) = oneshot::channel::<Result<R, BlockingError>>();
537 let span = tracing::info_span!(
538 "supervised_blocking_task",
539 task.name = %name,
540 task.wall_time_ms = tracing::field::Empty,
541 task.cpu_time_ms = tracing::field::Empty,
542 );
543
544 let semaphore = Arc::clone(&self.inner.blocking_semaphore);
545 let inner = Arc::clone(&self.inner);
546 let name_clone = Arc::clone(&name);
547 let completion_tx = self.inner.completion_tx.clone();
548
549 // Wrap the blocking spawn in an async task that first acquires a semaphore
550 // permit, bounding the number of concurrently running blocking tasks to 8.
551 let outer = tokio::spawn(async move {
552 let _permit = semaphore
553 .acquire_owned()
554 .await
555 .expect("blocking semaphore closed");
556
557 let name_for_measure = Arc::clone(&name_clone);
558 let join_handle = tokio::task::spawn_blocking(move || {
559 let _enter = span.enter();
560 measure_blocking(&name_for_measure, f)
561 });
562 let abort = join_handle.abort_handle();
563
564 // Update registry with the real abort handle now that spawn_blocking is live.
565 {
566 let mut state = inner.state.lock();
567 if let Some(entry) = state.tasks.get_mut(&name_clone) {
568 entry.abort_handle = abort;
569 }
570 }
571
572 let kind = match join_handle.await {
573 Ok(val) => {
574 let _ = tx.send(Ok(val));
575 CompletionKind::Normal
576 }
577 Err(e) if e.is_panic() => {
578 let _ = tx.send(Err(BlockingError::Panicked));
579 CompletionKind::Panicked
580 }
581 Err(_) => {
582 // Aborted — drop tx so rx returns SupervisorDropped.
583 CompletionKind::Cancelled
584 }
585 };
586 // _permit released here, freeing the semaphore slot.
587 let _ = completion_tx.send(Completion {
588 name: name_clone,
589 kind,
590 });
591 });
592 let abort = outer.abort_handle();
593
594 // Register in registry so snapshot/shutdown sees the task.
595 {
596 let mut state = self.inner.state.lock();
597 if let Some(old) = state.tasks.remove(&name) {
598 old.abort_handle.abort();
599 }
600 state.tasks.insert(
601 Arc::clone(&name),
602 TaskEntry {
603 name: Arc::clone(&name),
604 status: TaskStatus::Running,
605 started_at: Instant::now(),
606 restart_count: 0,
607 restart_policy: RestartPolicy::RunOnce,
608 abort_handle: abort.clone(),
609 factory: None,
610 },
611 );
612 }
613
614 BlockingHandle { rx, abort }
615 }
616
617 /// Spawn an async task that produces a typed result value (runs on tokio worker thread).
618 ///
619 /// Unlike [`spawn`][Self::spawn], no restart policy is supported — the task
620 /// runs once. The task is registered in the supervisor registry under the
621 /// provided `name` and is visible to [`snapshot`][Self::snapshot] and
622 /// [`shutdown_all`][Self::shutdown_all].
623 ///
624 /// For CPU-bound work that must not block tokio workers, use
625 /// [`spawn_blocking`][Self::spawn_blocking] instead.
626 ///
627 /// # Examples
628 ///
629 /// ```rust,no_run
630 /// use std::sync::Arc;
631 /// use tokio_util::sync::CancellationToken;
632 /// use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
633 ///
634 /// # #[tokio::main]
635 /// # async fn main() {
636 /// let cancel = CancellationToken::new();
637 /// let sup = TaskSupervisor::new(cancel.clone());
638 ///
639 /// let handle: BlockingHandle<u32> = sup.spawn_oneshot(Arc::from("compute"), || async { 42_u32 });
640 /// let result = handle.join().await.unwrap();
641 /// assert_eq!(result, 42);
642 /// # }
643 /// ```
644 pub fn spawn_oneshot<F, Fut, R>(&self, name: Arc<str>, factory: F) -> BlockingHandle<R>
645 where
646 F: FnOnce() -> Fut + Send + 'static,
647 Fut: Future<Output = R> + Send + 'static,
648 R: Send + 'static,
649 {
650 self.spawn_oneshot_classified(name, factory, |_| true)
651 }
652
653 /// Spawn an async task like [`spawn_oneshot`][Self::spawn_oneshot], but additionally
654 /// classify the produced value's own success/failure via `is_success` instead of only
655 /// the outer `JoinHandle` outcome (normal exit vs. panic vs. cancellation).
656 ///
657 /// This matters whenever `Fut::Output` itself encodes failure (typically
658 /// `Result<T, E>`): without a classifier, a task that runs to completion but produces
659 /// `Err(..)` is classified and logged identically to one that produced `Ok(..)` —
660 /// `spawn_oneshot` alone only observes whether the *outer* future resolved, never what
661 /// value it resolved to. `is_success` is called once, by reference, before the value is
662 /// sent to the returned [`BlockingHandle`]'s receiver — it never consumes or blocks on
663 /// the value, and the value itself is delivered to the caller unchanged either way.
664 ///
665 /// # Examples
666 ///
667 /// ```rust,no_run
668 /// use std::sync::Arc;
669 /// use tokio_util::sync::CancellationToken;
670 /// use zeph_common::task_supervisor::TaskSupervisor;
671 ///
672 /// # #[tokio::main]
673 /// # async fn main() {
674 /// let sup = TaskSupervisor::new(CancellationToken::new());
675 /// let handle = sup.spawn_oneshot_classified(
676 /// Arc::from("fallible-task"),
677 /// || async { Result::<u32, String>::Err("boom".to_string()) },
678 /// Result::is_ok,
679 /// );
680 /// let result = handle.join().await.unwrap();
681 /// assert!(result.is_err());
682 /// # }
683 /// ```
684 pub fn spawn_oneshot_classified<F, Fut, R>(
685 &self,
686 name: Arc<str>,
687 factory: F,
688 is_success: impl Fn(&R) -> bool + Send + 'static,
689 ) -> BlockingHandle<R>
690 where
691 F: FnOnce() -> Fut + Send + 'static,
692 Fut: Future<Output = R> + Send + 'static,
693 R: Send + 'static,
694 {
695 let (tx, rx) = oneshot::channel::<Result<R, BlockingError>>();
696 let cancel = self.inner.cancel.clone();
697 let span = tracing::info_span!("supervised_task", task.name = %name);
698 let join_handle: tokio::task::JoinHandle<Option<R>> = tokio::spawn(
699 async move {
700 let fut = factory();
701 tokio::select! {
702 result = fut => Some(result),
703 () = cancel.cancelled() => None,
704 }
705 }
706 .instrument(span),
707 );
708 let abort = join_handle.abort_handle();
709
710 {
711 let mut state = self.inner.state.lock();
712 if let Some(old) = state.tasks.remove(&name) {
713 old.abort_handle.abort();
714 }
715 state.tasks.insert(
716 Arc::clone(&name),
717 TaskEntry {
718 name: Arc::clone(&name),
719 status: TaskStatus::Running,
720 started_at: Instant::now(),
721 restart_count: 0,
722 restart_policy: RestartPolicy::RunOnce,
723 abort_handle: abort.clone(),
724 factory: None,
725 },
726 );
727 }
728
729 let completion_tx = self.inner.completion_tx.clone();
730 tokio::spawn(async move {
731 let kind = match join_handle.await {
732 Ok(Some(val)) => {
733 let kind = if is_success(&val) {
734 CompletionKind::Normal
735 } else {
736 CompletionKind::Failed
737 };
738 let _ = tx.send(Ok(val));
739 kind
740 }
741 Err(e) if e.is_panic() => {
742 let _ = tx.send(Err(BlockingError::Panicked));
743 CompletionKind::Panicked
744 }
745 _ => CompletionKind::Cancelled,
746 };
747 let _ = completion_tx.send(Completion { name, kind });
748 });
749 BlockingHandle { rx, abort }
750 }
751
752 /// Abort a task by name. No-op if no task with that name is registered.
753 pub fn abort(&self, name: &'static str) {
754 let state = self.inner.state.lock();
755 let key: Arc<str> = Arc::from(name);
756 if let Some(entry) = state.tasks.get(&key) {
757 entry.abort_handle.abort();
758 tracing::debug!(task.name = name, "task aborted via supervisor");
759 }
760 }
761
762 /// Gracefully shut down all supervised tasks.
763 ///
764 /// Cancels the supervisor's [`CancellationToken`] and waits up to `timeout`
765 /// for all tasks to exit. Tasks that do not exit within the timeout are
766 /// aborted forcefully and their registry entries updated to [`TaskStatus::Aborted`].
767 ///
768 /// # Note
769 ///
770 /// This cancels the token passed to [`TaskSupervisor::new`]. If you share
771 /// that token with other subsystems, they will be cancelled too. Use a child
772 /// token (`cancel.child_token()`) when the supervisor should not affect
773 /// unrelated components.
774 pub async fn shutdown_all(&self, timeout: Duration) {
775 self.inner.cancel.cancel();
776 let sleep = tokio::time::sleep(timeout);
777 tokio::pin!(sleep);
778 loop {
779 let active = self.active_count();
780 if active == 0 {
781 break;
782 }
783 // Subscribe before re-checking so we cannot miss a notification that
784 // fires between the active_count() call above and the select below.
785 let notified = self.inner.shutdown_notify.notified();
786 tokio::select! {
787 biased;
788 () = notified => {
789 // reap driver decremented active count — re-check at top of loop
790 }
791 () = &mut sleep => {
792 let mut remaining_names: Vec<Arc<str>> = Vec::new();
793 {
794 let mut state = self.inner.state.lock();
795 for entry in state.tasks.values_mut() {
796 if matches!(
797 entry.status,
798 TaskStatus::Running | TaskStatus::Restarting { .. }
799 ) {
800 remaining_names.push(Arc::clone(&entry.name));
801 entry.abort_handle.abort();
802 entry.status = TaskStatus::Aborted;
803 }
804 }
805 }
806 tracing::warn!(
807 remaining = active,
808 tasks = ?remaining_names,
809 "shutdown timeout — aborting remaining tasks"
810 );
811 break;
812 }
813 }
814 }
815 }
816
817 /// Return a point-in-time snapshot of all registered tasks.
818 ///
819 /// Suitable for TUI status panels and structured logging. The returned
820 /// list is sorted by `started_at` ascending.
821 #[must_use]
822 pub fn snapshot(&self) -> Vec<TaskSnapshot> {
823 let state = self.inner.state.lock();
824 let mut snaps: Vec<TaskSnapshot> = state
825 .tasks
826 .values()
827 .map(|e| TaskSnapshot {
828 name: Arc::clone(&e.name),
829 status: e.status.clone(),
830 started_at: e.started_at,
831 restart_count: e.restart_count,
832 })
833 .collect();
834 snaps.sort_by_key(|s| s.started_at);
835 snaps
836 }
837
838 /// Return the number of tasks currently in `Running` or `Restarting` state.
839 #[must_use]
840 pub fn active_count(&self) -> usize {
841 let state = self.inner.state.lock();
842 state
843 .tasks
844 .values()
845 .filter(|e| {
846 matches!(
847 e.status,
848 TaskStatus::Running | TaskStatus::Restarting { .. }
849 )
850 })
851 .count()
852 }
853
854 /// Return a clone of the supervisor's [`CancellationToken`].
855 ///
856 /// Callers can use this to check whether shutdown has been initiated.
857 #[must_use]
858 pub fn cancellation_token(&self) -> CancellationToken {
859 self.inner.cancel.clone()
860 }
861
862 // ── Internal helpers ──────────────────────────────────────────────────────
863
864 /// Spawn the actual tokio task. Returns `(AbortHandle, JoinHandle)`.
865 fn do_spawn(
866 name: &'static str,
867 factory: &BoxFactory,
868 cancel: CancellationToken,
869 ) -> (AbortHandle, tokio::task::JoinHandle<()>) {
870 let fut = factory();
871 let span = tracing::info_span!("supervised_task", task.name = name);
872 let jh = tokio::spawn(
873 async move {
874 tokio::select! {
875 () = fut => {},
876 () = cancel.cancelled() => {},
877 }
878 }
879 .instrument(span),
880 );
881 let abort = jh.abort_handle();
882 (abort, jh)
883 }
884
885 /// Wire a completion reporter: drives `jh` and sends the result to `completion_tx`.
886 fn wire_completion_reporter(
887 name: Arc<str>,
888 jh: tokio::task::JoinHandle<()>,
889 completion_tx: mpsc::UnboundedSender<Completion>,
890 ) {
891 tokio::spawn(async move {
892 let kind = match jh.await {
893 Ok(()) => CompletionKind::Normal,
894 Err(e) if e.is_panic() => CompletionKind::Panicked,
895 Err(_) => CompletionKind::Cancelled,
896 };
897 let _ = completion_tx.send(Completion { name, kind });
898 });
899 }
900
901 /// Spawn the reap driver. The driver processes completion events from the mpsc channel.
902 ///
903 /// After the cancellation token fires, the driver continues draining the channel
904 /// until the registry reports no active tasks (or the channel closes) — this
905 /// ensures that tasks which complete after cancel, however long that takes,
906 /// still have their registry entries updated, allowing `shutdown_all` to observe
907 /// `active_count() == 0` correctly and wake early. The driver enforces no
908 /// deadline of its own: [`TaskSupervisor::shutdown_all`]'s own `sleep(timeout)`
909 /// is the sole authority for giving up on a stuck task, regardless of how long
910 /// the gap is between the token being cancelled (typically out-of-band, by a
911 /// shutdown bridge or signal handler — see `src/runner.rs`, `src/serve/mod.rs`)
912 /// and `shutdown_all` actually being invoked. See the Phase 2 comment below for
913 /// why this is correct and terminates.
914 fn start_reap_driver(
915 inner: Arc<Inner>,
916 mut completion_rx: mpsc::UnboundedReceiver<Completion>,
917 cancel: CancellationToken,
918 ) {
919 tokio::spawn(async move {
920 // Phase 1: normal operation — process completions until cancel fires.
921 loop {
922 tokio::select! {
923 biased;
924 Some(completion) = completion_rx.recv() => {
925 Self::handle_completion(&inner, completion).await;
926 }
927 () = cancel.cancelled() => break,
928 }
929 }
930
931 // Phase 2: post-cancel drain — keep receiving completions until the
932 // registry reports no active tasks, or the channel closes. This prevents
933 // losing completions that arrive after tasks observe cancellation
934 // (#3161), for however long that takes.
935 //
936 // Deliberately no independent deadline here (#5926): an earlier version
937 // of this drain phase enforced its own short fallback timeout, which
938 // could — and at every real call site, reliably did — expire before
939 // `shutdown_all` was ever invoked (the shared `CancellationToken` is
940 // cancelled out-of-band by a shutdown bridge/signal handler well before
941 // `shutdown_all` runs; see `src/runner.rs`, `src/serve/mod.rs`), causing
942 // the driver to exit and silently drop any later completion. There is
943 // exactly one deadline authority for the whole shutdown sequence:
944 // `shutdown_all`'s own `sleep(timeout)` (below), which force-aborts any
945 // task still `Running`/`Restarting` under the registry lock directly —
946 // that abort itself generates a `Cancelled` completion for `spawn`/
947 // `spawn_oneshot` tasks via `wire_completion_reporter`, waking this loop
948 // promptly. The one exception is a `spawn_blocking` task force-aborted
949 // while its outer wrapper is still parked on the concurrency semaphore
950 // (see `spawn_blocking`, below) — no completion is sent in that case, but
951 // `shutdown_all` has already marked the entry `Aborted` directly under the
952 // state lock and returned correctly; only this now-idle reap-driver task
953 // may linger parked until process/test teardown, which is not a hang. This
954 // task holds its own `Arc<Inner>` (which also owns a `completion_tx`
955 // clone) for as long as it runs, so `completion_rx.recv()` can only
956 // resolve via a real completion or the channel genuinely closing — never a
957 // stale timeout.
958 let active = Self::has_active_tasks(&inner);
959 tracing::debug!(active, "reap driver entered post-cancel drain phase");
960 loop {
961 if !Self::has_active_tasks(&inner) {
962 break;
963 }
964 match completion_rx.recv().await {
965 Some(completion) => Self::handle_completion(&inner, completion).await,
966 None => break, // channel closed — unreachable in practice
967 }
968 }
969 tracing::debug!(
970 active = Self::has_active_tasks(&inner),
971 "reap driver drain phase complete"
972 );
973 });
974 }
975
976 /// Returns `true` if any task is in `Running` or `Restarting` state.
977 fn has_active_tasks(inner: &Arc<Inner>) -> bool {
978 let state = inner.state.lock();
979 state.tasks.values().any(|e| {
980 matches!(
981 e.status,
982 TaskStatus::Running | TaskStatus::Restarting { .. }
983 )
984 })
985 }
986
987 /// Process a single task completion event.
988 ///
989 /// Lock is never held across `.await`. Phase 1 classifies the completion
990 /// under lock; Phase 2 sleeps with exponential backoff without a lock;
991 /// Phase 3 spawns the next instance and updates the registry.
992 async fn handle_completion(inner: &Arc<Inner>, completion: Completion) {
993 // Short-circuit: once cancellation has fired, never schedule restarts.
994 // Without this, Restart-policy tasks re-register as Running, causing
995 // has_active_tasks() to stay true and the drain loop to spin until timeout.
996 if inner.cancel.is_cancelled() {
997 {
998 let mut state = inner.state.lock();
999 state.tasks.remove(&completion.name);
1000 }
1001 inner.shutdown_notify.notify_waiters();
1002 return;
1003 }
1004
1005 let Some((attempt, max, delay)) = Self::classify_completion(inner, &completion) else {
1006 // Task removed from registry (RunOnce completed or Restart exhausted) —
1007 // wake shutdown_all so it can re-check active_count immediately.
1008 inner.shutdown_notify.notify_waiters();
1009 return;
1010 };
1011
1012 tracing::warn!(
1013 task.name = %completion.name,
1014 attempt,
1015 max,
1016 delay_ms = delay.as_millis(),
1017 "restarting supervised task"
1018 );
1019
1020 if !delay.is_zero() {
1021 tokio::time::sleep(delay).await;
1022 }
1023
1024 Self::do_restart(inner, &completion.name, attempt);
1025 }
1026
1027 /// Phase 1: classify the completion under lock and return restart parameters if needed.
1028 ///
1029 /// Returns `Some((attempt, max, backoff_delay))` when a restart should be scheduled.
1030 fn classify_completion(
1031 inner: &Arc<Inner>,
1032 completion: &Completion,
1033 ) -> Option<(u32, u32, Duration)> {
1034 let mut state = inner.state.lock();
1035 let entry = state.tasks.get_mut(&completion.name)?;
1036
1037 match completion.kind {
1038 CompletionKind::Panicked => {
1039 tracing::warn!(task.name = %completion.name, "supervised task panicked");
1040 }
1041 CompletionKind::Failed => {
1042 tracing::warn!(
1043 task.name = %completion.name,
1044 "supervised task completed with an inner failure"
1045 );
1046 }
1047 CompletionKind::Normal => {
1048 tracing::info!(task.name = %completion.name, "supervised task completed");
1049 }
1050 CompletionKind::Cancelled => {
1051 tracing::debug!(task.name = %completion.name, "supervised task cancelled");
1052 }
1053 }
1054
1055 match entry.restart_policy {
1056 RestartPolicy::RunOnce => {
1057 entry.status = if completion.kind == CompletionKind::Failed {
1058 TaskStatus::Failed {
1059 reason: "task completed with an inner failure".to_string(),
1060 }
1061 } else {
1062 TaskStatus::Completed
1063 };
1064 state.tasks.remove(&completion.name);
1065 None
1066 }
1067 RestartPolicy::Restart { max, base_delay } => {
1068 // Only restart on panic — normal exit and cancellation are not errors.
1069 if completion.kind != CompletionKind::Panicked {
1070 entry.status = TaskStatus::Completed;
1071 state.tasks.remove(&completion.name);
1072 return None;
1073 }
1074 if entry.restart_count >= max {
1075 let reason = format!("panicked after {max} restart(s)");
1076 tracing::error!(
1077 task.name = %completion.name,
1078 attempts = max,
1079 "task failed permanently"
1080 );
1081 entry.status = TaskStatus::Failed { reason };
1082 None
1083 } else {
1084 let attempt = entry.restart_count + 1;
1085 entry.status = TaskStatus::Restarting { attempt, max };
1086 // Exponential backoff: base_delay * 2^(attempt-1), capped at MAX_RESTART_DELAY.
1087 let multiplier = 1_u32
1088 .checked_shl(attempt.saturating_sub(1))
1089 .unwrap_or(u32::MAX);
1090 let delay = base_delay.saturating_mul(multiplier).min(MAX_RESTART_DELAY);
1091 Some((attempt, max, delay))
1092 }
1093 }
1094 }
1095 // lock released here
1096 }
1097
1098 /// Phase 3: TOCTOU check, collect spawn params under lock, then spawn outside.
1099 fn do_restart(inner: &Arc<Inner>, name: &Arc<str>, attempt: u32) {
1100 let spawn_params = {
1101 let mut state = inner.state.lock();
1102 let Some(entry) = state.tasks.get_mut(name.as_ref()) else {
1103 tracing::debug!(
1104 task.name = %name,
1105 "task removed during restart delay — skipping"
1106 );
1107 return;
1108 };
1109 if !matches!(entry.status, TaskStatus::Restarting { .. }) {
1110 return;
1111 }
1112 let Some(factory) = &entry.factory else {
1113 return;
1114 };
1115 // Wrap factory() in catch_unwind to prevent a factory panic from crashing
1116 // the reap driver and orphaning the registry.
1117 match std::panic::catch_unwind(std::panic::AssertUnwindSafe(factory)) {
1118 Err(_) => {
1119 let reason = format!("factory panicked on restart attempt {attempt}");
1120 tracing::error!(task.name = %name, attempt, "factory panicked during restart");
1121 entry.status = TaskStatus::Failed { reason };
1122 None
1123 }
1124 Ok(fut) => Some((
1125 fut,
1126 inner.cancel.clone(),
1127 inner.completion_tx.clone(),
1128 name.clone(),
1129 )),
1130 }
1131 // lock released here
1132 };
1133
1134 let Some((fut, cancel, completion_tx, name)) = spawn_params else {
1135 return;
1136 };
1137
1138 let span = tracing::info_span!("supervised_task", task.name = %name);
1139 let jh = tokio::spawn(
1140 async move {
1141 tokio::select! {
1142 () = fut => {},
1143 () = cancel.cancelled() => {},
1144 }
1145 }
1146 .instrument(span),
1147 );
1148 let new_abort = jh.abort_handle();
1149
1150 {
1151 let mut state = inner.state.lock();
1152 if let Some(entry) = state.tasks.get_mut(name.as_ref()) {
1153 entry.restart_count = attempt;
1154 entry.status = TaskStatus::Running;
1155 entry.abort_handle = new_abort;
1156 }
1157 }
1158
1159 Self::wire_completion_reporter(name, jh, completion_tx);
1160 }
1161}
1162
1163// ── Task metrics helpers ──────────────────────────────────────────────────────
1164
1165/// Run `f` and record wall-time and CPU-time metrics via `metrics` crate.
1166#[inline]
1167fn measure_blocking<F, R>(name: &str, f: F) -> R
1168where
1169 F: FnOnce() -> R,
1170{
1171 use cpu_time::ThreadTime;
1172 let wall_start = std::time::Instant::now();
1173 let cpu_start = ThreadTime::now();
1174 let result = f();
1175 let wall_ms = wall_start.elapsed().as_secs_f64() * 1000.0;
1176 let cpu_ms = cpu_start.elapsed().as_secs_f64() * 1000.0;
1177 metrics::histogram!("zeph.task.wall_time_ms", "task" => name.to_owned()).record(wall_ms);
1178 metrics::histogram!("zeph.task.cpu_time_ms", "task" => name.to_owned()).record(cpu_ms);
1179 tracing::Span::current().record("task.wall_time_ms", wall_ms);
1180 tracing::Span::current().record("task.cpu_time_ms", cpu_ms);
1181 result
1182}
1183
1184// ── BlockingSpawner impl ──────────────────────────────────────────────────────
1185
1186impl BlockingSpawner for TaskSupervisor {
1187 /// Spawn a named blocking closure through the supervisor.
1188 ///
1189 /// The task is registered in the supervisor registry (visible in
1190 /// [`snapshot`][Self::snapshot] and subject to graceful shutdown) before
1191 /// the closure begins executing.
1192 fn spawn_blocking_named(
1193 &self,
1194 name: Arc<str>,
1195 f: Box<dyn FnOnce() + Send + 'static>,
1196 ) -> tokio::task::JoinHandle<()> {
1197 let handle = self.spawn_blocking(Arc::clone(&name), f);
1198 tokio::spawn(async move {
1199 if let Err(e) = handle.join().await {
1200 tracing::error!(task.name = %name, error = %e, "supervised blocking task failed");
1201 }
1202 })
1203 }
1204}
1205
1206// ── Unit tests ────────────────────────────────────────────────────────────────
1207
1208#[cfg(test)]
1209mod tests {
1210 use std::sync::Arc;
1211 use std::sync::atomic::{AtomicU32, Ordering};
1212 use std::time::Duration;
1213
1214 use tokio_util::sync::CancellationToken;
1215
1216 use super::*;
1217
1218 fn make_supervisor() -> (TaskSupervisor, CancellationToken) {
1219 let cancel = CancellationToken::new();
1220 let sup = TaskSupervisor::new(cancel.clone());
1221 (sup, cancel)
1222 }
1223
1224 #[tokio::test]
1225 async fn test_spawn_and_complete() {
1226 let (sup, _cancel) = make_supervisor();
1227
1228 let done = Arc::new(tokio::sync::Notify::new());
1229 let done2 = Arc::clone(&done);
1230
1231 sup.spawn(TaskDescriptor {
1232 name: "simple",
1233 restart: RestartPolicy::RunOnce,
1234 factory: move || {
1235 let d = Arc::clone(&done2);
1236 async move {
1237 d.notify_one();
1238 }
1239 },
1240 });
1241
1242 tokio::time::timeout(Duration::from_secs(2), done.notified())
1243 .await
1244 .expect("task should complete");
1245
1246 tokio::time::sleep(Duration::from_millis(50)).await;
1247 assert_eq!(
1248 sup.active_count(),
1249 0,
1250 "RunOnce task should be removed after completion"
1251 );
1252 }
1253
1254 #[tokio::test]
1255 async fn test_panic_capture() {
1256 let (sup, _cancel) = make_supervisor();
1257
1258 sup.spawn(TaskDescriptor {
1259 name: "panicking",
1260 restart: RestartPolicy::RunOnce,
1261 factory: || async { panic!("intentional test panic") },
1262 });
1263
1264 tokio::time::sleep(Duration::from_millis(200)).await;
1265
1266 let snaps = sup.snapshot();
1267 assert!(
1268 snaps.iter().all(|s| s.name.as_ref() != "panicking"),
1269 "entry should be reaped"
1270 );
1271 assert_eq!(
1272 sup.active_count(),
1273 0,
1274 "active count must be 0 after RunOnce panic"
1275 );
1276 }
1277
1278 /// Regression test for S2: Restart-policy tasks must only restart on panic,
1279 /// not on normal completion.
1280 #[tokio::test]
1281 async fn test_restart_only_on_panic() {
1282 let (sup, _cancel) = make_supervisor();
1283
1284 // Part 1: normal completion — must NOT restart.
1285 let normal_counter = Arc::new(AtomicU32::new(0));
1286 let nc = Arc::clone(&normal_counter);
1287 sup.spawn(TaskDescriptor {
1288 name: "normal-exit",
1289 restart: RestartPolicy::Restart {
1290 max: 3,
1291 base_delay: Duration::from_millis(10),
1292 },
1293 factory: move || {
1294 let c = Arc::clone(&nc);
1295 async move {
1296 c.fetch_add(1, Ordering::SeqCst);
1297 // Returns normally — no panic.
1298 }
1299 },
1300 });
1301
1302 tokio::time::sleep(Duration::from_millis(300)).await;
1303 assert_eq!(
1304 normal_counter.load(Ordering::SeqCst),
1305 1,
1306 "normal exit must not restart"
1307 );
1308 assert!(
1309 sup.snapshot()
1310 .iter()
1311 .all(|s| s.name.as_ref() != "normal-exit"),
1312 "entry removed after normal exit"
1313 );
1314
1315 // Part 2: panic — MUST restart up to max times.
1316 let panic_counter = Arc::new(AtomicU32::new(0));
1317 let pc = Arc::clone(&panic_counter);
1318 sup.spawn(TaskDescriptor {
1319 name: "panic-exit",
1320 restart: RestartPolicy::Restart {
1321 max: 2,
1322 base_delay: Duration::from_millis(10),
1323 },
1324 factory: move || {
1325 let c = Arc::clone(&pc);
1326 async move {
1327 c.fetch_add(1, Ordering::SeqCst);
1328 panic!("test panic");
1329 }
1330 },
1331 });
1332
1333 // initial + 2 restarts = 3 total
1334 tokio::time::sleep(Duration::from_millis(500)).await;
1335 assert!(
1336 panic_counter.load(Ordering::SeqCst) >= 3,
1337 "panicking task must restart max times"
1338 );
1339 let snap = sup
1340 .snapshot()
1341 .into_iter()
1342 .find(|s| s.name.as_ref() == "panic-exit");
1343 assert!(
1344 matches!(snap.unwrap().status, TaskStatus::Failed { .. }),
1345 "task must be Failed after exhausting restarts"
1346 );
1347 }
1348
1349 #[tokio::test]
1350 async fn test_restart_policy() {
1351 let (sup, _cancel) = make_supervisor();
1352
1353 let counter = Arc::new(AtomicU32::new(0));
1354 let counter2 = Arc::clone(&counter);
1355
1356 sup.spawn(TaskDescriptor {
1357 name: "restartable",
1358 restart: RestartPolicy::Restart {
1359 max: 2,
1360 base_delay: Duration::from_millis(10),
1361 },
1362 factory: move || {
1363 let c = Arc::clone(&counter2);
1364 async move {
1365 c.fetch_add(1, Ordering::SeqCst);
1366 panic!("always panic");
1367 }
1368 },
1369 });
1370
1371 tokio::time::sleep(Duration::from_millis(500)).await;
1372
1373 let runs = counter.load(Ordering::SeqCst);
1374 assert!(
1375 runs >= 3,
1376 "expected at least 3 invocations (initial + 2 restarts), got {runs}"
1377 );
1378
1379 let snaps = sup.snapshot();
1380 let snap = snaps.iter().find(|s| s.name.as_ref() == "restartable");
1381 assert!(snap.is_some(), "failed task should remain in registry");
1382 assert!(
1383 matches!(snap.unwrap().status, TaskStatus::Failed { .. }),
1384 "task should be Failed after exhausting retries"
1385 );
1386 }
1387
1388 /// Verify exponential backoff: delay doubles on each restart attempt.
1389 #[tokio::test]
1390 async fn test_exponential_backoff() {
1391 let (sup, _cancel) = make_supervisor();
1392
1393 let timestamps = Arc::new(parking_lot::Mutex::new(Vec::<std::time::Instant>::new()));
1394 let ts = Arc::clone(×tamps);
1395
1396 sup.spawn(TaskDescriptor {
1397 name: "backoff-task",
1398 restart: RestartPolicy::Restart {
1399 max: 3,
1400 base_delay: Duration::from_millis(50),
1401 },
1402 factory: move || {
1403 let t = Arc::clone(&ts);
1404 async move {
1405 t.lock().push(std::time::Instant::now());
1406 panic!("always panic");
1407 }
1408 },
1409 });
1410
1411 // Wait long enough for all restarts: 50 + 100 + 200 ms = 350 ms + overhead
1412 tokio::time::sleep(Duration::from_millis(800)).await;
1413
1414 let ts = timestamps.lock();
1415 assert!(
1416 ts.len() >= 3,
1417 "expected at least 3 invocations, got {}",
1418 ts.len()
1419 );
1420
1421 // Verify delays are roughly doubling (within 2x tolerance for CI jitter).
1422 if ts.len() >= 3 {
1423 let d1 = ts[1].duration_since(ts[0]);
1424 let d2 = ts[2].duration_since(ts[1]);
1425 // d2 should be at least 1.5x d1 (allowing for jitter).
1426 assert!(
1427 d2 >= d1.mul_f64(1.5),
1428 "expected exponential backoff: d1={d1:?} d2={d2:?}"
1429 );
1430 }
1431 }
1432
1433 #[tokio::test]
1434 async fn test_graceful_shutdown() {
1435 let (sup, _cancel) = make_supervisor();
1436
1437 for name in ["svc-a", "svc-b", "svc-c"] {
1438 sup.spawn(TaskDescriptor {
1439 name,
1440 restart: RestartPolicy::RunOnce,
1441 factory: || async {
1442 tokio::time::sleep(Duration::from_mins(1)).await;
1443 },
1444 });
1445 }
1446
1447 assert_eq!(sup.active_count(), 3);
1448
1449 tokio::time::timeout(
1450 Duration::from_secs(2),
1451 sup.shutdown_all(Duration::from_secs(1)),
1452 )
1453 .await
1454 .expect("shutdown should complete within timeout");
1455 }
1456
1457 /// Verify that force-aborted tasks get `TaskStatus::Aborted` in the registry (A2 fix).
1458 #[tokio::test]
1459 async fn test_force_abort_marks_aborted() {
1460 let cancel = CancellationToken::new();
1461 let sup = TaskSupervisor::new(cancel.clone());
1462
1463 sup.spawn(TaskDescriptor {
1464 name: "stubborn-for-abort",
1465 restart: RestartPolicy::RunOnce,
1466 factory: || async {
1467 // Does not cooperate with cancellation.
1468 std::future::pending::<()>().await;
1469 },
1470 });
1471
1472 // Use a very short timeout to trigger force-abort.
1473 sup.shutdown_all(Duration::from_millis(1)).await;
1474
1475 // Entry should be Aborted, not Running.
1476 let snaps = sup.snapshot();
1477 if let Some(snap) = snaps
1478 .iter()
1479 .find(|s| s.name.as_ref() == "stubborn-for-abort")
1480 {
1481 assert_eq!(
1482 snap.status,
1483 TaskStatus::Aborted,
1484 "force-aborted task must have Aborted status"
1485 );
1486 }
1487 // If entry was already reaped (cooperative cancel won), that's also acceptable.
1488 }
1489
1490 #[tokio::test]
1491 async fn test_registry_snapshot() {
1492 let (sup, _cancel) = make_supervisor();
1493
1494 for name in ["alpha", "beta"] {
1495 sup.spawn(TaskDescriptor {
1496 name,
1497 restart: RestartPolicy::RunOnce,
1498 factory: || async {
1499 tokio::time::sleep(Duration::from_secs(10)).await;
1500 },
1501 });
1502 }
1503
1504 let snaps = sup.snapshot();
1505 assert_eq!(snaps.len(), 2);
1506 let names: Vec<&str> = snaps.iter().map(|s| s.name.as_ref()).collect();
1507 assert!(names.contains(&"alpha"));
1508 assert!(names.contains(&"beta"));
1509 assert!(snaps.iter().all(|s| s.status == TaskStatus::Running));
1510 }
1511
1512 #[tokio::test]
1513 async fn test_blocking_returns_value() {
1514 let (sup, cancel) = make_supervisor();
1515
1516 let handle: BlockingHandle<u32> = sup.spawn_blocking(Arc::from("compute"), || 42_u32);
1517 let result = handle.join().await.expect("should return value");
1518 assert_eq!(result, 42);
1519 cancel.cancel();
1520 }
1521
1522 #[tokio::test]
1523 async fn test_blocking_panic() {
1524 let (sup, _cancel) = make_supervisor();
1525
1526 let handle: BlockingHandle<u32> =
1527 sup.spawn_blocking(Arc::from("panicking-compute"), || panic!("intentional"));
1528 let err = handle
1529 .join()
1530 .await
1531 .expect_err("should return error on panic");
1532 assert_eq!(err, BlockingError::Panicked);
1533 }
1534
1535 /// Verify `spawn_blocking` tasks appear in registry (M3 fix).
1536 #[tokio::test]
1537 async fn test_blocking_registered_in_registry() {
1538 let (sup, cancel) = make_supervisor();
1539
1540 let (tx, rx) = std::sync::mpsc::channel::<()>();
1541 let _handle: BlockingHandle<()> =
1542 sup.spawn_blocking(Arc::from("blocking-task"), move || {
1543 // Block until signalled.
1544 let _ = rx.recv();
1545 });
1546
1547 tokio::time::sleep(Duration::from_millis(10)).await;
1548 assert_eq!(
1549 sup.active_count(),
1550 1,
1551 "blocking task must appear in active_count"
1552 );
1553
1554 let _ = tx.send(());
1555 tokio::time::sleep(Duration::from_millis(100)).await;
1556 assert_eq!(
1557 sup.active_count(),
1558 0,
1559 "blocking task must be removed after completion"
1560 );
1561
1562 cancel.cancel();
1563 }
1564
1565 /// Verify `spawn_oneshot` tasks appear in registry (M3 fix).
1566 #[tokio::test]
1567 async fn test_oneshot_registered_in_registry() {
1568 let (sup, cancel) = make_supervisor();
1569
1570 let (tx, rx) = tokio::sync::oneshot::channel::<()>();
1571 let _handle: BlockingHandle<()> =
1572 sup.spawn_oneshot(Arc::from("oneshot-task"), move || async move {
1573 let _ = rx.await;
1574 });
1575
1576 tokio::time::sleep(Duration::from_millis(10)).await;
1577 assert_eq!(
1578 sup.active_count(),
1579 1,
1580 "oneshot task must appear in active_count"
1581 );
1582
1583 let _ = tx.send(());
1584 tokio::time::sleep(Duration::from_millis(50)).await;
1585 assert_eq!(
1586 sup.active_count(),
1587 0,
1588 "oneshot task must be removed after completion"
1589 );
1590
1591 cancel.cancel();
1592 }
1593
1594 /// `spawn_oneshot_classified` must deliver the factory's actual return value to the
1595 /// caller unchanged, regardless of how `is_success` classifies it — the classifier only
1596 /// affects internal registry/log bookkeeping, never the value received via `join()`.
1597 #[tokio::test]
1598 async fn test_oneshot_classified_delivers_actual_value_on_success_and_failure() {
1599 let (sup, _cancel) = make_supervisor();
1600
1601 let ok_handle = sup.spawn_oneshot_classified(
1602 Arc::from("classified-ok"),
1603 || async { Result::<u32, String>::Ok(7) },
1604 Result::is_ok,
1605 );
1606 assert_eq!(ok_handle.join().await.unwrap(), Ok(7));
1607
1608 let err_handle = sup.spawn_oneshot_classified(
1609 Arc::from("classified-err"),
1610 || async { Result::<u32, String>::Err("boom".to_string()) },
1611 Result::is_ok,
1612 );
1613 assert_eq!(
1614 err_handle.join().await.unwrap(),
1615 Err("boom".to_string()),
1616 "the real Err value must reach the caller unchanged even though it is \
1617 classified as CompletionKind::Failed internally"
1618 );
1619 }
1620
1621 /// Regression test for #6257: an `Err`-returning factory must be classified as
1622 /// `CompletionKind::Failed` (logged at `warn` as "completed with an inner failure"),
1623 /// not `CompletionKind::Normal` (logged at `info` as a plain completion). This is the
1624 /// distinction `spawn_agent_task` (zeph-subagent) relies on so a subagent task that
1625 /// returns `Err` after a genuine completion (e.g. a worktree-setup failure) is reported
1626 /// as a supervisor-level failure instead of a normal completion. The classification
1627 /// only surfaces via the reap driver's log lines — for `RunOnce` the registry entry's
1628 /// `Failed`/`Completed` status assignment and its removal happen in the same locked
1629 /// critical section (see `classify_completion`), so there is no window in which
1630 /// `snapshot()` could observe the transient status; the log line is the only externally
1631 /// observable signal of which branch was taken.
1632 #[tokio::test]
1633 #[tracing_test::traced_test]
1634 async fn test_oneshot_classified_logs_failed_for_inner_err() {
1635 let (sup, _cancel) = make_supervisor();
1636
1637 let handle = sup.spawn_oneshot_classified(
1638 Arc::from("classified-inner-err"),
1639 || async { Result::<u32, String>::Err("boom".to_string()) },
1640 Result::is_ok,
1641 );
1642 handle.join().await.unwrap().unwrap_err();
1643
1644 // Give the reap driver a moment to process the completion event.
1645 tokio::time::sleep(Duration::from_millis(50)).await;
1646
1647 assert!(
1648 logs_contain("completed with an inner failure"),
1649 "an Err value must be classified as CompletionKind::Failed, not Normal"
1650 );
1651 }
1652
1653 /// Counterpart to the above: an `Ok`-returning factory must be classified as
1654 /// `CompletionKind::Normal` (logged at `info`), never as an inner failure — confirms
1655 /// the classifier doesn't over-fire on the success path.
1656 #[tokio::test]
1657 #[tracing_test::traced_test]
1658 async fn test_oneshot_classified_logs_normal_for_inner_ok() {
1659 let (sup, _cancel) = make_supervisor();
1660
1661 let handle = sup.spawn_oneshot_classified(
1662 Arc::from("classified-inner-ok"),
1663 || async { Result::<u32, String>::Ok(1) },
1664 Result::is_ok,
1665 );
1666 handle.join().await.unwrap().unwrap();
1667
1668 tokio::time::sleep(Duration::from_millis(50)).await;
1669
1670 assert!(
1671 !logs_contain("completed with an inner failure"),
1672 "an Ok value must not be classified as CompletionKind::Failed"
1673 );
1674 }
1675
1676 /// Regression guard for #6257: `spawn_oneshot`'s ~20 existing call sites across the
1677 /// workspace must remain unaffected by the addition of `spawn_oneshot_classified` —
1678 /// `spawn_oneshot` delegates with an always-true classifier (`|_| true`), so a plain
1679 /// (non-`Result`) factory value must still always be classified `Normal`/logged as a
1680 /// plain completion, never as an inner failure.
1681 #[tokio::test]
1682 #[tracing_test::traced_test]
1683 async fn test_oneshot_backward_compat_never_logs_inner_failure() {
1684 let (sup, _cancel) = make_supervisor();
1685
1686 let handle: BlockingHandle<u32> =
1687 sup.spawn_oneshot(Arc::from("plain-oneshot"), || async { 42_u32 });
1688 assert_eq!(handle.join().await.unwrap(), 42);
1689
1690 tokio::time::sleep(Duration::from_millis(50)).await;
1691
1692 assert!(
1693 !logs_contain("completed with an inner failure"),
1694 "spawn_oneshot's default always-true classifier must never report an inner failure"
1695 );
1696 }
1697
1698 #[tokio::test]
1699 async fn test_restart_max_zero() {
1700 let (sup, _cancel) = make_supervisor();
1701
1702 let counter = Arc::new(AtomicU32::new(0));
1703 let counter2 = Arc::clone(&counter);
1704
1705 sup.spawn(TaskDescriptor {
1706 name: "zero-max",
1707 restart: RestartPolicy::Restart {
1708 max: 0,
1709 base_delay: Duration::from_millis(10),
1710 },
1711 factory: move || {
1712 let c = Arc::clone(&counter2);
1713 async move {
1714 c.fetch_add(1, Ordering::SeqCst);
1715 panic!("always panic");
1716 }
1717 },
1718 });
1719
1720 tokio::time::sleep(Duration::from_millis(200)).await;
1721
1722 assert_eq!(
1723 counter.load(Ordering::SeqCst),
1724 1,
1725 "max=0 should not restart"
1726 );
1727
1728 let snaps = sup.snapshot();
1729 let snap = snaps.iter().find(|s| s.name.as_ref() == "zero-max");
1730 assert!(snap.is_some(), "entry should remain as Failed");
1731 assert!(
1732 matches!(snap.unwrap().status, TaskStatus::Failed { .. }),
1733 "status should be Failed"
1734 );
1735 }
1736
1737 /// Stress test: spawn 50 tasks concurrently, all must complete and registry must be accurate.
1738 #[tokio::test]
1739 async fn test_concurrent_spawns() {
1740 // All task names must be 'static — pre-defined before any let statements.
1741 static NAMES: [&str; 50] = [
1742 "t00", "t01", "t02", "t03", "t04", "t05", "t06", "t07", "t08", "t09", "t10", "t11",
1743 "t12", "t13", "t14", "t15", "t16", "t17", "t18", "t19", "t20", "t21", "t22", "t23",
1744 "t24", "t25", "t26", "t27", "t28", "t29", "t30", "t31", "t32", "t33", "t34", "t35",
1745 "t36", "t37", "t38", "t39", "t40", "t41", "t42", "t43", "t44", "t45", "t46", "t47",
1746 "t48", "t49",
1747 ];
1748 let (sup, cancel) = make_supervisor();
1749
1750 let completed = Arc::new(AtomicU32::new(0));
1751 for name in &NAMES {
1752 let c = Arc::clone(&completed);
1753 sup.spawn(TaskDescriptor {
1754 name,
1755 restart: RestartPolicy::RunOnce,
1756 factory: move || {
1757 let c = Arc::clone(&c);
1758 async move {
1759 c.fetch_add(1, Ordering::SeqCst);
1760 }
1761 },
1762 });
1763 }
1764
1765 // Wait for all tasks to complete.
1766 tokio::time::timeout(Duration::from_secs(5), async {
1767 loop {
1768 if completed.load(Ordering::SeqCst) == 50 {
1769 break;
1770 }
1771 tokio::time::sleep(Duration::from_millis(10)).await;
1772 }
1773 })
1774 .await
1775 .expect("all 50 tasks should complete");
1776
1777 // Give reap driver time to process all completions.
1778 tokio::time::sleep(Duration::from_millis(100)).await;
1779 assert_eq!(sup.active_count(), 0, "all tasks must be reaped");
1780
1781 cancel.cancel();
1782 }
1783
1784 #[tokio::test]
1785 async fn test_shutdown_timeout_expiry() {
1786 let cancel = CancellationToken::new();
1787 let sup = TaskSupervisor::new(cancel.clone());
1788
1789 sup.spawn(TaskDescriptor {
1790 name: "stubborn",
1791 restart: RestartPolicy::RunOnce,
1792 factory: || async {
1793 tokio::time::sleep(Duration::from_mins(1)).await;
1794 },
1795 });
1796
1797 assert_eq!(sup.active_count(), 1);
1798
1799 tokio::time::timeout(
1800 Duration::from_secs(2),
1801 sup.shutdown_all(Duration::from_millis(50)),
1802 )
1803 .await
1804 .expect("shutdown_all should return even on timeout expiry");
1805
1806 assert!(
1807 cancel.is_cancelled(),
1808 "cancel token must be cancelled after shutdown"
1809 );
1810 }
1811
1812 #[tokio::test]
1813 async fn test_cancellation_token() {
1814 let cancel = CancellationToken::new();
1815 let sup = TaskSupervisor::new(cancel.clone());
1816
1817 assert!(!sup.cancellation_token().is_cancelled());
1818
1819 sup.shutdown_all(Duration::from_millis(100)).await;
1820
1821 assert!(
1822 sup.cancellation_token().is_cancelled(),
1823 "token must be cancelled after shutdown"
1824 );
1825 }
1826
1827 /// Regression test for #3161: after `shutdown_all`, all tasks must be reaped
1828 /// even when they complete *after* the cancel signal.
1829 ///
1830 /// The yield loop forces the reap driver to observe cancel and exit phase-1
1831 /// before the tasks send their completions — reliably reproducing the race.
1832 #[tokio::test]
1833 async fn test_shutdown_drains_post_cancel_completions() {
1834 let cancel = CancellationToken::new();
1835 let sup = TaskSupervisor::new(cancel.clone());
1836
1837 for name in [
1838 "loop-1", "loop-2", "loop-3", "loop-4", "loop-5", "loop-6", "loop-7",
1839 ] {
1840 let cancel_inner = cancel.clone();
1841 sup.spawn(TaskDescriptor {
1842 name,
1843 restart: RestartPolicy::RunOnce,
1844 factory: move || {
1845 let c = cancel_inner.clone();
1846 async move {
1847 c.cancelled().await;
1848 // Yield multiple times so the reap driver observes cancel first.
1849 for _ in 0..64 {
1850 tokio::task::yield_now().await;
1851 }
1852 }
1853 },
1854 });
1855 }
1856 assert_eq!(sup.active_count(), 7);
1857
1858 sup.shutdown_all(Duration::from_secs(2)).await;
1859
1860 assert_eq!(
1861 sup.active_count(),
1862 0,
1863 "all tasks must be reaped after shutdown (#3161)"
1864 );
1865 }
1866
1867 /// Regression test for #5926: `shutdown_all` must wait for a task that is
1868 /// still running when the caller-supplied `timeout` has not yet elapsed,
1869 /// rather than giving up on any independent, shorter deadline.
1870 ///
1871 /// `spawn_blocking` tasks don't observe the cancellation token — unlike
1872 /// `spawn`-based tasks, whose futures are dropped immediately by `do_spawn`'s
1873 /// `tokio::select!` against `cancel.cancelled()` — so a `spawn_blocking` task is
1874 /// the only way to have a task genuinely keep running past cancel and exercise
1875 /// the drain phase.
1876 #[tokio::test]
1877 async fn test_shutdown_all_reaps_blocking_task_finishing_after_cancel() {
1878 let (sup, _cancel) = make_supervisor();
1879
1880 let handle = sup.spawn_blocking(Arc::from("slow-blocking"), || {
1881 std::thread::sleep(Duration::from_millis(300));
1882 });
1883
1884 // Let the blocking task actually start before triggering shutdown.
1885 tokio::time::sleep(Duration::from_millis(20)).await;
1886
1887 let start = tokio::time::Instant::now();
1888 sup.shutdown_all(Duration::from_secs(2)).await;
1889 let elapsed = start.elapsed();
1890
1891 assert!(
1892 elapsed < Duration::from_secs(1),
1893 "shutdown_all must wake as soon as the blocking task actually finishes \
1894 (~300ms), not wait out the full 2s caller timeout; elapsed={elapsed:?}"
1895 );
1896
1897 // Successfully reaped tasks are removed from the registry entirely (see
1898 // `handle_completion`'s post-cancel short-circuit); only force-aborted tasks
1899 // remain with `TaskStatus::Aborted`. Absence here means the completion was
1900 // processed before shutdown_all's own timeout forced an abort.
1901 let snapshot = sup.snapshot();
1902 assert!(
1903 snapshot.iter().all(|t| t.name.as_ref() != "slow-blocking"),
1904 "task that completed cleanly after cancel must be reaped from the \
1905 registry, not left behind and force-aborted: {snapshot:?}"
1906 );
1907
1908 // `join` returns `Ok` only when the closure ran to completion and the
1909 // outer task wasn't aborted — the `BlockingHandle` equivalent of
1910 // `TaskStatus::Completed` vs `TaskStatus::Aborted`.
1911 handle
1912 .join()
1913 .await
1914 .expect("blocking task must complete, not be reported as aborted (#5926)");
1915 }
1916
1917 /// Regression test for #5926, critical topology (out-of-band cancel):
1918 /// production call sites (`src/runner.rs`, `src/serve/mod.rs`) always cancel
1919 /// the shared `CancellationToken` *before* calling `shutdown_all` — a
1920 /// separate shutdown bridge/signal handler owns the cancel, and `shutdown_all`
1921 /// is only invoked afterward as a waiter. A fix that only works when
1922 /// `shutdown_all` is the sole canceller would pass a naive test but stay
1923 /// broken in production. This test replicates the real topology explicitly:
1924 /// cancel first, let the reap driver start draining, *then* call
1925 /// `shutdown_all`, and confirm the caller's real timeout still governs the
1926 /// wait.
1927 #[tokio::test]
1928 async fn test_shutdown_all_reaps_task_after_out_of_band_cancel() {
1929 let cancel = CancellationToken::new();
1930 let sup = TaskSupervisor::new(cancel.clone());
1931
1932 let handle = sup.spawn_blocking(Arc::from("slow-blocking-oob"), || {
1933 std::thread::sleep(Duration::from_millis(300));
1934 });
1935 tokio::time::sleep(Duration::from_millis(20)).await;
1936
1937 // Out-of-band cancel — nothing to do with shutdown_all, mirrors a shutdown
1938 // bridge/signal handler cancelling the shared token directly.
1939 cancel.cancel();
1940 // Give the reap driver time to observe cancellation and enter Phase 2
1941 // before shutdown_all ever runs.
1942 tokio::time::sleep(Duration::from_millis(50)).await;
1943
1944 let start = tokio::time::Instant::now();
1945 sup.shutdown_all(Duration::from_secs(2)).await;
1946 let elapsed = start.elapsed();
1947
1948 assert!(
1949 elapsed < Duration::from_secs(1),
1950 "shutdown_all must wake once the blocking task actually finishes even \
1951 though the token was cancelled out-of-band before shutdown_all was \
1952 called; elapsed={elapsed:?}"
1953 );
1954
1955 handle.join().await.expect(
1956 "blocking task must complete, not be reported as aborted, even under \
1957 the out-of-band-cancel topology (#5926)",
1958 );
1959 }
1960
1961 /// Regression test for #5926 residual (critic finding S1): the reap driver's
1962 /// drain phase must never give up on a still-active, non-cooperative task on
1963 /// its own — only `shutdown_all`'s own `sleep(timeout)` may do that. An
1964 /// earlier version of the fix gave the drain phase its own short fallback
1965 /// deadline; if the gap between the out-of-band cancel and `shutdown_all`
1966 /// being invoked exceeded that fallback, the reap driver exited *before*
1967 /// `shutdown_all` ever ran, silently dropping the eventual completion and
1968 /// misreporting the task `Aborted`. This is realistic in production: e.g.
1969 /// `serve/mod.rs`'s `axum::serve(...).with_graceful_shutdown(...)` drains
1970 /// in-flight connections with no timeout of its own before `shutdown_all(30s)`
1971 /// is ever reached, and a long-running agent-turn request can hold that gap
1972 /// open well past any short fallback.
1973 ///
1974 /// Uses paused virtual time (`tokio::time::advance`) to jump the cancel-to-
1975 /// `shutdown_all` gap arbitrarily far forward without a real multi-second
1976 /// test. A `spawn_blocking` task runs on a real OS thread (unaffected by
1977 /// paused tokio time) and is held open via a channel until explicitly
1978 /// released, so its completion stays under exact manual control relative to
1979 /// the virtual-time advance and to when `shutdown_all` is actually called.
1980 #[tokio::test(start_paused = true)]
1981 async fn test_shutdown_all_reaps_task_after_long_pre_invocation_gap() {
1982 let cancel = CancellationToken::new();
1983 let sup = TaskSupervisor::new(cancel.clone());
1984
1985 let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
1986 let handle: BlockingHandle<()> = sup.spawn_blocking(Arc::from("held-task"), move || {
1987 let _ = release_rx.recv();
1988 });
1989
1990 // Out-of-band cancel — before shutdown_all is ever called.
1991 cancel.cancel();
1992 for _ in 0..8 {
1993 tokio::task::yield_now().await;
1994 }
1995
1996 // Advance virtual time by a large margin — deliberately larger than any
1997 // fallback window a naive fix might reintroduce — with `shutdown_all`
1998 // still not called and the task still held open. A drain phase with any
1999 // independent deadline of its own would have already exited by now.
2000 tokio::time::advance(Duration::from_mins(2)).await;
2001
2002 // *Now* the real caller shows up.
2003 let sup2 = sup.clone();
2004 let shutdown_jh =
2005 tokio::spawn(async move { sup2.shutdown_all(Duration::from_secs(30)).await });
2006 for _ in 0..8 {
2007 tokio::task::yield_now().await;
2008 }
2009
2010 // Let the real OS thread actually finish.
2011 let _ = release_tx.send(());
2012
2013 // Await directly (no wrapping virtual-time timeout): under `start_paused`,
2014 // an artificial timeout here would itself be subject to auto-advance and
2015 // cannot reliably distinguish outcomes. The real signal is the registry
2016 // state checked below, not how long this await took.
2017 shutdown_jh.await.expect("shutdown_all task must not panic");
2018
2019 // NOTE: `handle.join()` reflects `spawn_blocking`'s own internal oneshot
2020 // completion channel, which is fed directly by the real OS thread finishing
2021 // — entirely independent of the reap driver / registry path this test is
2022 // targeting. It succeeds either way and cannot distinguish the bug, so the
2023 // real assertion is the supervisor's registry state, not `handle.join()`.
2024 let snapshot = sup.snapshot();
2025 let entry = snapshot.iter().find(|t| t.name.as_ref() == "held-task");
2026 assert!(
2027 entry.is_none(),
2028 "a task finishing after shutdown_all was (belatedly) invoked must be \
2029 reaped from the registry (the post-cancel path removes completed \
2030 entries rather than marking them), not force-aborted, no matter how \
2031 long the cancel-to-shutdown_all gap was: {snapshot:?} (#5926 S1)"
2032 );
2033
2034 handle
2035 .join()
2036 .await
2037 .expect("blocking task must complete without panicking");
2038 }
2039
2040 #[tokio::test]
2041 async fn test_blocking_spawner_task_appears_in_snapshot() {
2042 // Verify that tasks spawned via BlockingSpawner appear in supervisor.snapshot().
2043 use crate::BlockingSpawner;
2044
2045 let cancel = CancellationToken::new();
2046 let sup = TaskSupervisor::new(cancel);
2047
2048 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
2049 let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
2050
2051 let handle = sup.spawn_blocking_named(
2052 Arc::from("chunk_file"),
2053 Box::new(move || {
2054 // Signal that the task has started.
2055 let _ = ready_tx.send(());
2056 // Block until test signals release.
2057 let _ = release_rx.blocking_recv();
2058 }),
2059 );
2060
2061 // Wait until the blocking task has actually started.
2062 ready_rx.await.expect("task should start");
2063
2064 let snapshot = sup.snapshot();
2065 assert!(
2066 snapshot.iter().any(|t| t.name.as_ref() == "chunk_file"),
2067 "chunk_file task must appear in supervisor snapshot"
2068 );
2069
2070 // Release the blocking task and await completion.
2071 let _ = release_tx.send(());
2072 handle.await.expect("task should complete");
2073 }
2074
2075 /// Verify that `measure_blocking` emits wall-time and CPU-time histograms.
2076 ///
2077 /// `measure_blocking` calls `metrics::histogram!` on the current thread.
2078 /// We test it directly using a `DebuggingRecorder` installed as the thread-local
2079 /// recorder via `metrics::with_local_recorder`.
2080 #[test]
2081 fn test_measure_blocking_emits_metrics() {
2082 use metrics_util::debugging::DebuggingRecorder;
2083
2084 let recorder = DebuggingRecorder::new();
2085 let snapshotter = recorder.snapshotter();
2086
2087 // Call measure_blocking inside the local recorder scope so histogram! calls
2088 // are captured. The closure runs synchronously on this thread.
2089 metrics::with_local_recorder(&recorder, || {
2090 measure_blocking("test_task", || std::hint::black_box(42_u64));
2091 });
2092
2093 let snapshot = snapshotter.snapshot();
2094 let metric_names: Vec<String> = snapshot
2095 .into_vec()
2096 .into_iter()
2097 .map(|(k, _, _, _)| k.key().name().to_owned())
2098 .collect();
2099
2100 assert!(
2101 metric_names.iter().any(|n| n == "zeph.task.wall_time_ms"),
2102 "expected zeph.task.wall_time_ms histogram; got: {metric_names:?}"
2103 );
2104 assert!(
2105 metric_names.iter().any(|n| n == "zeph.task.cpu_time_ms"),
2106 "expected zeph.task.cpu_time_ms histogram; got: {metric_names:?}"
2107 );
2108 }
2109
2110 /// Verify that `spawn_blocking` semaphore limits concurrent OS-thread tasks to 8.
2111 ///
2112 /// Spawns 16 tasks. Each holds a barrier until 8 are waiting; then releases in order.
2113 /// If more than 8 run concurrently the test would either deadlock (waiting for 9+ to reach
2114 /// the barrier) or the counter would exceed 8 — both are caught.
2115 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2116 async fn test_spawn_blocking_semaphore_cap() {
2117 let (sup, _cancel) = make_supervisor();
2118 let concurrent = Arc::new(AtomicU32::new(0));
2119 let max_concurrent = Arc::new(AtomicU32::new(0));
2120 let barrier = Arc::new(std::sync::Barrier::new(1)); // just a sync point
2121
2122 let mut handles = Vec::new();
2123 for i in 0u32..16 {
2124 let c = Arc::clone(&concurrent);
2125 let m = Arc::clone(&max_concurrent);
2126 let name: Arc<str> = Arc::from(format!("blocking-{i}").as_str());
2127 let h = sup.spawn_blocking(name, move || {
2128 let prev = c.fetch_add(1, Ordering::SeqCst);
2129 // Update observed maximum.
2130 let mut cur_max = m.load(Ordering::SeqCst);
2131 while prev + 1 > cur_max {
2132 match m.compare_exchange(cur_max, prev + 1, Ordering::SeqCst, Ordering::SeqCst)
2133 {
2134 Ok(_) => break,
2135 Err(x) => cur_max = x,
2136 }
2137 }
2138 // Simulate work.
2139 std::thread::sleep(std::time::Duration::from_millis(20));
2140 c.fetch_sub(1, Ordering::SeqCst);
2141 });
2142 handles.push(h);
2143 }
2144
2145 for h in handles {
2146 h.join().await.expect("blocking task should succeed");
2147 }
2148 drop(barrier);
2149
2150 let observed = max_concurrent.load(Ordering::SeqCst);
2151 assert!(
2152 observed <= 8,
2153 "observed {observed} concurrent blocking tasks; expected ≤ 8 (semaphore cap)"
2154 );
2155 }
2156}