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