Skip to main content

runledger_runtime/
supervisor.rs

1use std::borrow::Borrow;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::task::{Context, Poll};
7use std::time::Duration;
8
9use futures_util::stream::{FuturesUnordered, StreamExt};
10use tokio::runtime::Handle;
11use tokio::sync::watch;
12use tokio::task::{AbortHandle, JoinError, JoinHandle};
13use tokio::time::Instant;
14use tracing::{Instrument, debug, error, info_span, warn};
15
16use crate::catalog::JobCatalog;
17use crate::config::JobsConfig;
18use crate::reaper::run_reaper_loop;
19use crate::registry::JobRegistry;
20use crate::scheduler::run_scheduler_loop;
21use crate::worker::run_worker_loop;
22use crate::{Error, Result, RuntimeError, RuntimeLoopExit};
23
24const WORKER_TASK: &str = "worker";
25const SCHEDULER_TASK: &str = "scheduler";
26const REAPER_TASK: &str = "reaper";
27const MAX_ABORT_DRAIN_TIMEOUT: Duration = Duration::from_secs(1);
28
29/// Supervises the Runledger runtime loops spawned for a worker process.
30///
31/// A supervisor owns the worker, scheduler, and reaper task handles selected by
32/// [`SupervisorBuilder`]. Use [`Self::run_until_shutdown`] for a typical worker
33/// process that should exit on either an external shutdown signal or an internal
34/// runtime task failure.
35///
36/// Dropping a supervisor requests shutdown and detaches the task handles. Call
37/// [`Self::shutdown`] or [`Self::join`] when the owning process needs to observe
38/// panics or unexpected task exits.
39#[must_use]
40pub struct Supervisor {
41    shutdown_tx: watch::Sender<bool>,
42    shutdown_requested: Arc<AtomicBool>,
43    tasks: Vec<RuntimeTask>,
44}
45
46/// Builds a [`Supervisor`] with configurable runtime loops.
47///
48/// Worker, scheduler, and reaper loops are enabled by default. Call
49/// [`Self::with_registry`] or [`Self::with_catalog`] before [`Self::build`] when
50/// worker or reaper loops remain enabled.
51#[must_use]
52pub struct SupervisorBuilder<'a> {
53    pool: &'a runledger_postgres::DbPool,
54    runtime: Handle,
55    registry: Option<JobRegistry>,
56    registry_source: Option<RegistrySource>,
57    mixed_registry_sources: bool,
58    config: JobsConfig,
59    worker_enabled: bool,
60    scheduler_enabled: bool,
61    reaper_enabled: bool,
62}
63
64/// Cloneable handle for requesting supervisor shutdown from another task.
65#[derive(Clone)]
66pub struct SupervisorShutdown {
67    shutdown_tx: watch::Sender<bool>,
68    shutdown_requested: Arc<AtomicBool>,
69}
70
71struct RuntimeTask {
72    name: &'static str,
73    handle: JoinHandle<RuntimeTaskExit>,
74}
75
76#[derive(Clone, Copy, Debug, Eq, PartialEq)]
77enum RegistrySource {
78    Registry,
79    Catalog,
80}
81
82#[derive(Clone, Copy, Debug, Eq, PartialEq)]
83enum RuntimeTaskExit {
84    Completed,
85    InvalidConfig(crate::config::JobsConfigValidationError),
86    Shutdown,
87}
88
89#[derive(Clone, Copy, Debug, Eq, PartialEq)]
90enum DrainResult {
91    Drained,
92    TimedOut,
93}
94
95struct RuntimeTaskFuture {
96    name: &'static str,
97    future: Pin<Box<dyn Future<Output = RuntimeTaskExit> + Send>>,
98    started: bool,
99}
100
101type RuntimeTaskJoinResult = std::result::Result<RuntimeTaskExit, JoinError>;
102type JoinedRuntimeTask = (&'static str, RuntimeTaskJoinResult);
103
104impl Supervisor {
105    /// Returns a builder for a supervisor over a shared pool and runtime
106    /// configuration.
107    ///
108    /// This validates that the caller is inside the Tokio runtime that will own
109    /// spawned supervisor tasks.
110    pub fn builder(
111        pool: &runledger_postgres::DbPool,
112        config: JobsConfig,
113    ) -> std::result::Result<SupervisorBuilder<'_>, RuntimeError> {
114        let runtime =
115            Handle::try_current().map_err(|source| RuntimeError::MissingTokioRuntime { source })?;
116
117        Ok(SupervisorBuilder {
118            pool,
119            runtime,
120            registry: None,
121            registry_source: None,
122            mixed_registry_sources: false,
123            config,
124            worker_enabled: true,
125            scheduler_enabled: true,
126            reaper_enabled: true,
127        })
128    }
129
130    /// Returns a cloneable shutdown handle that can request shutdown without
131    /// owning the supervisor task joins.
132    #[must_use]
133    pub fn shutdown_handle(&self) -> SupervisorShutdown {
134        SupervisorShutdown {
135            shutdown_tx: self.shutdown_tx.clone(),
136            shutdown_requested: Arc::clone(&self.shutdown_requested),
137        }
138    }
139
140    /// Requests graceful shutdown of all supervised loops.
141    pub fn request_shutdown(&self) {
142        request_shutdown_signal(&self.shutdown_tx, self.shutdown_requested.as_ref());
143    }
144
145    /// Returns whether shutdown has been requested through this supervisor or a
146    /// clone of its shutdown handle.
147    #[must_use]
148    pub fn is_shutdown_requested(&self) -> bool {
149        self.shutdown_requested.load(Ordering::SeqCst)
150    }
151
152    /// Waits for all supervised loops to exit.
153    ///
154    /// With the default long-running loops, this method waits until shutdown is
155    /// requested through a [`SupervisorShutdown`] handle or until a task exits.
156    /// If a loop exits before shutdown was requested, the remaining loops are
157    /// asked to shut down and the first observed error is returned. Additional
158    /// task failures observed while draining are logged. This method does not
159    /// impose a deadline; use [`Self::shutdown_with_timeout`] when the caller
160    /// owns shutdown and needs a bounded wait.
161    pub async fn join(mut self) -> Result<()> {
162        let tasks = std::mem::take(&mut self.tasks);
163        self.join_tasks(tasks).await
164    }
165
166    /// Requests graceful shutdown and waits for all supervised loops to exit.
167    ///
168    /// If a loop exits before shutdown was requested, the remaining loops are
169    /// asked to shut down and the pre-existing task exit is reported, even when
170    /// that exit is only observed after shutdown begins. This method does not
171    /// impose a deadline. Use [`Self::shutdown_with_timeout`] when the owning
172    /// process needs a shutdown budget; externally timing out this consuming
173    /// future can detach still-running task handles.
174    pub async fn shutdown(mut self) -> Result<()> {
175        if let Some(error) = join_pre_shutdown_finished_tasks(&mut self.tasks).await {
176            self.request_shutdown();
177            let tasks = std::mem::take(&mut self.tasks);
178            drain_tasks(tasks).await;
179            return Err(Error::Runtime(error));
180        }
181
182        self.request_shutdown();
183        let tasks = std::mem::take(&mut self.tasks);
184        self.join_tasks(tasks).await
185    }
186
187    /// Waits until `shutdown` resolves or a supervised task fails, then exits.
188    ///
189    /// If `shutdown` resolves first, graceful shutdown is requested and the
190    /// supervisor waits up to `timeout` for all loops to exit. If a loop panics
191    /// or exits unexpectedly before `shutdown` resolves, shutdown is requested
192    /// for the remaining loops and the original task error is returned after
193    /// those loops drain or a timeout is reported. If shutdown is requested
194    /// through a [`SupervisorShutdown`] handle and every loop exits cleanly before
195    /// `shutdown` resolves, this returns successfully.
196    ///
197    /// This is the preferred method for worker binaries because it observes
198    /// internal task failures during normal operation while still applying a
199    /// bounded shutdown budget to cooperative process termination.
200    ///
201    /// If `timeout` is too large to represent as a runtime deadline, this returns
202    /// [`RuntimeError::ShutdownTimeoutTooLarge`] immediately. A zero timeout
203    /// requests shutdown, aborts tasks without waiting for cooperative exits, and
204    /// reports [`RuntimeError::ShutdownTimeout`].
205    ///
206    /// If the initial timeout validation fails before `shutdown` resolves, the
207    /// supervisor is still dropped, so shutdown is requested, but task handles
208    /// are not aborted or drained. If a deadline overflow is detected after
209    /// shutdown begins, remaining tasks are aborted and drained before returning.
210    pub async fn run_until_shutdown<F>(mut self, shutdown: F, timeout: Duration) -> Result<()>
211    where
212        F: Future<Output = ()>,
213    {
214        // Validate the shutdown budget before waiting on a signal, then
215        // recompute the deadline when shutdown actually begins.
216        let _ = shutdown_deadline(timeout)?;
217        let tasks = std::mem::take(&mut self.tasks);
218        if tasks.is_empty() {
219            shutdown.await;
220            self.request_shutdown();
221            return Ok(());
222        }
223
224        let mut abort_handles = Some(task_abort_handles(&tasks));
225        let mut joined = join_runtime_tasks(tasks);
226        let mut shutdown = std::pin::pin!(shutdown);
227
228        loop {
229            tokio::select! {
230                _ = shutdown.as_mut() => {
231                    self.request_shutdown();
232                    let abort_handles = abort_handles.take().expect("abort handles are consumed on return");
233                    // The budget was validated before waiting; recompute here so
234                    // the timeout starts when shutdown begins, and abort instead
235                    // of detaching if this late recompute somehow overflows.
236                    let deadline = match shutdown_deadline(timeout) {
237                        Ok(deadline) => deadline,
238                        Err(error) => {
239                            abort_and_drain_joined_tasks_or_log(
240                                &mut joined,
241                                abort_handles,
242                                abort_drain_timeout(timeout),
243                            )
244                            .await;
245                            return Err(error.into());
246                        }
247                    };
248                    return self
249                        .join_joined_tasks_with_timeout(
250                            &mut joined,
251                            abort_handles,
252                            timeout,
253                            deadline,
254                        )
255                        .await;
256                }
257                joined_result = joined.next() => {
258                    let Some((task, result)) = joined_result else {
259                        return Ok(());
260                    };
261                    let Some(error) = classify_task_result(task, result) else {
262                        continue;
263                    };
264
265                    self.request_shutdown();
266                    let abort_handles = abort_handles.take().expect("abort handles are consumed on return");
267                    // Start the drain budget when shutdown begins. If this
268                    // pathological recompute overflows, abort already-running
269                    // tasks rather than dropping their handles detached.
270                    let deadline = match shutdown_deadline(timeout) {
271                        Ok(deadline) => deadline,
272                        Err(error) => {
273                            abort_and_drain_joined_tasks_or_log(
274                                &mut joined,
275                                abort_handles,
276                                abort_drain_timeout(timeout),
277                            )
278                            .await;
279                            return Err(error.into());
280                        }
281                    };
282                    return drain_after_task_error_with_timeout(
283                        &mut joined,
284                        abort_handles,
285                        timeout,
286                        deadline,
287                        error,
288                    )
289                    .await;
290                }
291            }
292        }
293    }
294
295    /// Requests graceful shutdown and waits up to `timeout` for all supervised
296    /// loops to exit.
297    ///
298    /// If a loop had already exited before this method begins shutdown, that
299    /// failure is returned after the remaining loops have had the same shutdown
300    /// budget to exit cooperatively. If the timeout expires, remaining tasks are
301    /// aborted and drained with a bounded cleanup attempt before a timeout error
302    /// is returned. Abort cleanup can make total wall-clock time exceed `timeout`
303    /// by up to one second, or `timeout`, whichever is smaller. A zero timeout
304    /// requests shutdown, immediately aborts tasks that did not already finish,
305    /// and reports [`RuntimeError::ShutdownTimeout`].
306    ///
307    /// If `timeout` is too large to represent as a runtime deadline, this returns
308    /// [`RuntimeError::ShutdownTimeoutTooLarge`] immediately. The supervisor is
309    /// still dropped, so shutdown is requested, but task handles are not aborted
310    /// or drained.
311    pub async fn shutdown_with_timeout(mut self, timeout: Duration) -> Result<()> {
312        let deadline = shutdown_deadline(timeout)?;
313
314        if let Some(error) = join_pre_shutdown_finished_tasks(&mut self.tasks).await {
315            self.request_shutdown();
316            let tasks = std::mem::take(&mut self.tasks);
317            let abort_handles = task_abort_handles(&tasks);
318            let mut joined = join_runtime_tasks(tasks);
319
320            return drain_after_task_error_with_timeout(
321                &mut joined,
322                abort_handles,
323                timeout,
324                deadline,
325                error,
326            )
327            .await;
328        }
329
330        self.request_shutdown();
331        let tasks = std::mem::take(&mut self.tasks);
332        self.join_tasks_with_timeout(tasks, timeout, deadline).await
333    }
334
335    async fn join_tasks(&self, tasks: Vec<RuntimeTask>) -> Result<()> {
336        let mut joined = join_runtime_tasks(tasks);
337
338        while let Some((task, result)) = joined.next().await {
339            if let Some(error) = classify_task_result(task, result) {
340                self.request_shutdown();
341                drain_joined_tasks(&mut joined).await;
342                return Err(Error::Runtime(error));
343            }
344        }
345
346        Ok(())
347    }
348
349    async fn join_tasks_with_timeout(
350        &self,
351        tasks: Vec<RuntimeTask>,
352        timeout: Duration,
353        deadline: Instant,
354    ) -> Result<()> {
355        let abort_handles = task_abort_handles(&tasks);
356        let mut joined = join_runtime_tasks(tasks);
357
358        self.join_joined_tasks_with_timeout(&mut joined, abort_handles, timeout, deadline)
359            .await
360    }
361
362    async fn join_joined_tasks_with_timeout(
363        &self,
364        joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
365        abort_handles: Vec<AbortHandle>,
366        timeout: Duration,
367        deadline: Instant,
368    ) -> Result<()> {
369        loop {
370            match tokio::time::timeout_at(deadline, joined.next()).await {
371                Ok(Some((task, result))) => {
372                    if let Some(error) = classify_task_result(task, result) {
373                        self.request_shutdown();
374                        return drain_after_task_error_with_timeout(
375                            joined,
376                            abort_handles,
377                            timeout,
378                            deadline,
379                            error,
380                        )
381                        .await;
382                    }
383                }
384                Ok(None) => return Ok(()),
385                Err(_) => {
386                    abort_and_drain_joined_tasks_or_log(
387                        joined,
388                        abort_handles,
389                        abort_drain_timeout(timeout),
390                    )
391                    .await;
392                    return Err(Error::Runtime(RuntimeError::ShutdownTimeout { timeout }));
393                }
394            }
395        }
396    }
397
398    #[cfg(test)]
399    fn from_tasks_for_tests(tasks: Vec<RuntimeTask>) -> Self {
400        // These synthetic tasks do not receive this channel; tests using this
401        // helper either finish without shutdown or exercise timeout abort paths.
402        let (shutdown_tx, _) = watch::channel(false);
403        Self {
404            shutdown_tx,
405            shutdown_requested: Arc::new(AtomicBool::new(false)),
406            tasks,
407        }
408    }
409}
410
411impl Drop for Supervisor {
412    fn drop(&mut self) {
413        if !self.tasks.is_empty() {
414            warn!(
415                task_count = self.tasks.len(),
416                "dropping jobs runtime supervisor before joining tasks; tasks may continue detached after shutdown is requested and later panics will not be observed"
417            );
418        }
419        // Drop cannot await task handles, so this only nudges loops to exit.
420        self.request_shutdown();
421    }
422}
423
424impl<'a> SupervisorBuilder<'a> {
425    /// Registers the handlers used by worker execution and reaper terminal hooks.
426    ///
427    /// A registry is required when worker or reaper loops are enabled. Scheduler-only
428    /// supervisors can be built without one.
429    #[must_use = "builder methods return an updated builder value"]
430    pub fn with_registry(mut self, registry: JobRegistry) -> Self {
431        self.mixed_registry_sources |= self.registry_source == Some(RegistrySource::Catalog);
432        self.registry_source = Some(RegistrySource::Registry);
433        self.registry = Some(registry);
434        self
435    }
436
437    /// Registers handlers from a [`JobCatalog`].
438    ///
439    /// This does not sync database job definitions. Call
440    /// [`JobCatalog::sync_definitions`] before starting the supervisor or
441    /// creating schedules. Pass `&catalog` when the caller will continue using
442    /// the catalog for schedule, enqueue, or workflow helpers after building the
443    /// supervisor.
444    ///
445    /// # Registry Source
446    ///
447    /// Calling this and [`Self::with_registry`] on the same builder is rejected
448    /// by [`Self::build`]. Choose one registration source per builder.
449    #[must_use = "builder methods return an updated builder value"]
450    pub fn with_catalog(mut self, catalog: impl Borrow<JobCatalog>) -> Self {
451        self.mixed_registry_sources |= self.registry_source == Some(RegistrySource::Registry);
452        self.registry_source = Some(RegistrySource::Catalog);
453        self.registry = Some(catalog.borrow().to_registry());
454        self
455    }
456
457    /// Disables worker job claiming and execution for this supervisor.
458    #[must_use = "builder methods return an updated builder value"]
459    pub fn disable_worker(mut self) -> Self {
460        self.worker_enabled = false;
461        self
462    }
463
464    /// Disables cron schedule materialization for this supervisor.
465    #[must_use = "builder methods return an updated builder value"]
466    pub fn disable_scheduler(mut self) -> Self {
467        self.scheduler_enabled = false;
468        self
469    }
470
471    /// Disables expired-lease reaping for this supervisor.
472    #[must_use = "builder methods return an updated builder value"]
473    pub fn disable_reaper(mut self) -> Self {
474        self.reaper_enabled = false;
475        self
476    }
477
478    /// Starts the enabled runtime loops and returns the owning supervisor.
479    ///
480    /// Returns an error when worker or reaper loops are enabled without a job
481    /// registry.
482    pub fn build(self) -> std::result::Result<Supervisor, RuntimeError> {
483        let Self {
484            pool,
485            runtime,
486            registry,
487            registry_source: _,
488            mixed_registry_sources,
489            config,
490            worker_enabled,
491            scheduler_enabled,
492            reaper_enabled,
493        } = self;
494
495        config
496            .validate()
497            .map_err(|source| RuntimeError::InvalidJobsConfig { source })?;
498
499        if mixed_registry_sources {
500            return Err(RuntimeError::MixedRegistrySources);
501        }
502
503        let registry = match registry {
504            Some(registry) => registry,
505            None if worker_enabled || reaper_enabled => {
506                return Err(RuntimeError::MissingRegistry {
507                    worker_enabled,
508                    reaper_enabled,
509                });
510            }
511            None => JobRegistry::new(),
512        };
513
514        let (shutdown_tx, shutdown_rx) = watch::channel(false);
515        let shutdown_requested = Arc::new(AtomicBool::new(false));
516        let mut tasks = Vec::new();
517
518        if worker_enabled {
519            tasks.push(RuntimeTask::spawn_on(&runtime, WORKER_TASK, {
520                let pool = pool.clone();
521                let registry = registry.clone();
522                let config = config.clone();
523                let shutdown_rx = shutdown_rx.clone();
524                async move { run_worker_loop(pool, registry, config, shutdown_rx).await }
525            }));
526        }
527
528        if scheduler_enabled {
529            tasks.push(RuntimeTask::spawn_on(&runtime, SCHEDULER_TASK, {
530                let pool = pool.clone();
531                let config = config.clone();
532                let shutdown_rx = shutdown_rx.clone();
533                async move { run_scheduler_loop(pool, config, shutdown_rx).await }
534            }));
535        }
536
537        if reaper_enabled {
538            let pool = pool.clone();
539            let registry = registry.clone();
540            let config = config.clone();
541            let shutdown_rx = shutdown_rx.clone();
542            tasks.push(RuntimeTask::spawn_on(&runtime, REAPER_TASK, async move {
543                run_reaper_loop(pool, registry, config, shutdown_rx).await
544            }));
545        }
546
547        Ok(Supervisor {
548            shutdown_tx,
549            shutdown_requested,
550            tasks,
551        })
552    }
553}
554
555impl SupervisorShutdown {
556    /// Requests graceful shutdown of all loops watched by the supervisor.
557    pub fn request_shutdown(&self) {
558        request_shutdown_signal(&self.shutdown_tx, self.shutdown_requested.as_ref());
559    }
560
561    /// Returns whether shutdown has been requested.
562    #[must_use]
563    pub fn is_shutdown_requested(&self) -> bool {
564        self.shutdown_requested.load(Ordering::SeqCst)
565    }
566}
567
568impl RuntimeTask {
569    fn spawn_on<F>(runtime: &Handle, name: &'static str, future: F) -> Self
570    where
571        F: Future<Output = RuntimeLoopExit> + Send + 'static,
572    {
573        let span = info_span!("runledger_runtime_supervisor_task", task = name);
574        Self {
575            name,
576            handle: runtime.spawn(
577                RuntimeTaskFuture::new(name, async move { future.await.into() }).instrument(span),
578            ),
579        }
580    }
581
582    #[cfg(test)]
583    fn spawn<F>(name: &'static str, future: F) -> Self
584    where
585        F: Future<Output = RuntimeTaskExit> + Send + 'static,
586    {
587        Self {
588            name,
589            handle: tokio::spawn(RuntimeTaskFuture::new(name, future)),
590        }
591    }
592
593    async fn await_result(self) -> RuntimeTaskJoinResult {
594        self.handle.await
595    }
596}
597
598impl RuntimeTaskFuture {
599    fn new<F>(name: &'static str, future: F) -> Self
600    where
601        F: Future<Output = RuntimeTaskExit> + Send + 'static,
602    {
603        Self {
604            name,
605            future: Box::pin(future),
606            started: false,
607        }
608    }
609}
610
611impl Future for RuntimeTaskFuture {
612    type Output = RuntimeTaskExit;
613
614    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
615        let task = self.as_mut().get_mut();
616        if !task.started {
617            task.started = true;
618            debug!(task = task.name, "supervised runtime task started");
619        }
620
621        match task.future.as_mut().poll(cx) {
622            Poll::Pending => Poll::Pending,
623            Poll::Ready(exit) => {
624                debug!(task = task.name, ?exit, "supervised runtime task completed");
625                Poll::Ready(exit)
626            }
627        }
628    }
629}
630
631impl From<RuntimeLoopExit> for RuntimeTaskExit {
632    fn from(exit: RuntimeLoopExit) -> Self {
633        match exit {
634            RuntimeLoopExit::Shutdown => Self::Shutdown,
635            RuntimeLoopExit::InvalidConfig(source) => Self::InvalidConfig(source),
636            RuntimeLoopExit::Completed => Self::Completed,
637        }
638    }
639}
640
641fn request_shutdown_signal(shutdown_tx: &watch::Sender<bool>, shutdown_requested: &AtomicBool) {
642    if !shutdown_requested.swap(true, Ordering::SeqCst) {
643        // Mark the observable flag before notifying receivers so status readers
644        // agree that shutdown has begun as soon as request_shutdown returns.
645        // Safe to ignore: send failure means every shutdown receiver was dropped,
646        // so there are no supervised loops left to notify.
647        let _ = shutdown_tx.send(true);
648    }
649}
650
651fn take_finished_tasks(tasks: &mut Vec<RuntimeTask>) -> Vec<RuntimeTask> {
652    let mut finished = Vec::new();
653    let mut index = 0;
654    while index < tasks.len() {
655        if tasks[index].handle.is_finished() {
656            // Preserving order is not required for draining, and swap_remove
657            // avoids shifting still-running task handles.
658            finished.push(tasks.swap_remove(index));
659        } else {
660            index += 1;
661        }
662    }
663    finished
664}
665
666async fn join_pre_shutdown_finished_tasks(tasks: &mut Vec<RuntimeTask>) -> Option<RuntimeError> {
667    let finished = take_finished_tasks(tasks);
668    let mut first_error = None;
669
670    for task in finished {
671        let task_name = task.name;
672        let result = task.await_result().await;
673        let Some(error) = classify_task_result(task_name, result) else {
674            continue;
675        };
676
677        if first_error.is_none() {
678            first_error = Some(error);
679        } else {
680            log_drained_task_error(error);
681        }
682    }
683
684    first_error
685}
686
687fn join_runtime_tasks(
688    tasks: Vec<RuntimeTask>,
689) -> FuturesUnordered<impl Future<Output = JoinedRuntimeTask>> {
690    tasks
691        .into_iter()
692        .map(|task| async move {
693            let name = task.name;
694            (name, task.await_result().await)
695        })
696        .collect()
697}
698
699async fn drain_tasks(tasks: Vec<RuntimeTask>) {
700    let mut joined = join_runtime_tasks(tasks);
701    drain_joined_tasks(&mut joined).await;
702}
703
704fn task_abort_handles(tasks: &[RuntimeTask]) -> Vec<AbortHandle> {
705    tasks
706        .iter()
707        .map(|task| task.handle.abort_handle())
708        .collect()
709}
710
711fn shutdown_deadline(timeout: Duration) -> std::result::Result<Instant, RuntimeError> {
712    Instant::now()
713        .checked_add(timeout)
714        .ok_or(RuntimeError::ShutdownTimeoutTooLarge { timeout })
715}
716
717fn abort_drain_timeout(timeout: Duration) -> Duration {
718    timeout.min(MAX_ABORT_DRAIN_TIMEOUT)
719}
720
721async fn drain_joined_tasks(
722    joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
723) {
724    while let Some((task, result)) = joined.next().await {
725        if let Some(error) = classify_task_result(task, result) {
726            log_drained_task_error(error);
727        }
728    }
729}
730
731async fn drain_joined_tasks_until_deadline(
732    joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
733    deadline: Instant,
734) -> DrainResult {
735    loop {
736        match tokio::time::timeout_at(deadline, joined.next()).await {
737            Ok(Some((task, result))) => {
738                if let Some(error) = classify_task_result(task, result) {
739                    log_drained_task_error(error);
740                }
741            }
742            Ok(None) => return DrainResult::Drained,
743            Err(_) => return DrainResult::TimedOut,
744        }
745    }
746}
747
748async fn drain_after_task_error_with_timeout(
749    joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
750    abort_handles: Vec<AbortHandle>,
751    timeout: Duration,
752    deadline: Instant,
753    error: RuntimeError,
754) -> Result<()> {
755    if matches!(
756        drain_joined_tasks_until_deadline(joined, deadline).await,
757        DrainResult::Drained
758    ) {
759        return Err(error.into());
760    }
761
762    abort_and_drain_joined_tasks_or_log(joined, abort_handles, abort_drain_timeout(timeout)).await;
763    Err(RuntimeError::ShutdownTimeoutAfterTaskError {
764        timeout,
765        source: Box::new(error),
766    }
767    .into())
768}
769
770async fn abort_and_drain_joined_tasks_with_timeout(
771    joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
772    abort_handles: Vec<AbortHandle>,
773    timeout: Duration,
774) -> DrainResult {
775    for abort_handle in abort_handles {
776        abort_handle.abort();
777    }
778
779    match tokio::time::timeout(timeout, drain_aborted_joined_tasks(joined)).await {
780        Ok(()) => DrainResult::Drained,
781        Err(_) => DrainResult::TimedOut,
782    }
783}
784
785async fn abort_and_drain_joined_tasks_or_log(
786    joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
787    abort_handles: Vec<AbortHandle>,
788    timeout: Duration,
789) {
790    if matches!(
791        abort_and_drain_joined_tasks_with_timeout(joined, abort_handles, timeout).await,
792        DrainResult::TimedOut
793    ) {
794        log_abort_drain_timeout(timeout);
795    }
796}
797
798async fn drain_aborted_joined_tasks(
799    joined: &mut FuturesUnordered<impl Future<Output = JoinedRuntimeTask>>,
800) {
801    while let Some((task, result)) = joined.next().await {
802        match result {
803            Ok(_) => {}
804            Err(source) if source.is_cancelled() => {
805                // Cancellation is expected after this helper aborts the
806                // remaining tasks; the caller returns the timeout or earlier task
807                // failure that triggered the abort.
808            }
809            Err(source) => {
810                log_drained_task_error(RuntimeError::TaskJoin { task, source });
811            }
812        }
813    }
814}
815
816fn log_drained_task_error(error: RuntimeError) {
817    error!(
818        %error,
819        "supervised runtime task failed while draining after an earlier failure"
820    );
821}
822
823fn log_abort_drain_timeout(timeout: Duration) {
824    warn!(
825        ?timeout,
826        "timed out draining aborted supervisor tasks; later task failures may be unobserved"
827    );
828}
829
830fn classify_task_result(task: &'static str, result: RuntimeTaskJoinResult) -> Option<RuntimeError> {
831    match result {
832        Ok(RuntimeTaskExit::Shutdown) => {
833            debug!(task, "supervised runtime task joined after shutdown");
834            None
835        }
836        Ok(RuntimeTaskExit::Completed) => {
837            debug!(task, "supervised runtime task exited before shutdown");
838            Some(RuntimeError::TaskExitedUnexpectedly { task })
839        }
840        Ok(RuntimeTaskExit::InvalidConfig(source)) => {
841            debug!(
842                task,
843                "supervised runtime task rejected invalid config after build validation"
844            );
845            Some(RuntimeError::InvalidJobsConfig { source })
846        }
847        Err(source) => {
848            debug!(
849                task,
850                is_cancelled = source.is_cancelled(),
851                is_panic = source.is_panic(),
852                "supervised runtime task join failed"
853            );
854            Some(RuntimeError::TaskJoin { task, source })
855        }
856    }
857}
858
859#[cfg(test)]
860mod tests {
861    use std::sync::Arc;
862    use std::sync::atomic::Ordering;
863    use std::time::Duration;
864
865    use sqlx::postgres::PgPoolOptions;
866    use tokio::time::timeout;
867
868    use super::*;
869    use crate::Error;
870
871    const UNUSED_LAZY_POOL_URL: &str = "postgres://postgres:postgres@127.0.0.1:65535/runledger";
872
873    struct DropFlag(Arc<AtomicBool>);
874
875    impl Drop for DropFlag {
876        fn drop(&mut self) {
877            self.0.store(true, Ordering::SeqCst);
878        }
879    }
880
881    struct CompleteAfterPollSignal {
882        entered_tx: Option<std::sync::mpsc::Sender<()>>,
883        release_rx: std::sync::mpsc::Receiver<()>,
884        exit: RuntimeTaskExit,
885    }
886
887    impl Future for CompleteAfterPollSignal {
888        type Output = RuntimeTaskExit;
889
890        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
891            let task = self.as_mut().get_mut();
892            if let Some(entered_tx) = task.entered_tx.take() {
893                entered_tx
894                    .send(())
895                    .expect("completion poll entry signal should be received");
896            }
897            task.release_rx
898                .recv()
899                .expect("completion poll should be released");
900            Poll::Ready(task.exit)
901        }
902    }
903
904    fn lazy_pool() -> runledger_postgres::DbPool {
905        PgPoolOptions::new()
906            // The disable-only tests never acquire this pool; this URL is only
907            // a valid PgPool value for supervisor wiring assertions.
908            .connect_lazy(UNUSED_LAZY_POOL_URL)
909            .expect("construct lazy pool")
910    }
911
912    fn test_config() -> JobsConfig {
913        JobsConfig {
914            worker_id: "supervisor-test-worker".to_string(),
915            poll_interval: Duration::from_millis(25),
916            claim_batch_size: 4,
917            lease_ttl_seconds: 10,
918            max_global_concurrency: 4,
919            reaper_interval: Duration::from_millis(50),
920            schedule_poll_interval: Duration::from_millis(50),
921            reaper_retry_delay_ms: 1_000,
922        }
923    }
924
925    fn empty_builder(pool: &runledger_postgres::DbPool) -> SupervisorBuilder<'_> {
926        Supervisor::builder(pool, test_config()).expect("supervisor builder has runtime")
927    }
928
929    fn missing_registry_flags(builder: SupervisorBuilder<'_>) -> (bool, bool) {
930        match builder.build() {
931            Err(RuntimeError::MissingRegistry {
932                worker_enabled,
933                reaper_enabled,
934            }) => (worker_enabled, reaper_enabled),
935            Ok(_) => panic!("missing registry should be a build error"),
936            Err(other) => panic!("expected missing registry error, got {other:?}"),
937        }
938    }
939
940    fn test_task<F>(name: &'static str, future: F) -> RuntimeTask
941    where
942        F: Future<Output = ()> + Send + 'static,
943    {
944        test_task_with_exit(name, RuntimeTaskExit::Completed, future)
945    }
946
947    fn test_shutdown_task<F>(name: &'static str, future: F) -> RuntimeTask
948    where
949        F: Future<Output = ()> + Send + 'static,
950    {
951        test_task_with_exit(name, RuntimeTaskExit::Shutdown, future)
952    }
953
954    fn test_task_with_exit<F>(name: &'static str, exit: RuntimeTaskExit, future: F) -> RuntimeTask
955    where
956        F: Future<Output = ()> + Send + 'static,
957    {
958        RuntimeTask::spawn(name, async move {
959            future.await;
960            exit
961        })
962    }
963
964    fn supervisor_from_shutdown_channel(
965        shutdown_tx: watch::Sender<bool>,
966        shutdown_requested: Arc<AtomicBool>,
967        tasks: Vec<RuntimeTask>,
968    ) -> Supervisor {
969        Supervisor {
970            shutdown_tx,
971            shutdown_requested,
972            tasks,
973        }
974    }
975
976    fn task_names(supervisor: &Supervisor) -> Vec<&'static str> {
977        supervisor.tasks.iter().map(|task| task.name).collect()
978    }
979
980    async fn abort_supervisor_tasks(mut supervisor: Supervisor) {
981        let tasks = std::mem::take(&mut supervisor.tasks);
982        for task in tasks {
983            task.handle.abort();
984            let _ = task.handle.await;
985        }
986    }
987
988    #[tokio::test]
989    async fn builder_defaults_enable_all_loops() {
990        let pool = lazy_pool();
991        let builder = empty_builder(&pool);
992
993        assert!(builder.worker_enabled);
994        assert!(builder.scheduler_enabled);
995        assert!(builder.reaper_enabled);
996        assert!(builder.registry.is_none());
997        assert_eq!(builder.registry_source, None);
998        assert!(!builder.mixed_registry_sources);
999    }
1000
1001    #[tokio::test]
1002    async fn builder_accepts_registry_for_worker_and_reaper_loops() {
1003        let pool = lazy_pool();
1004        let builder = empty_builder(&pool).with_registry(JobRegistry::new());
1005
1006        assert!(builder.registry.is_some());
1007        assert_eq!(builder.registry_source, Some(RegistrySource::Registry));
1008        assert!(!builder.mixed_registry_sources);
1009    }
1010
1011    #[tokio::test]
1012    async fn builder_rejects_mixed_registry_sources() {
1013        let pool = lazy_pool();
1014        let registry_then_catalog = empty_builder(&pool)
1015            .with_registry(JobRegistry::new())
1016            .with_catalog(JobCatalog::new())
1017            .disable_worker()
1018            .disable_reaper()
1019            .build();
1020        let Err(registry_then_catalog) = registry_then_catalog else {
1021            panic!("mixed registry sources should be rejected");
1022        };
1023        assert!(matches!(
1024            registry_then_catalog,
1025            RuntimeError::MixedRegistrySources
1026        ));
1027
1028        let catalog_then_registry = empty_builder(&pool)
1029            .with_catalog(JobCatalog::new())
1030            .with_registry(JobRegistry::new())
1031            .disable_worker()
1032            .disable_reaper()
1033            .build();
1034        let Err(catalog_then_registry) = catalog_then_registry else {
1035            panic!("mixed registry sources should be rejected");
1036        };
1037        assert!(matches!(
1038            catalog_then_registry,
1039            RuntimeError::MixedRegistrySources
1040        ));
1041    }
1042
1043    #[tokio::test]
1044    async fn builder_requires_registry_when_worker_or_reaper_is_enabled() {
1045        let pool = lazy_pool();
1046
1047        assert_eq!(missing_registry_flags(empty_builder(&pool)), (true, true));
1048        assert_eq!(
1049            missing_registry_flags(empty_builder(&pool).disable_scheduler().disable_reaper()),
1050            (true, false)
1051        );
1052        assert_eq!(
1053            missing_registry_flags(empty_builder(&pool).disable_worker().disable_scheduler()),
1054            (false, true)
1055        );
1056    }
1057
1058    #[tokio::test]
1059    async fn builder_rejects_invalid_direct_config_values_before_spawning_loops() {
1060        let cases = [
1061            {
1062                let mut config = test_config();
1063                config.max_global_concurrency = 0;
1064                (
1065                    config,
1066                    crate::config::JobsConfigValidationError::InvalidMaxGlobalConcurrency,
1067                )
1068            },
1069            {
1070                let mut config = test_config();
1071                config.claim_batch_size = 0;
1072                (
1073                    config,
1074                    crate::config::JobsConfigValidationError::InvalidClaimBatchSize { actual: 0 },
1075                )
1076            },
1077            {
1078                let mut config = test_config();
1079                config.lease_ttl_seconds = 0;
1080                (
1081                    config,
1082                    crate::config::JobsConfigValidationError::InvalidLeaseTtlSeconds { actual: 0 },
1083                )
1084            },
1085        ];
1086
1087        for (config, expected) in cases {
1088            let pool = lazy_pool();
1089            let result = Supervisor::builder(&pool, config)
1090                .expect("supervisor builder has runtime")
1091                .disable_worker()
1092                .disable_scheduler()
1093                .disable_reaper()
1094                .build();
1095            let Err(error) = result else {
1096                panic!("invalid direct config should be rejected");
1097            };
1098
1099            match error {
1100                RuntimeError::InvalidJobsConfig { source } => {
1101                    assert_eq!(source, expected);
1102                }
1103                other => panic!("expected invalid jobs config error, got {other:?}"),
1104            }
1105        }
1106    }
1107
1108    #[test]
1109    fn builder_requires_tokio_runtime_before_cloning_pool() {
1110        let runtime = tokio::runtime::Runtime::new().expect("construct Tokio runtime");
1111        let pool = runtime.block_on(async { lazy_pool() });
1112        let error = match Supervisor::builder(&pool, test_config()) {
1113            Err(error) => error,
1114            Ok(builder) => {
1115                drop(builder);
1116                runtime.block_on(async {
1117                    pool.close().await;
1118                });
1119                std::mem::forget(pool);
1120                panic!("missing Tokio runtime should be a builder error");
1121            }
1122        };
1123
1124        // The builder was intentionally called outside a runtime to exercise
1125        // the pre-clone runtime check. Close and drop the pool inside the
1126        // temporary runtime so sqlx's own drop precondition does not contaminate
1127        // this assertion.
1128        runtime.block_on(async {
1129            pool.close().await;
1130        });
1131        std::mem::forget(pool);
1132        match error {
1133            RuntimeError::MissingTokioRuntime { .. } => {}
1134            other => panic!("expected missing Tokio runtime error, got {other:?}"),
1135        }
1136    }
1137
1138    #[tokio::test]
1139    async fn builder_can_disable_each_loop() {
1140        let pool = lazy_pool();
1141        let builder = empty_builder(&pool)
1142            .disable_worker()
1143            .disable_scheduler()
1144            .disable_reaper();
1145
1146        assert!(!builder.worker_enabled);
1147        assert!(!builder.scheduler_enabled);
1148        assert!(!builder.reaper_enabled);
1149    }
1150
1151    #[tokio::test]
1152    async fn builder_spawns_only_enabled_tasks() {
1153        let pool = lazy_pool();
1154
1155        let all_disabled = empty_builder(&pool)
1156            .disable_worker()
1157            .disable_scheduler()
1158            .disable_reaper()
1159            .build()
1160            .expect("all-disabled supervisor should build");
1161        assert_eq!(task_names(&all_disabled), Vec::<&'static str>::new());
1162        abort_supervisor_tasks(all_disabled).await;
1163
1164        let scheduler_only = empty_builder(&pool)
1165            .disable_worker()
1166            .disable_reaper()
1167            .build()
1168            .expect("scheduler-only supervisor should not require registry");
1169        assert_eq!(task_names(&scheduler_only), vec![SCHEDULER_TASK]);
1170        abort_supervisor_tasks(scheduler_only).await;
1171
1172        let worker_only = empty_builder(&pool)
1173            .with_registry(JobRegistry::new())
1174            .disable_scheduler()
1175            .disable_reaper()
1176            .build()
1177            .expect("worker-only supervisor should build with registry");
1178        assert_eq!(task_names(&worker_only), vec![WORKER_TASK]);
1179        abort_supervisor_tasks(worker_only).await;
1180
1181        let reaper_only = empty_builder(&pool)
1182            .with_registry(JobRegistry::new())
1183            .disable_worker()
1184            .disable_scheduler()
1185            .build()
1186            .expect("reaper-only supervisor should build with registry");
1187        assert_eq!(task_names(&reaper_only), vec![REAPER_TASK]);
1188        abort_supervisor_tasks(reaper_only).await;
1189
1190        let all_enabled = empty_builder(&pool)
1191            .with_registry(JobRegistry::new())
1192            .build()
1193            .expect("all-enabled supervisor should build with registry");
1194        assert_eq!(
1195            task_names(&all_enabled),
1196            vec![WORKER_TASK, SCHEDULER_TASK, REAPER_TASK]
1197        );
1198        abort_supervisor_tasks(all_enabled).await;
1199    }
1200
1201    #[tokio::test]
1202    async fn all_disabled_supervisor_join_and_shutdown_succeed() {
1203        Supervisor::builder(&lazy_pool(), test_config())
1204            .expect("supervisor builder has runtime")
1205            .disable_worker()
1206            .disable_scheduler()
1207            .disable_reaper()
1208            .build()
1209            .expect("all-disabled supervisor should build")
1210            .join()
1211            .await
1212            .expect("all-disabled supervisor should join");
1213
1214        Supervisor::builder(&lazy_pool(), test_config())
1215            .expect("supervisor builder has runtime")
1216            .disable_worker()
1217            .disable_scheduler()
1218            .disable_reaper()
1219            .build()
1220            .expect("all-disabled supervisor should build")
1221            .shutdown()
1222            .await
1223            .expect("all-disabled supervisor should shut down");
1224    }
1225
1226    #[tokio::test]
1227    async fn shutdown_handle_can_request_shutdown_before_join() {
1228        let supervisor = Supervisor::builder(&lazy_pool(), test_config())
1229            .expect("supervisor builder has runtime")
1230            .disable_worker()
1231            .disable_scheduler()
1232            .disable_reaper()
1233            .build()
1234            .expect("all-disabled supervisor should build");
1235        let shutdown = supervisor.shutdown_handle();
1236        let cloned_shutdown = shutdown.clone();
1237
1238        cloned_shutdown.request_shutdown();
1239
1240        assert!(shutdown.is_shutdown_requested());
1241        assert!(supervisor.is_shutdown_requested());
1242        supervisor
1243            .join()
1244            .await
1245            .expect("supervisor should join after shutdown handle request");
1246    }
1247
1248    #[tokio::test]
1249    async fn shutdown_after_shutdown_handle_request_allows_clean_task_exit() {
1250        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1251        let shutdown_requested = Arc::new(AtomicBool::new(false));
1252        let supervisor = supervisor_from_shutdown_channel(
1253            shutdown_tx,
1254            Arc::clone(&shutdown_requested),
1255            vec![test_shutdown_task("cooperative-loop", async move {
1256                while !*shutdown_rx.borrow() {
1257                    if shutdown_rx.changed().await.is_err() {
1258                        break;
1259                    }
1260                }
1261            })],
1262        );
1263        let shutdown = supervisor.shutdown_handle();
1264
1265        shutdown.request_shutdown();
1266
1267        supervisor
1268            .shutdown()
1269            .await
1270            .expect("clean exit after requested shutdown should succeed");
1271    }
1272
1273    #[tokio::test]
1274    async fn run_until_shutdown_requests_shutdown_when_signal_resolves() {
1275        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1276        let shutdown_requested = Arc::new(AtomicBool::new(false));
1277        let (signal_tx, signal_rx) = tokio::sync::oneshot::channel();
1278        let supervisor = supervisor_from_shutdown_channel(
1279            shutdown_tx,
1280            Arc::clone(&shutdown_requested),
1281            vec![test_shutdown_task(
1282                "run-until-cooperative-loop",
1283                async move {
1284                    while !*shutdown_rx.borrow() {
1285                        if shutdown_rx.changed().await.is_err() {
1286                            break;
1287                        }
1288                    }
1289                },
1290            )],
1291        );
1292
1293        signal_tx.send(()).expect("signal receiver should be alive");
1294
1295        supervisor
1296            .run_until_shutdown(
1297                async move {
1298                    signal_rx.await.expect("shutdown signal should be sent");
1299                },
1300                Duration::from_secs(1),
1301            )
1302            .await
1303            .expect("resolved shutdown signal should shut down cleanly");
1304        assert!(shutdown_requested.load(Ordering::SeqCst));
1305    }
1306
1307    #[tokio::test]
1308    async fn run_until_shutdown_with_no_tasks_waits_for_signal() {
1309        let supervisor = Supervisor::from_tasks_for_tests(Vec::new());
1310        let (signal_tx, signal_rx) = tokio::sync::oneshot::channel();
1311        let mut run = tokio::spawn(supervisor.run_until_shutdown(
1312            async move {
1313                signal_rx.await.expect("shutdown signal should be sent");
1314            },
1315            Duration::from_secs(1),
1316        ));
1317
1318        assert!(
1319            timeout(Duration::from_millis(50), &mut run).await.is_err(),
1320            "all-disabled supervisor should wait for the shutdown signal"
1321        );
1322
1323        signal_tx.send(()).expect("signal receiver should be alive");
1324        run.await
1325            .expect("run-until-shutdown task should join")
1326            .expect("all-disabled supervisor should complete after signal");
1327    }
1328
1329    #[tokio::test]
1330    async fn run_until_shutdown_reports_task_exit_before_signal() {
1331        let supervisor =
1332            Supervisor::from_tasks_for_tests(vec![test_task("run-until-early-loop", async {})]);
1333
1334        let error = timeout(
1335            Duration::from_secs(1),
1336            supervisor.run_until_shutdown(std::future::pending::<()>(), Duration::from_secs(1)),
1337        )
1338        .await
1339        .expect("task exit should be reported before external signal")
1340        .expect_err("early task exit should fail run-until shutdown");
1341
1342        match error {
1343            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1344                assert_eq!(task, "run-until-early-loop");
1345            }
1346            other => panic!("expected unexpected task exit, got {other:?}"),
1347        }
1348    }
1349
1350    #[tokio::test]
1351    async fn run_until_shutdown_times_out_and_aborts_after_signal() {
1352        let dropped = Arc::new(AtomicBool::new(false));
1353        let drop_flag = DropFlag(Arc::clone(&dropped));
1354        let supervisor =
1355            Supervisor::from_tasks_for_tests(vec![test_task("run-until-stubborn-loop", async {
1356                let _drop_flag = drop_flag;
1357                std::future::pending::<()>().await;
1358            })]);
1359
1360        let error = supervisor
1361            .run_until_shutdown(async {}, Duration::from_millis(50))
1362            .await
1363            .expect_err("stubborn task should time out after shutdown signal");
1364
1365        match error {
1366            Error::Runtime(RuntimeError::ShutdownTimeout { timeout }) => {
1367                assert_eq!(timeout, Duration::from_millis(50));
1368            }
1369            other => panic!("expected shutdown timeout error, got {other:?}"),
1370        }
1371        assert!(dropped.load(Ordering::SeqCst));
1372    }
1373
1374    #[tokio::test]
1375    async fn run_until_shutdown_reports_task_exit_after_signal_before_deadline() {
1376        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1377        let shutdown_requested = Arc::new(AtomicBool::new(false));
1378        let supervisor = supervisor_from_shutdown_channel(
1379            shutdown_tx,
1380            Arc::clone(&shutdown_requested),
1381            vec![test_task("run-until-bad-shutdown-loop", async move {
1382                while !*shutdown_rx.borrow() {
1383                    if shutdown_rx.changed().await.is_err() {
1384                        break;
1385                    }
1386                }
1387            })],
1388        );
1389
1390        let error = supervisor
1391            .run_until_shutdown(async {}, Duration::from_secs(1))
1392            .await
1393            .expect_err("task completion after signal should still be reported");
1394
1395        match error {
1396            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1397                assert_eq!(task, "run-until-bad-shutdown-loop");
1398            }
1399            other => panic!("expected unexpected task exit, got {other:?}"),
1400        }
1401    }
1402
1403    #[tokio::test]
1404    async fn dropping_supervisor_requests_shutdown_signal() {
1405        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1406        let mut observed_shutdown = shutdown_rx.clone();
1407        let shutdown_requested = Arc::new(AtomicBool::new(false));
1408        let supervisor = supervisor_from_shutdown_channel(
1409            shutdown_tx,
1410            Arc::clone(&shutdown_requested),
1411            vec![test_shutdown_task("drop-shutdown-loop", async move {
1412                while !*shutdown_rx.borrow() {
1413                    if shutdown_rx.changed().await.is_err() {
1414                        break;
1415                    }
1416                }
1417            })],
1418        );
1419
1420        drop(supervisor);
1421
1422        timeout(Duration::from_secs(1), observed_shutdown.changed())
1423            .await
1424            .expect("drop should promptly notify shutdown")
1425            .expect("shutdown sender should notify before closing");
1426        assert!(*observed_shutdown.borrow());
1427        assert!(shutdown_requested.load(Ordering::SeqCst));
1428    }
1429
1430    #[tokio::test]
1431    async fn join_reports_task_that_exited_before_late_shutdown_request() {
1432        let (shutdown_tx, _) = watch::channel(false);
1433        let shutdown_requested = Arc::new(AtomicBool::new(false));
1434        let supervisor = supervisor_from_shutdown_channel(
1435            shutdown_tx,
1436            Arc::clone(&shutdown_requested),
1437            vec![test_task("early-before-late-signal", async {})],
1438        );
1439
1440        while !supervisor.tasks[0].handle.is_finished() {
1441            tokio::task::yield_now().await;
1442        }
1443
1444        let shutdown = supervisor.shutdown_handle();
1445        shutdown.request_shutdown();
1446
1447        let error = supervisor
1448            .join()
1449            .await
1450            .expect_err("task exit before shutdown request should still be reported");
1451
1452        match error {
1453            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1454                assert_eq!(task, "early-before-late-signal");
1455            }
1456            other => panic!("expected unexpected task exit, got {other:?}"),
1457        }
1458    }
1459
1460    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1461    async fn join_reports_task_exit_when_shutdown_races_completion_poll() {
1462        let (entered_tx, entered_rx) = std::sync::mpsc::channel();
1463        let (release_tx, release_rx) = std::sync::mpsc::channel();
1464        let (shutdown_tx, _) = watch::channel(false);
1465        let shutdown_requested = Arc::new(AtomicBool::new(false));
1466        let supervisor = supervisor_from_shutdown_channel(
1467            shutdown_tx,
1468            Arc::clone(&shutdown_requested),
1469            vec![RuntimeTask::spawn(
1470                "race-completion",
1471                CompleteAfterPollSignal {
1472                    entered_tx: Some(entered_tx),
1473                    release_rx,
1474                    exit: RuntimeTaskExit::Completed,
1475                },
1476            )],
1477        );
1478        entered_rx
1479            .recv_timeout(Duration::from_secs(1))
1480            .expect("task should enter its completion poll");
1481        let shutdown = supervisor.shutdown_handle();
1482
1483        shutdown.request_shutdown();
1484        release_tx
1485            .send(())
1486            .expect("completion poll release should be received");
1487
1488        let error = supervisor
1489            .join()
1490            .await
1491            .expect_err("task exit that began before shutdown should be reported");
1492
1493        match error {
1494            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1495                assert_eq!(task, "race-completion");
1496            }
1497            other => panic!("expected unexpected task exit, got {other:?}"),
1498        }
1499    }
1500
1501    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1502    async fn join_allows_shutdown_exit_when_shutdown_races_completion_poll() {
1503        let (entered_tx, entered_rx) = std::sync::mpsc::channel();
1504        let (release_tx, release_rx) = std::sync::mpsc::channel();
1505        let (shutdown_tx, _) = watch::channel(false);
1506        let shutdown_requested = Arc::new(AtomicBool::new(false));
1507        let supervisor = supervisor_from_shutdown_channel(
1508            shutdown_tx,
1509            Arc::clone(&shutdown_requested),
1510            vec![RuntimeTask::spawn(
1511                "shutdown-race-completion",
1512                CompleteAfterPollSignal {
1513                    entered_tx: Some(entered_tx),
1514                    release_rx,
1515                    exit: RuntimeTaskExit::Shutdown,
1516                },
1517            )],
1518        );
1519        entered_rx
1520            .recv_timeout(Duration::from_secs(1))
1521            .expect("task should enter its completion poll");
1522        let shutdown = supervisor.shutdown_handle();
1523
1524        shutdown.request_shutdown();
1525        release_tx
1526            .send(())
1527            .expect("completion poll release should be received");
1528
1529        supervisor
1530            .join()
1531            .await
1532            .expect("task that reports shutdown should join cleanly");
1533    }
1534
1535    #[tokio::test]
1536    async fn panic_after_shutdown_request_is_reported() {
1537        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1538        let shutdown_requested = Arc::new(AtomicBool::new(false));
1539        let supervisor = supervisor_from_shutdown_channel(
1540            shutdown_tx,
1541            Arc::clone(&shutdown_requested),
1542            vec![test_shutdown_task("panic-after-shutdown", async move {
1543                while !*shutdown_rx.borrow() {
1544                    if shutdown_rx.changed().await.is_err() {
1545                        return;
1546                    }
1547                }
1548                panic!("forced post-shutdown panic");
1549            })],
1550        );
1551        let shutdown = supervisor.shutdown_handle();
1552
1553        shutdown.request_shutdown();
1554
1555        let error = supervisor
1556            .shutdown()
1557            .await
1558            .expect_err("panic after requested shutdown should fail");
1559
1560        match error {
1561            Error::Runtime(RuntimeError::TaskJoin { task, source }) => {
1562                assert_eq!(task, "panic-after-shutdown");
1563                assert!(source.is_panic());
1564            }
1565            other => panic!("expected task join error, got {other:?}"),
1566        }
1567    }
1568
1569    #[tokio::test]
1570    async fn early_normal_task_exit_is_unexpected() {
1571        let supervisor = Supervisor::from_tasks_for_tests(vec![test_task("test-loop", async {})]);
1572
1573        let error = supervisor
1574            .join()
1575            .await
1576            .expect_err("early normal exit should fail");
1577
1578        match error {
1579            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1580                assert_eq!(task, "test-loop");
1581            }
1582            other => panic!("expected unexpected task exit, got {other:?}"),
1583        }
1584    }
1585
1586    #[tokio::test]
1587    async fn invalid_config_task_exit_preserves_validation_source() {
1588        let expected =
1589            crate::config::JobsConfigValidationError::InvalidClaimBatchSize { actual: 0 };
1590        let supervisor = Supervisor::from_tasks_for_tests(vec![test_task_with_exit(
1591            "invalid-config-loop",
1592            RuntimeTaskExit::InvalidConfig(expected),
1593            async {},
1594        )]);
1595
1596        let error = supervisor
1597            .join()
1598            .await
1599            .expect_err("invalid config task exit should fail");
1600
1601        match error {
1602            Error::Runtime(RuntimeError::InvalidJobsConfig { source }) => {
1603                assert_eq!(source, expected);
1604            }
1605            other => panic!("expected invalid jobs config error, got {other:?}"),
1606        }
1607    }
1608
1609    #[tokio::test]
1610    async fn shutdown_reports_task_that_exited_before_shutdown_request() {
1611        let supervisor = Supervisor::from_tasks_for_tests(vec![test_task("early-loop", async {})]);
1612
1613        while !supervisor.tasks[0].handle.is_finished() {
1614            tokio::task::yield_now().await;
1615        }
1616
1617        let error = supervisor
1618            .shutdown()
1619            .await
1620            .expect_err("pre-shutdown task exit should fail");
1621
1622        match error {
1623            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1624                assert_eq!(task, "early-loop");
1625            }
1626            other => panic!("expected unexpected task exit, got {other:?}"),
1627        }
1628    }
1629
1630    #[tokio::test]
1631    async fn pre_shutdown_sweep_consumes_all_already_finished_tasks() {
1632        let mut tasks = vec![
1633            test_task("finished-a", async {}),
1634            test_task("pending", async {
1635                std::future::pending::<()>().await;
1636            }),
1637            test_task("finished-b", async {}),
1638        ];
1639
1640        while tasks
1641            .iter()
1642            .filter(|task| task.name != "pending")
1643            .any(|task| !task.handle.is_finished())
1644        {
1645            tokio::task::yield_now().await;
1646        }
1647
1648        let error = join_pre_shutdown_finished_tasks(&mut tasks)
1649            .await
1650            .expect("finished tasks should produce a pre-shutdown error");
1651
1652        match error {
1653            RuntimeError::TaskExitedUnexpectedly { task } => {
1654                assert!(
1655                    matches!(task, "finished-a" | "finished-b"),
1656                    "unexpected first finished task: {task}"
1657                );
1658            }
1659            other => panic!("expected unexpected task exit, got {other:?}"),
1660        }
1661        assert_eq!(tasks.len(), 1);
1662        assert_eq!(tasks[0].name, "pending");
1663
1664        let pending = tasks.pop().expect("pending task remains");
1665        pending.handle.abort();
1666        // The task was deliberately aborted; this test only needs the join
1667        // handle drained before returning.
1668        let _ = pending.handle.await;
1669    }
1670
1671    #[tokio::test]
1672    async fn pre_shutdown_sweep_allows_explicit_shutdown_exit() {
1673        let mut tasks = vec![test_shutdown_task("finished-after-signal", async {})];
1674        while !tasks[0].handle.is_finished() {
1675            tokio::task::yield_now().await;
1676        }
1677
1678        let error = join_pre_shutdown_finished_tasks(&mut tasks).await;
1679
1680        assert!(error.is_none());
1681        assert!(tasks.is_empty());
1682    }
1683
1684    #[tokio::test]
1685    async fn shutdown_with_timeout_aborts_and_drains_stubborn_task() {
1686        let dropped = Arc::new(AtomicBool::new(false));
1687        let drop_flag = DropFlag(Arc::clone(&dropped));
1688        let supervisor =
1689            Supervisor::from_tasks_for_tests(vec![test_task("stubborn-loop", async move {
1690                let _drop_flag = drop_flag;
1691                std::future::pending::<()>().await;
1692            })]);
1693
1694        let error = supervisor
1695            .shutdown_with_timeout(Duration::from_millis(50))
1696            .await
1697            .expect_err("stubborn task should time out shutdown");
1698
1699        match error {
1700            Error::Runtime(RuntimeError::ShutdownTimeout { timeout }) => {
1701                assert_eq!(timeout, Duration::from_millis(50));
1702            }
1703            other => panic!("expected shutdown timeout error, got {other:?}"),
1704        }
1705        assert!(dropped.load(Ordering::SeqCst));
1706    }
1707
1708    #[tokio::test]
1709    async fn shutdown_with_timeout_rejects_unrepresentable_deadline() {
1710        let error = Supervisor::from_tasks_for_tests(Vec::new())
1711            .shutdown_with_timeout(Duration::MAX)
1712            .await
1713            .expect_err("unrepresentable timeout should fail instead of panicking");
1714
1715        match error {
1716            Error::Runtime(RuntimeError::ShutdownTimeoutTooLarge { timeout }) => {
1717                assert_eq!(timeout, Duration::MAX);
1718            }
1719            other => panic!("expected oversized timeout error, got {other:?}"),
1720        }
1721    }
1722
1723    #[tokio::test]
1724    async fn shutdown_with_zero_timeout_aborts_immediately() {
1725        let supervisor =
1726            Supervisor::from_tasks_for_tests(vec![test_task("zero-timeout-pending-loop", async {
1727                std::future::pending::<()>().await;
1728            })]);
1729
1730        let error = supervisor
1731            .shutdown_with_timeout(Duration::ZERO)
1732            .await
1733            .expect_err("zero timeout should report an immediate shutdown timeout");
1734
1735        match error {
1736            Error::Runtime(RuntimeError::ShutdownTimeout { timeout }) => {
1737                assert_eq!(timeout, Duration::ZERO);
1738            }
1739            other => panic!("expected shutdown timeout error, got {other:?}"),
1740        }
1741    }
1742
1743    #[tokio::test]
1744    async fn shutdown_with_timeout_succeeds_when_task_exits_cooperatively() {
1745        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1746        let shutdown_requested = Arc::new(AtomicBool::new(false));
1747        let supervisor = supervisor_from_shutdown_channel(
1748            shutdown_tx,
1749            Arc::clone(&shutdown_requested),
1750            vec![test_shutdown_task("cooperative-timeout-loop", async move {
1751                while !*shutdown_rx.borrow() {
1752                    if shutdown_rx.changed().await.is_err() {
1753                        break;
1754                    }
1755                }
1756            })],
1757        );
1758
1759        supervisor
1760            .shutdown_with_timeout(Duration::from_secs(1))
1761            .await
1762            .expect("cooperative task should shut down before timeout");
1763    }
1764
1765    #[tokio::test]
1766    async fn shutdown_with_timeout_pre_shutdown_error_allows_remaining_task_to_exit() {
1767        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1768        let shutdown_requested = Arc::new(AtomicBool::new(false));
1769        let dropped = Arc::new(AtomicBool::new(false));
1770        let drop_flag = DropFlag(Arc::clone(&dropped));
1771        let tasks = vec![
1772            test_task("finished-before-shutdown", async {}),
1773            test_shutdown_task("cooperative-after-error", async move {
1774                let _drop_flag = drop_flag;
1775                while !*shutdown_rx.borrow() {
1776                    if shutdown_rx.changed().await.is_err() {
1777                        break;
1778                    }
1779                }
1780            }),
1781        ];
1782
1783        while !tasks[0].handle.is_finished() {
1784            tokio::task::yield_now().await;
1785        }
1786
1787        let supervisor =
1788            supervisor_from_shutdown_channel(shutdown_tx, Arc::clone(&shutdown_requested), tasks);
1789        let error = supervisor
1790            .shutdown_with_timeout(Duration::from_secs(1))
1791            .await
1792            .expect_err("pre-shutdown task exit should fail");
1793
1794        match error {
1795            Error::Runtime(RuntimeError::TaskExitedUnexpectedly { task }) => {
1796                assert_eq!(task, "finished-before-shutdown");
1797            }
1798            other => panic!("expected pre-shutdown task exit, got {other:?}"),
1799        }
1800        assert!(dropped.load(Ordering::SeqCst));
1801    }
1802
1803    #[tokio::test]
1804    async fn shutdown_with_timeout_reports_timeout_after_pre_shutdown_error() {
1805        let tasks = vec![
1806            test_task("finished-before-shutdown", async {}),
1807            test_task("pending-after-error", async {
1808                std::future::pending::<()>().await;
1809            }),
1810        ];
1811
1812        while !tasks[0].handle.is_finished() {
1813            tokio::task::yield_now().await;
1814        }
1815
1816        let error = Supervisor::from_tasks_for_tests(tasks)
1817            .shutdown_with_timeout(Duration::from_millis(1))
1818            .await
1819            .expect_err("pre-shutdown task exit with stuck drain should time out");
1820
1821        match error {
1822            Error::Runtime(RuntimeError::ShutdownTimeoutAfterTaskError { timeout, source }) => {
1823                assert_eq!(timeout, Duration::from_millis(1));
1824                match *source {
1825                    RuntimeError::TaskExitedUnexpectedly { task } => {
1826                        assert_eq!(task, "finished-before-shutdown");
1827                    }
1828                    other => panic!("expected pre-shutdown task exit source, got {other:?}"),
1829                }
1830            }
1831            other => panic!("expected shutdown timeout after task error, got {other:?}"),
1832        }
1833    }
1834
1835    #[tokio::test]
1836    async fn shutdown_with_timeout_reports_task_error_when_remaining_task_misses_deadline() {
1837        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
1838        let shutdown_requested = Arc::new(AtomicBool::new(false));
1839        let supervisor = supervisor_from_shutdown_channel(
1840            shutdown_tx,
1841            Arc::clone(&shutdown_requested),
1842            vec![
1843                test_shutdown_task("panic-after-timeout-shutdown", async move {
1844                    while !*shutdown_rx.borrow() {
1845                        if shutdown_rx.changed().await.is_err() {
1846                            return;
1847                        }
1848                    }
1849                    panic!("forced live shutdown panic");
1850                }),
1851                test_shutdown_task("pending-after-timeout-panic", async {
1852                    std::future::pending::<()>().await;
1853                }),
1854            ],
1855        );
1856
1857        let error = supervisor
1858            .shutdown_with_timeout(Duration::from_millis(50))
1859            .await
1860            .expect_err("task failure with stuck drain should preserve task error source");
1861
1862        match error {
1863            Error::Runtime(RuntimeError::ShutdownTimeoutAfterTaskError { timeout, source }) => {
1864                assert_eq!(timeout, Duration::from_millis(50));
1865                match *source {
1866                    RuntimeError::TaskJoin { task, source } => {
1867                        assert_eq!(task, "panic-after-timeout-shutdown");
1868                        assert!(source.is_panic());
1869                    }
1870                    other => panic!("expected task join source, got {other:?}"),
1871                }
1872            }
1873            other => panic!("expected timeout after task join error, got {other:?}"),
1874        }
1875    }
1876
1877    #[tokio::test]
1878    async fn panicked_task_maps_to_task_join_error() {
1879        let supervisor = Supervisor::from_tasks_for_tests(vec![test_task("panic-loop", async {
1880            panic!("forced supervisor test panic");
1881        })]);
1882
1883        let error = supervisor
1884            .join()
1885            .await
1886            .expect_err("panicked task should fail");
1887
1888        match error {
1889            Error::Runtime(RuntimeError::TaskJoin { task, source }) => {
1890                assert_eq!(task, "panic-loop");
1891                assert!(source.is_panic());
1892            }
1893            other => panic!("expected task join error, got {other:?}"),
1894        }
1895    }
1896}