1use std::borrow::Borrow;
2use std::future::Future;
3use std::sync::Arc;
4use std::time::Duration;
5
6use tokio::runtime::Handle;
7use tracing::warn;
8
9use crate::catalog::JobCatalog;
10use crate::config::{IntentPromoterConfig, JobsConfig};
11use crate::observer::{JobLifecycleObserver, JobLifecycleObservers};
12use crate::registry::JobRegistry;
13use crate::scheduler::run_scheduler_loop;
14use crate::shutdown::{ShutdownHandle, ShutdownSignal};
15use crate::task_group::TaskGroup;
16use crate::{Result, RuntimeError};
17
18const WORKER_TASK: &str = "worker";
19const INTENT_PROMOTER_TASK: &str = "intent_promoter";
20const SCHEDULER_TASK: &str = "scheduler";
21const REAPER_TASK: &str = "reaper";
22
23#[must_use]
34pub struct Supervisor {
35 shutdown: ShutdownSignal,
36 tasks: TaskGroup,
37}
38
39#[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 observers: Vec<Arc<dyn JobLifecycleObserver>>,
60 worker_enabled: bool,
61 intent_promoter_enabled: bool,
62 intent_promoter_config: Option<IntentPromoterConfig>,
63 scheduler_enabled: bool,
64 reaper_enabled: bool,
65}
66
67#[derive(Clone)]
69pub struct SupervisorShutdown {
70 handle: ShutdownHandle,
71}
72
73#[derive(Clone, Copy, Debug, Eq, PartialEq)]
74enum RegistrySource {
75 Registry,
76 Catalog,
77}
78
79impl Supervisor {
80 pub fn builder_from_env(
86 pool: &runledger_postgres::DbPool,
87 ) -> std::result::Result<SupervisorBuilder<'_>, RuntimeError> {
88 let config = JobsConfig::from_env();
89 let intent_promoter_config =
90 IntentPromoterConfig::from_env_with_jobs_config_defaults(&config);
91
92 Self::builder(pool, config)
93 .map(|builder| builder.with_intent_promoter_config(intent_promoter_config))
94 }
95
96 pub fn builder(
102 pool: &runledger_postgres::DbPool,
103 config: JobsConfig,
104 ) -> std::result::Result<SupervisorBuilder<'_>, RuntimeError> {
105 let runtime =
106 Handle::try_current().map_err(|source| RuntimeError::MissingTokioRuntime { source })?;
107
108 Ok(SupervisorBuilder {
109 pool,
110 runtime,
111 registry: None,
112 registry_source: None,
113 mixed_registry_sources: false,
114 config,
115 observers: Vec::new(),
116 worker_enabled: true,
117 intent_promoter_enabled: true,
118 intent_promoter_config: None,
119 scheduler_enabled: true,
120 reaper_enabled: true,
121 })
122 }
123
124 #[must_use]
127 pub fn shutdown_handle(&self) -> SupervisorShutdown {
128 SupervisorShutdown {
129 handle: self.shutdown.handle(),
130 }
131 }
132
133 pub fn request_shutdown(&self) {
135 self.shutdown.request();
136 }
137
138 #[must_use]
141 pub fn is_shutdown_requested(&self) -> bool {
142 self.shutdown.is_requested()
143 }
144
145 pub async fn join(mut self) -> Result<()> {
155 let shutdown = self.shutdown.clone();
156 self.tasks.join(&shutdown).await
157 }
158
159 pub async fn shutdown(mut self) -> Result<()> {
168 let shutdown = self.shutdown.clone();
169 self.tasks.shutdown(&shutdown).await
170 }
171
172 pub async fn run_until_shutdown<F>(mut self, shutdown: F, timeout: Duration) -> Result<()>
196 where
197 F: Future<Output = ()>,
198 {
199 let shutdown_signal = self.shutdown.clone();
200 self.tasks
201 .run_until_shutdown(shutdown, timeout, &shutdown_signal)
202 .await
203 }
204
205 pub async fn shutdown_with_timeout(mut self, timeout: Duration) -> Result<()> {
222 let shutdown = self.shutdown.clone();
223 self.tasks.shutdown_with_timeout(timeout, &shutdown).await
224 }
225}
226
227impl Drop for Supervisor {
228 fn drop(&mut self) {
229 if !self.tasks.is_empty() {
230 warn!(
231 task_count = self.tasks.len(),
232 "dropping jobs runtime supervisor before joining tasks; tasks may continue detached after shutdown is requested and later panics will not be observed"
233 );
234 }
235 self.request_shutdown();
237 }
238}
239
240impl<'a> SupervisorBuilder<'a> {
241 #[must_use = "builder methods return an updated builder value"]
246 pub fn with_registry(mut self, registry: JobRegistry) -> Self {
247 self.mixed_registry_sources |= self.registry_source == Some(RegistrySource::Catalog);
248 self.registry_source = Some(RegistrySource::Registry);
249 self.registry = Some(registry);
250 self
251 }
252
253 #[must_use = "builder methods return an updated builder value"]
266 pub fn with_catalog(mut self, catalog: impl Borrow<JobCatalog>) -> Self {
267 self.mixed_registry_sources |= self.registry_source == Some(RegistrySource::Registry);
268 self.registry_source = Some(RegistrySource::Catalog);
269 self.registry = Some(catalog.borrow().to_registry());
270 self
271 }
272
273 #[must_use = "builder methods return an updated builder value"]
276 pub fn disable_worker(mut self) -> Self {
277 self.worker_enabled = false;
278 self.intent_promoter_enabled = false;
279 self
280 }
281
282 #[must_use = "builder methods return an updated builder value"]
290 pub fn disable_intent_promoter(mut self) -> Self {
291 self.intent_promoter_enabled = false;
292 self
293 }
294
295 #[must_use = "builder methods return an updated builder value"]
300 pub fn with_intent_promoter_config(mut self, config: IntentPromoterConfig) -> Self {
301 self.intent_promoter_config = Some(config);
302 self
303 }
304
305 #[must_use = "builder methods return an updated builder value"]
307 pub fn disable_scheduler(mut self) -> Self {
308 self.scheduler_enabled = false;
309 self
310 }
311
312 #[must_use = "builder methods return an updated builder value"]
314 pub fn disable_reaper(mut self) -> Self {
315 self.reaper_enabled = false;
316 self
317 }
318
319 #[must_use = "builder methods return an updated builder value"]
324 pub fn with_job_lifecycle_observer(
325 mut self,
326 observer: impl JobLifecycleObserver + 'static,
327 ) -> Self {
328 self.observers.push(Arc::new(observer));
329 self
330 }
331
332 pub fn build(self) -> std::result::Result<Supervisor, RuntimeError> {
337 let Self {
338 pool,
339 runtime,
340 registry,
341 registry_source: _,
342 mixed_registry_sources,
343 config,
344 observers,
345 worker_enabled,
346 intent_promoter_enabled,
347 intent_promoter_config,
348 scheduler_enabled,
349 reaper_enabled,
350 } = self;
351
352 config
353 .validate()
354 .map_err(|source| RuntimeError::InvalidJobsConfig { source })?;
355 let intent_promoter_config = intent_promoter_config
356 .unwrap_or_else(|| IntentPromoterConfig::from_jobs_config(&config));
357 if intent_promoter_enabled {
358 intent_promoter_config
359 .validate()
360 .map_err(|source| RuntimeError::InvalidJobsConfig { source })?;
361 }
362
363 if mixed_registry_sources {
364 return Err(RuntimeError::MixedRegistrySources);
365 }
366
367 let registry = match registry {
368 Some(registry) => registry,
369 None if worker_enabled || reaper_enabled => {
370 return Err(RuntimeError::MissingRegistry {
371 worker_enabled,
372 reaper_enabled,
373 });
374 }
375 None => JobRegistry::new(),
376 };
377
378 let (shutdown, shutdown_rx) = ShutdownSignal::channel();
379 let mut tasks = TaskGroup::new();
380 let observers = JobLifecycleObservers::from_arc_observers(observers);
381
382 if intent_promoter_enabled {
383 tasks.spawn_on(&runtime, INTENT_PROMOTER_TASK, {
384 let pool = pool.clone();
385 let registry = registry.clone();
386 let shutdown_rx = shutdown_rx.clone();
387 async move {
388 crate::intent_promoter::run_intent_promoter_loop_with_config(
389 pool,
390 registry,
391 intent_promoter_config,
392 shutdown_rx,
393 )
394 .await
395 }
396 });
397 }
398
399 if worker_enabled {
400 tasks.spawn_on(&runtime, WORKER_TASK, {
401 let pool = pool.clone();
402 let registry = registry.clone();
403 let config = config.clone();
404 let shutdown_rx = shutdown_rx.clone();
405 let observers = observers.clone();
406 async move {
407 crate::worker::run_worker_loop_with_observer(
408 pool,
409 registry,
410 config,
411 shutdown_rx,
412 observers,
413 )
414 .await
415 }
416 });
417 }
418
419 if scheduler_enabled {
420 tasks.spawn_on(&runtime, SCHEDULER_TASK, {
421 let pool = pool.clone();
422 let config = config.clone();
423 let shutdown_rx = shutdown_rx.clone();
424 async move { run_scheduler_loop(pool, config, shutdown_rx).await }
425 });
426 }
427
428 if reaper_enabled {
429 let pool = pool.clone();
430 let registry = registry.clone();
431 let config = config.clone();
432 let shutdown_rx = shutdown_rx.clone();
433 let observers = observers.clone();
434 tasks.spawn_on(&runtime, REAPER_TASK, async move {
435 crate::reaper::run_reaper_loop_with_observer(
436 pool,
437 registry,
438 config,
439 shutdown_rx,
440 observers,
441 )
442 .await
443 });
444 }
445
446 Ok(Supervisor { shutdown, tasks })
447 }
448}
449
450impl SupervisorShutdown {
451 pub fn request_shutdown(&self) {
453 self.handle.request();
454 }
455
456 #[must_use]
458 pub fn is_shutdown_requested(&self) -> bool {
459 self.handle.is_requested()
460 }
461}
462
463#[cfg(test)]
464mod tests {
465 use std::time::Duration;
466
467 use sqlx::postgres::PgPoolOptions;
468 use tokio::time::timeout;
469
470 use super::*;
471
472 const UNUSED_LAZY_POOL_URL: &str = "postgres://postgres:postgres@127.0.0.1:65535/runledger";
473
474 fn lazy_pool() -> runledger_postgres::DbPool {
475 PgPoolOptions::new()
476 .connect_lazy(UNUSED_LAZY_POOL_URL)
479 .expect("construct lazy pool")
480 }
481
482 fn test_config() -> JobsConfig {
483 JobsConfig {
484 worker_id: "supervisor-test-worker".to_string(),
485 poll_interval: Duration::from_millis(25),
486 claim_batch_size: 4,
487 lease_ttl_seconds: 10,
488 max_global_concurrency: 4,
489 reaper_interval: Duration::from_millis(50),
490 schedule_poll_interval: Duration::from_millis(50),
491 reaper_retry_delay_ms: 1_000,
492 }
493 }
494
495 fn empty_builder(pool: &runledger_postgres::DbPool) -> SupervisorBuilder<'_> {
496 Supervisor::builder(pool, test_config()).expect("supervisor builder has runtime")
497 }
498
499 fn missing_registry_flags(builder: SupervisorBuilder<'_>) -> (bool, bool) {
500 match builder.build() {
501 Err(RuntimeError::MissingRegistry {
502 worker_enabled,
503 reaper_enabled,
504 }) => (worker_enabled, reaper_enabled),
505 Ok(_) => panic!("missing registry should be a build error"),
506 Err(other) => panic!("expected missing registry error, got {other:?}"),
507 }
508 }
509
510 fn task_names(supervisor: &Supervisor) -> Vec<&'static str> {
511 supervisor.tasks.names_for_tests()
512 }
513
514 async fn abort_supervisor_tasks(mut supervisor: Supervisor) {
515 supervisor.tasks.abort_all_for_tests().await;
516 }
517
518 #[tokio::test]
519 async fn builder_defaults_enable_all_loops() {
520 let pool = lazy_pool();
521 let builder = empty_builder(&pool);
522
523 assert!(builder.worker_enabled);
524 assert!(builder.intent_promoter_enabled);
525 assert_eq!(builder.intent_promoter_config, None);
526 assert!(builder.scheduler_enabled);
527 assert!(builder.reaper_enabled);
528 assert!(builder.registry.is_none());
529 assert_eq!(builder.registry_source, None);
530 assert!(!builder.mixed_registry_sources);
531 }
532
533 #[tokio::test]
534 async fn environment_builder_explicitly_configures_intent_promoter() {
535 let pool = lazy_pool();
536 let builder = Supervisor::builder_from_env(&pool).expect("build supervisor from env");
537
538 assert!(builder.intent_promoter_config.is_some());
539 }
540
541 #[tokio::test]
542 async fn builder_accepts_registry_for_worker_and_reaper_loops() {
543 let pool = lazy_pool();
544 let builder = empty_builder(&pool).with_registry(JobRegistry::new());
545
546 assert!(builder.registry.is_some());
547 assert_eq!(builder.registry_source, Some(RegistrySource::Registry));
548 assert!(!builder.mixed_registry_sources);
549 }
550
551 #[tokio::test]
552 async fn builder_rejects_mixed_registry_sources() {
553 let pool = lazy_pool();
554 let registry_then_catalog = empty_builder(&pool)
555 .with_registry(JobRegistry::new())
556 .with_catalog(JobCatalog::new())
557 .disable_worker()
558 .disable_reaper()
559 .build();
560 let Err(registry_then_catalog) = registry_then_catalog else {
561 panic!("mixed registry sources should be rejected");
562 };
563 assert!(matches!(
564 registry_then_catalog,
565 RuntimeError::MixedRegistrySources
566 ));
567
568 let catalog_then_registry = empty_builder(&pool)
569 .with_catalog(JobCatalog::new())
570 .with_registry(JobRegistry::new())
571 .disable_worker()
572 .disable_reaper()
573 .build();
574 let Err(catalog_then_registry) = catalog_then_registry else {
575 panic!("mixed registry sources should be rejected");
576 };
577 assert!(matches!(
578 catalog_then_registry,
579 RuntimeError::MixedRegistrySources
580 ));
581 }
582
583 #[tokio::test]
584 async fn builder_requires_registry_when_worker_or_reaper_is_enabled() {
585 let pool = lazy_pool();
586
587 assert_eq!(missing_registry_flags(empty_builder(&pool)), (true, true));
588 assert_eq!(
589 missing_registry_flags(empty_builder(&pool).disable_scheduler().disable_reaper()),
590 (true, false)
591 );
592 assert_eq!(
593 missing_registry_flags(empty_builder(&pool).disable_worker().disable_scheduler()),
594 (false, true)
595 );
596 }
597
598 #[tokio::test]
599 async fn builder_rejects_invalid_direct_config_values_before_spawning_loops() {
600 let cases = [
601 {
602 let mut config = test_config();
603 config.max_global_concurrency = 0;
604 (
605 config,
606 crate::config::JobsConfigValidationError::InvalidMaxGlobalConcurrency,
607 )
608 },
609 {
610 let mut config = test_config();
611 config.claim_batch_size = 0;
612 (
613 config,
614 crate::config::JobsConfigValidationError::InvalidClaimBatchSize { actual: 0 },
615 )
616 },
617 {
618 let mut config = test_config();
619 config.lease_ttl_seconds = 0;
620 (
621 config,
622 crate::config::JobsConfigValidationError::InvalidLeaseTtlSeconds { actual: 0 },
623 )
624 },
625 ];
626
627 for (config, expected) in cases {
628 let pool = lazy_pool();
629 let result = Supervisor::builder(&pool, config)
630 .expect("supervisor builder has runtime")
631 .disable_worker()
632 .disable_scheduler()
633 .disable_reaper()
634 .build();
635 let Err(error) = result else {
636 panic!("invalid direct config should be rejected");
637 };
638
639 match error {
640 RuntimeError::InvalidJobsConfig { source } => {
641 assert_eq!(source, expected);
642 }
643 other => panic!("expected invalid jobs config error, got {other:?}"),
644 }
645 }
646 }
647
648 #[test]
649 fn builder_requires_tokio_runtime_before_cloning_pool() {
650 let runtime = tokio::runtime::Runtime::new().expect("construct Tokio runtime");
651 let pool = runtime.block_on(async { lazy_pool() });
652 let error = match Supervisor::builder(&pool, test_config()) {
653 Err(error) => error,
654 Ok(builder) => {
655 drop(builder);
656 runtime.block_on(async {
657 pool.close().await;
658 });
659 std::mem::forget(pool);
660 panic!("missing Tokio runtime should be a builder error");
661 }
662 };
663
664 runtime.block_on(async {
669 pool.close().await;
670 });
671 std::mem::forget(pool);
672 match error {
673 RuntimeError::MissingTokioRuntime { .. } => {}
674 other => panic!("expected missing Tokio runtime error, got {other:?}"),
675 }
676 }
677
678 #[tokio::test]
679 async fn builder_can_disable_each_loop() {
680 let pool = lazy_pool();
681 let builder = empty_builder(&pool)
682 .disable_worker()
683 .disable_scheduler()
684 .disable_reaper();
685
686 assert!(!builder.worker_enabled);
687 assert!(!builder.intent_promoter_enabled);
688 assert!(!builder.scheduler_enabled);
689 assert!(!builder.reaper_enabled);
690
691 let worker_without_promoter = empty_builder(&pool).disable_intent_promoter();
692 assert!(worker_without_promoter.worker_enabled);
693 assert!(!worker_without_promoter.intent_promoter_enabled);
694
695 let promoter_config = IntentPromoterConfig::new(Duration::from_secs(2), 7);
696 let customized = empty_builder(&pool).with_intent_promoter_config(promoter_config);
697 assert_eq!(customized.intent_promoter_config, Some(promoter_config));
698 }
699
700 #[tokio::test]
701 async fn builder_spawns_only_enabled_tasks() {
702 let pool = lazy_pool();
703
704 let all_disabled = empty_builder(&pool)
705 .disable_worker()
706 .disable_scheduler()
707 .disable_reaper()
708 .build()
709 .expect("all-disabled supervisor should build");
710 assert_eq!(task_names(&all_disabled), Vec::<&'static str>::new());
711 abort_supervisor_tasks(all_disabled).await;
712
713 let scheduler_only = empty_builder(&pool)
714 .disable_worker()
715 .disable_reaper()
716 .build()
717 .expect("scheduler-only supervisor should not require registry");
718 assert_eq!(task_names(&scheduler_only), vec![SCHEDULER_TASK]);
719 abort_supervisor_tasks(scheduler_only).await;
720
721 let worker_only = empty_builder(&pool)
722 .with_registry(JobRegistry::new())
723 .disable_scheduler()
724 .disable_reaper()
725 .build()
726 .expect("worker-only supervisor should build with registry");
727 assert_eq!(
728 task_names(&worker_only),
729 vec![INTENT_PROMOTER_TASK, WORKER_TASK]
730 );
731 abort_supervisor_tasks(worker_only).await;
732
733 let worker_without_promoter = empty_builder(&pool)
734 .with_registry(JobRegistry::new())
735 .disable_intent_promoter()
736 .disable_scheduler()
737 .disable_reaper()
738 .build()
739 .expect("worker should run without intent promotion");
740 assert_eq!(task_names(&worker_without_promoter), vec![WORKER_TASK]);
741 abort_supervisor_tasks(worker_without_promoter).await;
742
743 let reaper_only = empty_builder(&pool)
744 .with_registry(JobRegistry::new())
745 .disable_worker()
746 .disable_scheduler()
747 .build()
748 .expect("reaper-only supervisor should build with registry");
749 assert_eq!(task_names(&reaper_only), vec![REAPER_TASK]);
750 abort_supervisor_tasks(reaper_only).await;
751
752 let all_enabled = empty_builder(&pool)
753 .with_registry(JobRegistry::new())
754 .build()
755 .expect("all-enabled supervisor should build with registry");
756 assert_eq!(
757 task_names(&all_enabled),
758 vec![
759 INTENT_PROMOTER_TASK,
760 WORKER_TASK,
761 SCHEDULER_TASK,
762 REAPER_TASK
763 ]
764 );
765 abort_supervisor_tasks(all_enabled).await;
766 }
767
768 #[tokio::test]
769 async fn all_disabled_supervisor_join_and_shutdown_succeed() {
770 Supervisor::builder(&lazy_pool(), test_config())
771 .expect("supervisor builder has runtime")
772 .disable_worker()
773 .disable_scheduler()
774 .disable_reaper()
775 .build()
776 .expect("all-disabled supervisor should build")
777 .join()
778 .await
779 .expect("all-disabled supervisor should join");
780
781 Supervisor::builder(&lazy_pool(), test_config())
782 .expect("supervisor builder has runtime")
783 .disable_worker()
784 .disable_scheduler()
785 .disable_reaper()
786 .build()
787 .expect("all-disabled supervisor should build")
788 .shutdown()
789 .await
790 .expect("all-disabled supervisor should shut down");
791 }
792
793 #[tokio::test]
794 async fn shutdown_handle_can_request_shutdown_before_join() {
795 let supervisor = Supervisor::builder(&lazy_pool(), test_config())
796 .expect("supervisor builder has runtime")
797 .disable_worker()
798 .disable_scheduler()
799 .disable_reaper()
800 .build()
801 .expect("all-disabled supervisor should build");
802 let shutdown = supervisor.shutdown_handle();
803 let cloned_shutdown = shutdown.clone();
804
805 cloned_shutdown.request_shutdown();
806
807 assert!(shutdown.is_shutdown_requested());
808 assert!(supervisor.is_shutdown_requested());
809 supervisor
810 .join()
811 .await
812 .expect("supervisor should join after shutdown handle request");
813 }
814
815 #[tokio::test]
816 async fn run_until_shutdown_with_no_tasks_waits_for_signal() {
817 let supervisor = Supervisor::builder(&lazy_pool(), test_config())
818 .expect("supervisor builder has runtime")
819 .disable_worker()
820 .disable_scheduler()
821 .disable_reaper()
822 .build()
823 .expect("all-disabled supervisor should build");
824 let (signal_tx, signal_rx) = tokio::sync::oneshot::channel();
825 let mut run = tokio::spawn(supervisor.run_until_shutdown(
826 async move {
827 signal_rx.await.expect("shutdown signal should be sent");
828 },
829 Duration::from_secs(1),
830 ));
831
832 assert!(
833 timeout(Duration::from_millis(50), &mut run).await.is_err(),
834 "all-disabled supervisor should wait for the shutdown signal"
835 );
836
837 signal_tx.send(()).expect("signal receiver should be alive");
838 run.await
839 .expect("run-until-shutdown task should join")
840 .expect("all-disabled supervisor should complete after signal");
841 }
842}