aster_forge_tasks/
runtime.rs

1//! Runtime worker loops for background task systems.
2//!
3//! This module contains the product-neutral runtime mechanics shared by services that run
4//! background work: a shutdown-aware task container, periodic worker loops, panic recovery for one
5//! recorded iteration, jittered sleep calculation, and the adaptive idle backoff used by database
6//! dispatchers. Product crates keep ownership of their task names, runtime configuration, outcome
7//! records, wakeup sources, and persistence layer by passing small closures into these helpers.
8
9use std::any::Any;
10use std::future::Future;
11use std::panic::AssertUnwindSafe;
12use std::time::Duration;
13
14use chrono::{DateTime, Utc};
15use futures::FutureExt;
16use rand::RngExt;
17use tokio::task::JoinSet;
18use tokio_util::sync::CancellationToken;
19use tracing::Instrument;
20
21/// Default grace period used when shutting down background workers.
22pub const BACKGROUND_TASK_SHUTDOWN_GRACE: Duration = Duration::from_secs(30);
23/// Minimum backoff used after a dispatch iteration returns an error.
24pub const BACKGROUND_TASK_DISPATCH_ERROR_BACKOFF_CAP: Duration = Duration::from_secs(5);
25
26/// Reason a dispatch loop is about to run an iteration.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum BackgroundTaskDispatchTrigger {
29    /// Initial run immediately after worker startup.
30    Startup,
31    /// Regular timer-based polling.
32    Timer,
33    /// Product wakeup signal, usually emitted after enqueueing a task.
34    Wakeup,
35}
36
37/// Activity summary returned by one dispatch iteration.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct BackgroundTaskDispatchIteration {
40    has_activity: bool,
41    failed: bool,
42}
43
44impl BackgroundTaskDispatchIteration {
45    /// Creates an idle dispatch iteration.
46    #[must_use]
47    pub const fn idle() -> Self {
48        Self {
49            has_activity: false,
50            failed: false,
51        }
52    }
53
54    /// Creates a dispatch iteration that claimed or completed work.
55    #[must_use]
56    pub const fn active() -> Self {
57        Self {
58            has_activity: true,
59            failed: false,
60        }
61    }
62
63    /// Creates a dispatch iteration that failed.
64    #[must_use]
65    pub const fn failed() -> Self {
66        Self {
67            has_activity: false,
68            failed: true,
69        }
70    }
71
72    /// Returns whether the iteration performed task work.
73    #[must_use]
74    pub const fn has_activity(self) -> bool {
75        self.has_activity
76    }
77
78    /// Returns whether the iteration failed.
79    #[must_use]
80    pub const fn failed_to_dispatch(self) -> bool {
81        self.failed
82    }
83}
84
85/// Adaptive idle backoff for background task dispatch workers.
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87pub struct BackgroundTaskDispatchBackoff {
88    idle_interval: Duration,
89    last_error: bool,
90}
91
92impl BackgroundTaskDispatchBackoff {
93    /// Creates a dispatch backoff state using the current runtime intervals.
94    #[must_use]
95    pub fn new(base_interval: Duration, max_interval: Duration) -> Self {
96        Self {
97            idle_interval: effective_dispatch_base_interval(base_interval, max_interval),
98            last_error: false,
99        }
100    }
101
102    /// Returns the sleep duration for the next dispatch loop wait.
103    #[must_use]
104    pub fn sleep_duration(&self, base_interval: Duration, max_interval: Duration) -> Duration {
105        let base_interval = effective_dispatch_base_interval(base_interval, max_interval);
106        let max_interval = effective_dispatch_max_interval(base_interval, max_interval);
107        if self.last_error {
108            return base_interval.max(BACKGROUND_TASK_DISPATCH_ERROR_BACKOFF_CAP);
109        }
110        self.idle_interval.max(base_interval).min(max_interval)
111    }
112
113    /// Records the last dispatch iteration and updates the idle/error backoff state.
114    pub fn record_iteration(
115        &mut self,
116        trigger: BackgroundTaskDispatchTrigger,
117        iteration: BackgroundTaskDispatchIteration,
118        base_interval: Duration,
119        max_interval: Duration,
120    ) {
121        let base_interval = effective_dispatch_base_interval(base_interval, max_interval);
122        let max_interval = effective_dispatch_max_interval(base_interval, max_interval);
123        if iteration.failed {
124            self.idle_interval = base_interval;
125            self.last_error = true;
126            return;
127        }
128        if iteration.has_activity || matches!(trigger, BackgroundTaskDispatchTrigger::Wakeup) {
129            self.idle_interval = base_interval;
130            self.last_error = false;
131            return;
132        }
133        self.idle_interval = self
134            .idle_interval
135            .max(base_interval)
136            .saturating_mul(2)
137            .min(max_interval);
138        self.last_error = false;
139    }
140}
141
142/// Shutdown-aware collection of spawned background workers.
143pub struct BackgroundTasks {
144    shutdown_token: CancellationToken,
145    handles: JoinSet<()>,
146    shutdown_grace: Duration,
147}
148
149impl BackgroundTasks {
150    /// Creates a task collection with a fresh shutdown token and the default shutdown grace.
151    #[must_use]
152    pub fn new() -> Self {
153        Self::with_shutdown_token(CancellationToken::new())
154    }
155
156    /// Creates a task collection using an externally owned shutdown token.
157    #[must_use]
158    pub fn with_shutdown_token(shutdown_token: CancellationToken) -> Self {
159        Self::with_shutdown_token_and_grace(shutdown_token, BACKGROUND_TASK_SHUTDOWN_GRACE)
160    }
161
162    /// Creates a task collection using an externally owned token and custom shutdown grace.
163    #[must_use]
164    pub fn with_shutdown_token_and_grace(
165        shutdown_token: CancellationToken,
166        shutdown_grace: Duration,
167    ) -> Self {
168        Self {
169            shutdown_token,
170            handles: JoinSet::new(),
171            shutdown_grace,
172        }
173    }
174
175    /// Returns a clone of the shutdown token observed by all workers in this collection.
176    #[must_use]
177    pub fn shutdown_token(&self) -> CancellationToken {
178        self.shutdown_token.clone()
179    }
180
181    /// Spawns a worker into the collection.
182    pub fn push<F>(&mut self, task: F)
183    where
184        F: Future<Output = ()> + Send + 'static,
185    {
186        self.handles.spawn(task);
187    }
188
189    /// Requests shutdown, waits for cooperative exit, and aborts remaining workers after grace.
190    pub async fn shutdown(self) {
191        let BackgroundTasks {
192            shutdown_token,
193            mut handles,
194            shutdown_grace,
195        } = self;
196        shutdown_token.cancel();
197
198        let graceful_shutdown = async {
199            drain_task_handles(&mut handles).await;
200        };
201        if tokio::time::timeout(shutdown_grace, graceful_shutdown)
202            .await
203            .is_err()
204        {
205            let aborted = handles.len();
206            handles.abort_all();
207            tracing::warn!(
208                aborted,
209                grace_secs = shutdown_grace.as_secs(),
210                "background tasks did not stop before the shutdown grace period; aborting remaining workers"
211            );
212            drain_task_handles(&mut handles).await;
213        }
214    }
215}
216
217/// Drains the join set, logging workers that exited by panic.
218///
219/// Cancellation is an expected outcome once `abort_all` has run, so only
220/// panics are reported. Returns how many workers panicked.
221async fn drain_task_handles(handles: &mut JoinSet<()>) -> usize {
222    let mut panicked = 0;
223    while let Some(result) = handles.join_next().await {
224        match result {
225            Err(error) if error.is_panic() => {
226                panicked += 1;
227                tracing::error!(%error, "background task worker panicked");
228            }
229            Ok(()) | Err(_) => {}
230        }
231    }
232    panicked
233}
234
235impl Default for BackgroundTasks {
236    fn default() -> Self {
237        Self::new()
238    }
239}
240
241/// Runs a singleton background task group behind a Forge runtime lease.
242///
243/// This helper keeps the repeated product glue in one place: acquire a
244/// process-level singleton lease, start a [`BackgroundTasks`] group with the
245/// lease-scoped shutdown token, and shut that group down when the lease is lost
246/// or the process is terminating. Products still own the lease id, lease store,
247/// worker list, and runtime owner id.
248pub async fn run_leased_background_tasks<Store, StartFn>(
249    store: Store,
250    config: aster_forge_runtime::RuntimeLeaseConfig,
251    shutdown_token: CancellationToken,
252    start_background_tasks: StartFn,
253) where
254    Store: aster_forge_runtime::RuntimeLeaseStore,
255    StartFn: FnMut(CancellationToken) -> BackgroundTasks + Send,
256{
257    aster_forge_runtime::run_runtime_lease_supervisor(
258        store,
259        config,
260        shutdown_token,
261        start_background_tasks,
262        |background_tasks| async move {
263            background_tasks.shutdown().await;
264        },
265    )
266    .await;
267}
268
269/// Product callbacks used by panic-protected recorded task iterations.
270pub struct RecordedTaskHooks<TaskFn, PanicFn, RecordFn> {
271    /// Runs the product task body.
272    pub task_fn: TaskFn,
273    /// Converts a panic payload message into the product's runtime outcome type.
274    pub panic_outcome: PanicFn,
275    /// Persists or observes one runtime task outcome.
276    pub record_outcome: RecordFn,
277}
278
279impl<TaskFn, PanicFn, RecordFn> RecordedTaskHooks<TaskFn, PanicFn, RecordFn> {
280    /// Creates recorded task hooks from product callbacks.
281    pub const fn new(task_fn: TaskFn, panic_outcome: PanicFn, record_outcome: RecordFn) -> Self {
282        Self {
283            task_fn,
284            panic_outcome,
285            record_outcome,
286        }
287    }
288}
289
290/// Configuration for one periodic runtime task worker.
291pub struct PeriodicTask<Name, State, IntervalFn, TaskFn, PanicFn, RecordFn> {
292    /// Product task identifier.
293    pub name: Name,
294    /// Stable task name used in tracing spans.
295    pub task_name: &'static str,
296    /// Reads the latest product-configured interval.
297    pub interval_fn: IntervalFn,
298    /// Optional upper bound for positive jitter.
299    pub jitter_cap: Option<Duration>,
300    /// Shared shutdown token.
301    pub shutdown_token: CancellationToken,
302    /// Product runtime state passed to callbacks.
303    pub state: State,
304    /// Product callbacks for execution, panic conversion, and recording.
305    pub hooks: RecordedTaskHooks<TaskFn, PanicFn, RecordFn>,
306}
307
308/// Runs a periodic task until shutdown.
309///
310/// The first iteration runs immediately unless the token is already cancelled. Each later
311/// iteration sleeps using the latest product-provided interval and optional jitter cap. Panics in
312/// one iteration are converted into a product outcome through `panic_outcome` and then recorded by
313/// `record_outcome`; they do not kill the worker loop.
314pub async fn run_periodic_task<
315    Name,
316    State,
317    IntervalFn,
318    TaskFn,
319    TaskFut,
320    PanicFn,
321    RecordFn,
322    RecordFut,
323    Outcome,
324>(
325    task: PeriodicTask<Name, State, IntervalFn, TaskFn, PanicFn, RecordFn>,
326) where
327    Name: Copy + Send + 'static,
328    State: Clone + Send + Sync + 'static,
329    IntervalFn: Fn(&State) -> Duration + Send + Sync + 'static,
330    TaskFn: Fn(State) -> TaskFut + Send + Sync + 'static,
331    TaskFut: Future<Output = Outcome> + Send,
332    PanicFn: Fn(String) -> Outcome + Send + Sync + 'static,
333    RecordFn:
334        Fn(State, Name, DateTime<Utc>, DateTime<Utc>, Outcome) -> RecordFut + Send + Sync + 'static,
335    RecordFut: Future<Output = ()> + Send,
336    Outcome: Send + 'static,
337{
338    let PeriodicTask {
339        name,
340        task_name,
341        interval_fn,
342        jitter_cap,
343        shutdown_token,
344        state,
345        hooks,
346    } = task;
347    let RecordedTaskHooks {
348        task_fn,
349        panic_outcome,
350        record_outcome,
351    } = hooks;
352
353    if shutdown_token.is_cancelled() {
354        return;
355    }
356    run_recorded_task_iteration(
357        name,
358        task_name,
359        state.clone(),
360        &task_fn,
361        &panic_outcome,
362        &record_outcome,
363    )
364    .instrument(tracing::info_span!("bg_task", task.name = task_name))
365    .await;
366
367    loop {
368        let sleep_duration = periodic_sleep_duration(interval_fn(&state), jitter_cap);
369        tokio::select! {
370            biased;
371            () = shutdown_token.cancelled() => break,
372            () = tokio::time::sleep(sleep_duration) => {}
373        }
374
375        if shutdown_token.is_cancelled() {
376            break;
377        }
378
379        run_recorded_task_iteration(
380            name,
381            task_name,
382            state.clone(),
383            &task_fn,
384            &panic_outcome,
385            &record_outcome,
386        )
387        .instrument(tracing::info_span!("bg_task", task.name = task_name))
388        .await;
389    }
390}
391
392/// Runs one panic-protected task iteration and records its outcome.
393pub async fn run_recorded_task_iteration<
394    Name,
395    State,
396    TaskFn,
397    TaskFut,
398    PanicFn,
399    RecordFn,
400    RecordFut,
401    Outcome,
402>(
403    name: Name,
404    task_name: &'static str,
405    state: State,
406    task_fn: &TaskFn,
407    panic_outcome: &PanicFn,
408    record_outcome: &RecordFn,
409) where
410    Name: Copy + Send + 'static,
411    State: Clone + Send + Sync + 'static,
412    TaskFn: Fn(State) -> TaskFut + Send + Sync + 'static,
413    TaskFut: Future<Output = Outcome> + Send,
414    PanicFn: Fn(String) -> Outcome + Send + Sync + 'static,
415    RecordFn:
416        Fn(State, Name, DateTime<Utc>, DateTime<Utc>, Outcome) -> RecordFut + Send + Sync + 'static,
417    RecordFut: Future<Output = ()> + Send,
418    Outcome: Send + 'static,
419{
420    let started_at = Utc::now();
421    let outcome = match AssertUnwindSafe(task_fn(state.clone()))
422        .catch_unwind()
423        .await
424    {
425        Ok(outcome) => outcome,
426        Err(panic) => {
427            let panic_message = panic_payload_message(&panic);
428            tracing::error!("background task '{task_name}' panicked: {panic_message}");
429            panic_outcome(panic_message)
430        }
431    };
432    let finished_at = Utc::now();
433
434    record_outcome(state, name, started_at, finished_at, outcome).await;
435}
436
437/// Runs a wakeable dispatch loop with adaptive idle backoff.
438///
439/// Product crates provide the wakeup future and one dispatch iteration closure. The iteration
440/// closure is responsible for claim/execute logic, panic recovery if desired, metrics, and
441/// persistence of runtime task history.
442pub async fn run_dispatch_worker<State, BaseFn, MaxFn, WakeFn, WakeFut, DispatchFn, DispatchFut>(
443    task_name: &'static str,
444    shutdown_token: CancellationToken,
445    state: State,
446    base_interval_fn: BaseFn,
447    max_interval_fn: MaxFn,
448    wakeup: WakeFn,
449    dispatch_iteration: DispatchFn,
450) where
451    State: Clone + Send + Sync + 'static,
452    BaseFn: Fn(&State) -> Duration + Send + Sync + 'static,
453    MaxFn: Fn(&State) -> Duration + Send + Sync + 'static,
454    WakeFn: Fn(State) -> WakeFut + Send + Sync + 'static,
455    WakeFut: Future<Output = ()> + Send,
456    DispatchFn: Fn(State, CancellationToken) -> DispatchFut + Send + Sync + 'static,
457    DispatchFut: Future<Output = BackgroundTaskDispatchIteration> + Send,
458{
459    let mut backoff =
460        BackgroundTaskDispatchBackoff::new(base_interval_fn(&state), max_interval_fn(&state));
461    if shutdown_token.is_cancelled() {
462        return;
463    }
464    let iteration = dispatch_iteration(state.clone(), shutdown_token.clone())
465        .instrument(tracing::info_span!("bg_task", task.name = task_name))
466        .await;
467    backoff.record_iteration(
468        BackgroundTaskDispatchTrigger::Startup,
469        iteration,
470        base_interval_fn(&state),
471        max_interval_fn(&state),
472    );
473
474    loop {
475        let sleep_duration =
476            backoff.sleep_duration(base_interval_fn(&state), max_interval_fn(&state));
477        let trigger = tokio::select! {
478            biased;
479            () = shutdown_token.cancelled() => break,
480            () = wakeup(state.clone()) => BackgroundTaskDispatchTrigger::Wakeup,
481            () = tokio::time::sleep(sleep_duration) => BackgroundTaskDispatchTrigger::Timer,
482        };
483
484        if shutdown_token.is_cancelled() {
485            break;
486        }
487
488        let iteration = dispatch_iteration(state.clone(), shutdown_token.clone())
489            .instrument(tracing::info_span!("bg_task", task.name = task_name))
490            .await;
491        backoff.record_iteration(
492            trigger,
493            iteration,
494            base_interval_fn(&state),
495            max_interval_fn(&state),
496        );
497    }
498}
499
500/// Returns a periodic delay with bounded positive jitter.
501///
502/// A zero base interval would hot-spin the runner loop — and, for scheduled
503/// tasks, the database — so it is raised to the same one-second floor the
504/// dispatch worker enforces in [`effective_dispatch_base_interval`]. Non-zero
505/// sub-second intervals pass through untouched: they are a deliberate caller
506/// choice (fast tests, high-frequency in-memory polls), not the accident this
507/// floor guards against.
508#[must_use]
509pub fn periodic_sleep_duration(base_interval: Duration, jitter_cap: Option<Duration>) -> Duration {
510    let base_interval = if base_interval.is_zero() {
511        Duration::from_secs(1)
512    } else {
513        base_interval
514    };
515
516    let Some(jitter_cap) = jitter_cap else {
517        return base_interval;
518    };
519
520    let max_jitter_ms = effective_jitter_cap(base_interval, jitter_cap).as_millis();
521    if max_jitter_ms == 0 {
522        return base_interval;
523    }
524
525    let max_jitter_ms = u128_to_u64_saturating(max_jitter_ms.min(u128::from(u64::MAX)));
526    let jitter_ms = rand::rng().random_range(0..=max_jitter_ms);
527    base_interval.saturating_add(Duration::from_millis(jitter_ms))
528}
529
530/// Returns the effective jitter cap for one periodic interval.
531#[must_use]
532pub fn effective_jitter_cap(base_interval: Duration, jitter_cap: Duration) -> Duration {
533    let bounded_ms =
534        u128_to_u64_saturating(base_interval.as_millis().min(u128::from(u64::MAX))) / 10;
535    jitter_cap.min(Duration::from_millis(bounded_ms))
536}
537
538/// Returns the effective dispatch base interval, enforcing a one-second minimum.
539#[must_use]
540pub fn effective_dispatch_base_interval(
541    base_interval: Duration,
542    _max_interval: Duration,
543) -> Duration {
544    if base_interval.is_zero() {
545        return Duration::from_secs(1);
546    }
547    base_interval
548}
549
550/// Returns the effective maximum dispatch interval.
551#[must_use]
552pub fn effective_dispatch_max_interval(
553    base_interval: Duration,
554    max_interval: Duration,
555) -> Duration {
556    max_interval.max(base_interval)
557}
558
559pub(crate) fn panic_payload_message(panic: &Box<dyn Any + Send>) -> String {
560    if let Some(message) = panic.downcast_ref::<&str>() {
561        (*message).to_string()
562    } else if let Some(message) = panic.downcast_ref::<String>() {
563        message.clone()
564    } else {
565        "unknown panic payload".to_string()
566    }
567}
568
569fn u128_to_u64_saturating(value: u128) -> u64 {
570    u64::try_from(value).unwrap_or(u64::MAX)
571}
572
573#[cfg(test)]
574mod tests {
575    use std::sync::Arc;
576    use std::sync::atomic::{AtomicUsize, Ordering};
577
578    use chrono::{DateTime, Utc};
579    use tokio::sync::{Notify, oneshot};
580    use tokio_util::sync::CancellationToken;
581
582    use super::{
583        BACKGROUND_TASK_DISPATCH_ERROR_BACKOFF_CAP, BackgroundTaskDispatchBackoff,
584        BackgroundTaskDispatchIteration, BackgroundTaskDispatchTrigger, BackgroundTasks,
585        PeriodicTask, RecordedTaskHooks, drain_task_handles, effective_jitter_cap,
586        periodic_sleep_duration, run_dispatch_worker, run_leased_background_tasks,
587        run_periodic_task, run_recorded_task_iteration,
588    };
589    use aster_forge_runtime::{
590        RuntimeLeaseAcquire, RuntimeLeaseClaim, RuntimeLeaseConfig, RuntimeLeaseStore,
591    };
592    use std::time::Duration;
593
594    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
595    enum TestTaskName {
596        Cleanup,
597    }
598
599    #[derive(Debug, Clone, PartialEq, Eq)]
600    enum TestOutcome {
601        Succeeded,
602        Failed(String),
603    }
604
605    #[expect(
606        clippy::trivially_copy_pass_by_ref,
607        reason = "The periodic task interval callback receives shared state by reference; this fixture uses unit state."
608    )]
609    fn test_interval(_: &()) -> Duration {
610        Duration::from_mins(1)
611    }
612
613    #[test]
614    fn periodic_sleep_duration_is_unchanged_without_jitter() {
615        let base = Duration::from_secs(5);
616        assert_eq!(periodic_sleep_duration(base, None), base);
617    }
618
619    #[test]
620    fn periodic_sleep_duration_raises_zero_interval_to_one_second_floor() {
621        // A zero interval would hot-spin the runner loop; the scheduled variant would
622        // tight-loop against the database. The dispatch worker already enforces the
623        // same floor via effective_dispatch_base_interval.
624        assert_eq!(
625            periodic_sleep_duration(Duration::ZERO, None),
626            Duration::from_secs(1)
627        );
628
629        // Jitter is derived from the clamped base, so the delay lands in [1s, 1.1s].
630        for _ in 0..64 {
631            let delay = periodic_sleep_duration(Duration::ZERO, Some(Duration::from_secs(30)));
632            assert!(delay >= Duration::from_secs(1));
633            assert!(delay <= Duration::from_millis(1_100));
634        }
635
636        // Non-zero sub-second intervals are a deliberate choice (e.g. tests) and pass through.
637        let tiny = Duration::from_millis(50);
638        assert_eq!(periodic_sleep_duration(tiny, None), tiny);
639    }
640
641    #[tokio::test]
642    async fn drain_task_handles_counts_panics_and_ignores_cancellation() {
643        let mut handles = tokio::task::JoinSet::new();
644        handles.spawn(async { panic!("worker exploded") });
645        handles.spawn(async {});
646        let abort_handle = handles.spawn(async { std::future::pending::<()>().await });
647        abort_handle.abort();
648
649        // The panicked worker is the only one reported: the aborted worker ends in
650        // cancellation, which is the expected outcome after abort_all.
651        let panicked = drain_task_handles(&mut handles).await;
652
653        assert_eq!(panicked, 1);
654    }
655
656    #[test]
657    fn periodic_sleep_duration_caps_jitter_to_ten_percent_of_interval() {
658        let base = Duration::from_secs(5);
659        let cap = Duration::from_secs(30);
660
661        for _ in 0..64 {
662            let delay = periodic_sleep_duration(base, Some(cap));
663            assert!(delay >= base);
664            assert!(delay <= base + Duration::from_millis(500));
665        }
666    }
667
668    #[test]
669    fn periodic_sleep_duration_uses_requested_cap_when_it_is_smaller() {
670        let base = Duration::from_hours(1);
671        let cap = Duration::from_secs(30);
672
673        for _ in 0..64 {
674            let delay = periodic_sleep_duration(base, Some(cap));
675            assert!(delay >= base);
676            assert!(delay <= base + cap);
677        }
678    }
679
680    #[test]
681    fn effective_jitter_cap_handles_zero_interval() {
682        assert_eq!(
683            effective_jitter_cap(Duration::ZERO, Duration::from_secs(30)),
684            Duration::ZERO
685        );
686    }
687
688    #[tokio::test]
689    async fn shutdown_only_awaits_each_handle_once() {
690        let mut tasks = BackgroundTasks::new();
691        tasks.push(async {});
692
693        tasks.shutdown().await;
694    }
695
696    #[tokio::test]
697    async fn external_shutdown_token_stops_background_worker_before_shutdown_join() {
698        let shutdown_token = CancellationToken::new();
699        let mut tasks = BackgroundTasks::with_shutdown_token(shutdown_token.clone());
700        let (stopped_tx, stopped_rx) = oneshot::channel();
701
702        tasks.push({
703            let shutdown_token = shutdown_token.clone();
704            async move {
705                shutdown_token.cancelled().await;
706                let _ = stopped_tx.send(());
707            }
708        });
709
710        shutdown_token.cancel();
711        tokio::time::timeout(Duration::from_millis(50), stopped_rx)
712            .await
713            .expect("background worker should observe external shutdown")
714            .expect("background worker should report shutdown");
715
716        tasks.shutdown().await;
717    }
718
719    #[derive(Clone)]
720    struct AlwaysAcquireLeaseStore;
721
722    #[async_trait::async_trait]
723    impl RuntimeLeaseStore for AlwaysAcquireLeaseStore {
724        type Error = std::convert::Infallible;
725
726        async fn try_acquire(
727            &self,
728            _claim: RuntimeLeaseClaim<'_>,
729        ) -> Result<RuntimeLeaseAcquire, Self::Error> {
730            Ok(RuntimeLeaseAcquire::Acquired)
731        }
732
733        async fn renew(
734            &self,
735            _lease_id: &str,
736            _owner_id: &str,
737            _now: DateTime<Utc>,
738            _expires_at: DateTime<Utc>,
739        ) -> Result<bool, Self::Error> {
740            Ok(true)
741        }
742
743        async fn release(&self, _lease_id: &str, _owner_id: &str) -> Result<(), Self::Error> {
744            Ok(())
745        }
746    }
747
748    #[tokio::test]
749    async fn leased_background_tasks_shutdown_stops_owned_worker_group() {
750        let shutdown_token = CancellationToken::new();
751        let (started_tx, started_rx) = oneshot::channel();
752        let (stopped_tx, stopped_rx) = oneshot::channel();
753        let started_tx = Arc::new(std::sync::Mutex::new(Some(started_tx)));
754        let stopped_tx = Arc::new(std::sync::Mutex::new(Some(stopped_tx)));
755        let handle = tokio::spawn(run_leased_background_tasks(
756            AlwaysAcquireLeaseStore,
757            RuntimeLeaseConfig::new("test.background", "runtime-a")
758                .ttl(Duration::from_millis(100))
759                .renew_interval(Duration::from_millis(10)),
760            shutdown_token.clone(),
761            move |leased_shutdown_token| {
762                let mut tasks = BackgroundTasks::with_shutdown_token(leased_shutdown_token.clone());
763                let started_tx = started_tx.clone();
764                let stopped_tx = stopped_tx.clone();
765                tasks.push(async move {
766                    if let Some(started_tx) = started_tx
767                        .lock()
768                        .expect("test sender mutex should not be poisoned")
769                        .take()
770                    {
771                        let _ = started_tx.send(());
772                    }
773                    leased_shutdown_token.cancelled().await;
774                    if let Some(stopped_tx) = stopped_tx
775                        .lock()
776                        .expect("test sender mutex should not be poisoned")
777                        .take()
778                    {
779                        let _ = stopped_tx.send(());
780                    }
781                });
782                tasks
783            },
784        ));
785
786        tokio::time::timeout(Duration::from_millis(100), started_rx)
787            .await
788            .expect("leased worker should start")
789            .expect("leased worker should report startup");
790        shutdown_token.cancel();
791        tokio::time::timeout(Duration::from_millis(100), stopped_rx)
792            .await
793            .expect("leased worker should observe shutdown")
794            .expect("leased worker should report shutdown");
795        handle.await.expect("lease supervisor should stop cleanly");
796    }
797
798    #[tokio::test]
799    async fn shutdown_aborts_workers_after_custom_grace() {
800        let mut tasks = BackgroundTasks::with_shutdown_token_and_grace(
801            CancellationToken::new(),
802            Duration::from_millis(1),
803        );
804        let calls = Arc::new(AtomicUsize::new(0));
805
806        tasks.push({
807            let calls = calls.clone();
808            async move {
809                calls.fetch_add(1, Ordering::SeqCst);
810                futures::future::pending::<()>().await;
811            }
812        });
813
814        tasks.shutdown().await;
815        assert_eq!(calls.load(Ordering::SeqCst), 1);
816    }
817
818    #[tokio::test]
819    async fn recorded_iteration_records_success() {
820        let recorded = Arc::new(AtomicUsize::new(0));
821
822        run_recorded_task_iteration(
823            TestTaskName::Cleanup,
824            "cleanup",
825            (),
826            &|()| async { TestOutcome::Succeeded },
827            &TestOutcome::Failed,
828            &{
829                let recorded = recorded.clone();
830                move |(), name, started_at: DateTime<Utc>, finished_at, outcome| {
831                    let recorded = recorded.clone();
832                    async move {
833                        assert_eq!(name, TestTaskName::Cleanup);
834                        assert!(finished_at >= started_at);
835                        assert_eq!(outcome, TestOutcome::Succeeded);
836                        recorded.fetch_add(1, Ordering::SeqCst);
837                    }
838                }
839            },
840        )
841        .await;
842
843        assert_eq!(recorded.load(Ordering::SeqCst), 1);
844    }
845
846    #[tokio::test]
847    async fn recorded_iteration_converts_panic_to_failure_outcome() {
848        let recorded = Arc::new(AtomicUsize::new(0));
849
850        run_recorded_task_iteration(
851            TestTaskName::Cleanup,
852            "cleanup",
853            (),
854            &|()| async {
855                panic!("boom");
856                #[expect(
857                    unreachable_code,
858                    reason = "The test intentionally panics to exercise runtime task failure handling."
859                )]
860                TestOutcome::Succeeded
861            },
862            &TestOutcome::Failed,
863            &{
864                let recorded = recorded.clone();
865                move |(), _name, _started_at, _finished_at, outcome| {
866                    let recorded = recorded.clone();
867                    async move {
868                        assert_eq!(outcome, TestOutcome::Failed("boom".to_string()));
869                        recorded.fetch_add(1, Ordering::SeqCst);
870                    }
871                }
872            },
873        )
874        .await;
875
876        assert_eq!(recorded.load(Ordering::SeqCst), 1);
877    }
878
879    #[tokio::test]
880    async fn pre_cancelled_shutdown_token_skips_periodic_startup_iteration() {
881        let shutdown_token = CancellationToken::new();
882        let calls = Arc::new(AtomicUsize::new(0));
883        shutdown_token.cancel();
884        let task_fn = {
885            let calls = calls.clone();
886            move |()| {
887                let calls = calls.clone();
888                async move {
889                    calls.fetch_add(1, Ordering::SeqCst);
890                    TestOutcome::Succeeded
891                }
892            }
893        };
894
895        run_periodic_task(PeriodicTask {
896            name: TestTaskName::Cleanup,
897            task_name: "cleanup",
898            interval_fn: test_interval,
899            jitter_cap: None,
900            shutdown_token,
901            state: (),
902            hooks: RecordedTaskHooks::new(
903                task_fn,
904                TestOutcome::Failed,
905                |(), _name, _started_at, _finished_at, _outcome| async {},
906            ),
907        })
908        .await;
909
910        assert_eq!(calls.load(Ordering::SeqCst), 0);
911    }
912
913    #[tokio::test]
914    async fn dispatch_worker_runs_startup_and_wakeup_iterations() {
915        #[derive(Clone)]
916        struct State {
917            notify: Arc<Notify>,
918            calls: Arc<AtomicUsize>,
919        }
920
921        let shutdown_token = CancellationToken::new();
922        let state = State {
923            notify: Arc::new(Notify::new()),
924            calls: Arc::new(AtomicUsize::new(0)),
925        };
926        let calls = state.calls.clone();
927
928        let worker = tokio::spawn(run_dispatch_worker(
929            "dispatch",
930            shutdown_token.clone(),
931            state.clone(),
932            |_| Duration::from_mins(1),
933            |_| Duration::from_mins(2),
934            |state: State| async move {
935                state.notify.notified().await;
936            },
937            |state: State, _shutdown| async move {
938                state.calls.fetch_add(1, Ordering::SeqCst);
939                BackgroundTaskDispatchIteration::idle()
940            },
941        ));
942
943        while calls.load(Ordering::SeqCst) == 0 {
944            tokio::task::yield_now().await;
945        }
946        state.notify.notify_one();
947        while calls.load(Ordering::SeqCst) < 2 {
948            tokio::task::yield_now().await;
949        }
950
951        shutdown_token.cancel();
952        worker.await.expect("dispatch worker should stop cleanly");
953        assert_eq!(calls.load(Ordering::SeqCst), 2);
954    }
955
956    #[test]
957    fn background_task_dispatch_zero_base_interval_uses_minimum_delay() {
958        let base = Duration::ZERO;
959        let max = Duration::from_secs(30);
960        let mut backoff = BackgroundTaskDispatchBackoff::new(base, max);
961
962        assert_eq!(backoff.sleep_duration(base, max), Duration::from_secs(1));
963
964        backoff.record_iteration(
965            BackgroundTaskDispatchTrigger::Timer,
966            BackgroundTaskDispatchIteration::idle(),
967            base,
968            max,
969        );
970        assert_eq!(backoff.sleep_duration(base, max), Duration::from_secs(2));
971    }
972
973    #[test]
974    fn background_task_dispatch_backoff_grows_on_idle_and_caps() {
975        let base = Duration::from_secs(5);
976        let max = Duration::from_secs(30);
977        let mut backoff = BackgroundTaskDispatchBackoff::new(base, max);
978
979        assert_eq!(backoff.sleep_duration(base, max), base);
980
981        backoff.record_iteration(
982            BackgroundTaskDispatchTrigger::Timer,
983            BackgroundTaskDispatchIteration::idle(),
984            base,
985            max,
986        );
987        assert_eq!(backoff.sleep_duration(base, max), Duration::from_secs(10));
988
989        backoff.record_iteration(
990            BackgroundTaskDispatchTrigger::Timer,
991            BackgroundTaskDispatchIteration::idle(),
992            base,
993            max,
994        );
995        assert_eq!(backoff.sleep_duration(base, max), Duration::from_secs(20));
996
997        backoff.record_iteration(
998            BackgroundTaskDispatchTrigger::Timer,
999            BackgroundTaskDispatchIteration::idle(),
1000            base,
1001            max,
1002        );
1003        assert_eq!(backoff.sleep_duration(base, max), max);
1004
1005        backoff.record_iteration(
1006            BackgroundTaskDispatchTrigger::Timer,
1007            BackgroundTaskDispatchIteration::idle(),
1008            base,
1009            max,
1010        );
1011        assert_eq!(backoff.sleep_duration(base, max), max);
1012    }
1013
1014    #[test]
1015    fn background_task_dispatch_backoff_resets_on_wakeup_and_activity() {
1016        let base = Duration::from_secs(5);
1017        let max = Duration::from_mins(1);
1018        let mut backoff = BackgroundTaskDispatchBackoff::new(base, max);
1019
1020        backoff.record_iteration(
1021            BackgroundTaskDispatchTrigger::Timer,
1022            BackgroundTaskDispatchIteration::idle(),
1023            base,
1024            max,
1025        );
1026        backoff.record_iteration(
1027            BackgroundTaskDispatchTrigger::Timer,
1028            BackgroundTaskDispatchIteration::idle(),
1029            base,
1030            max,
1031        );
1032        assert_eq!(backoff.sleep_duration(base, max), Duration::from_secs(20));
1033
1034        backoff.record_iteration(
1035            BackgroundTaskDispatchTrigger::Wakeup,
1036            BackgroundTaskDispatchIteration::idle(),
1037            base,
1038            max,
1039        );
1040        assert_eq!(backoff.sleep_duration(base, max), base);
1041
1042        backoff.record_iteration(
1043            BackgroundTaskDispatchTrigger::Timer,
1044            BackgroundTaskDispatchIteration::idle(),
1045            base,
1046            max,
1047        );
1048        assert_eq!(backoff.sleep_duration(base, max), Duration::from_secs(10));
1049
1050        backoff.record_iteration(
1051            BackgroundTaskDispatchTrigger::Timer,
1052            BackgroundTaskDispatchIteration::active(),
1053            base,
1054            max,
1055        );
1056        assert_eq!(backoff.sleep_duration(base, max), base);
1057    }
1058
1059    #[test]
1060    fn background_task_dispatch_backoff_never_polls_faster_than_normal_after_error() {
1061        let base = Duration::from_secs(30);
1062        let max = Duration::from_mins(2);
1063        let mut backoff = BackgroundTaskDispatchBackoff::new(base, max);
1064
1065        backoff.record_iteration(
1066            BackgroundTaskDispatchTrigger::Timer,
1067            BackgroundTaskDispatchIteration::failed(),
1068            base,
1069            max,
1070        );
1071        assert_eq!(backoff.sleep_duration(base, max), base);
1072
1073        let short_base = Duration::from_secs(1);
1074        let mut short_backoff = BackgroundTaskDispatchBackoff::new(short_base, max);
1075        short_backoff.record_iteration(
1076            BackgroundTaskDispatchTrigger::Timer,
1077            BackgroundTaskDispatchIteration::failed(),
1078            short_base,
1079            max,
1080        );
1081        assert_eq!(
1082            short_backoff.sleep_duration(short_base, max),
1083            BACKGROUND_TASK_DISPATCH_ERROR_BACKOFF_CAP
1084        );
1085
1086        backoff.record_iteration(
1087            BackgroundTaskDispatchTrigger::Timer,
1088            BackgroundTaskDispatchIteration::idle(),
1089            base,
1090            max,
1091        );
1092        assert_eq!(backoff.sleep_duration(base, max), Duration::from_mins(1));
1093    }
1094}
aster_forge_tasks/runtime.rs

aster_forge_tasks/
runtime.rs
