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#[must_use]
40pub struct Supervisor {
41 shutdown_tx: watch::Sender<bool>,
42 shutdown_requested: Arc<AtomicBool>,
43 tasks: Vec<RuntimeTask>,
44}
45
46#[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#[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 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 #[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 pub fn request_shutdown(&self) {
142 request_shutdown_signal(&self.shutdown_tx, self.shutdown_requested.as_ref());
143 }
144
145 #[must_use]
148 pub fn is_shutdown_requested(&self) -> bool {
149 self.shutdown_requested.load(Ordering::SeqCst)
150 }
151
152 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 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 pub async fn run_until_shutdown<F>(mut self, shutdown: F, timeout: Duration) -> Result<()>
211 where
212 F: Future<Output = ()>,
213 {
214 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 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 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 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 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 self.request_shutdown();
421 }
422}
423
424impl<'a> SupervisorBuilder<'a> {
425 #[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 #[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 #[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 #[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 #[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 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 pub fn request_shutdown(&self) {
558 request_shutdown_signal(&self.shutdown_tx, self.shutdown_requested.as_ref());
559 }
560
561 #[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 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 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 }
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 .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 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 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}