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

aster_forge_tasks/
runtime.rs
