aster_forge_tasks/
schedule.rs

1//! Scheduled runtime task catalog and runner primitives.
2//!
3//! A scheduled task is a product-owned runtime job with a stable name and interval. Forge keeps
4//! the reusable coordination contract here: products register catalog entries, a store atomically
5//! claims due firings, and the runner records one panic-protected execution before advancing the
6//! next due timestamp. Concrete persistence is supplied by another crate, typically
7//! `aster_forge_db`.
8
9use std::future::Future;
10use std::marker::PhantomData;
11use std::time::Duration;
12
13use chrono::{DateTime, Utc};
14use futures::FutureExt;
15use tokio_util::sync::CancellationToken;
16use tracing::Instrument;
17
18use crate::runtime::panic_payload_message;
19use crate::{
20    BackgroundTasks, RecordedTaskHooks, RegisteredRuntimeTaskKind, periodic_sleep_duration,
21};
22
23/// One scheduled runtime task entry registered by a product runtime.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct ScheduledTaskCatalogEntry<'a> {
26    /// Product namespace.
27    pub namespace: &'a str,
28    /// Stable task wire name.
29    pub task_name: &'a str,
30    /// Operator-facing display name.
31    pub display_name: &'a str,
32    /// First due timestamp used when inserting a new catalog row.
33    pub first_run_at: DateTime<Utc>,
34}
35
36/// Request to atomically claim one due scheduled task firing.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct ScheduledTaskClaimRequest<'a> {
39    /// Product namespace.
40    pub namespace: &'a str,
41    /// Stable task wire name.
42    pub task_name: &'a str,
43    /// Process-unique runtime owner id.
44    pub owner_id: &'a str,
45    /// Current timestamp.
46    pub now: DateTime<Utc>,
47    /// Claim TTL. Another runtime may reclaim after this duration.
48    pub claim_ttl: Duration,
49}
50
51/// Claimed scheduled task firing.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct ScheduledTaskClaim {
54    /// Stable row identifier.
55    pub task_id: String,
56    /// Product namespace.
57    pub namespace: String,
58    /// Stable task wire name.
59    pub task_name: String,
60    /// Runtime owner id that owns this claim.
61    pub owner_id: String,
62    /// Due timestamp that was claimed.
63    pub scheduled_at: DateTime<Utc>,
64    /// Claim acquisition timestamp.
65    pub claimed_at: DateTime<Utc>,
66    /// Claim expiry timestamp.
67    pub claim_expires_at: DateTime<Utc>,
68}
69
70/// Completion update for a claimed scheduled task.
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub struct ScheduledTaskCompletion {
73    /// Claimed firing to complete.
74    pub claim: ScheduledTaskClaim,
75    /// Runtime completion timestamp.
76    pub finished_at: DateTime<Utc>,
77    /// Next due timestamp after this completion.
78    pub next_run_at: DateTime<Utc>,
79}
80
81/// Renewal update for a claim whose task body is still running.
82///
83/// The store must apply the same ownership predicate as completion (task id,
84/// owner id, and claim acquisition timestamp), so a renewal can never revive a
85/// claim another runtime has already reclaimed.
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87pub struct ScheduledTaskClaimRenewal<'a> {
88    /// Owned claim to renew.
89    pub claim: &'a ScheduledTaskClaim,
90    /// Renewal timestamp.
91    pub now: DateTime<Utc>,
92    /// Fresh claim TTL applied from `now`.
93    pub claim_ttl: Duration,
94}
95
96/// Persistence contract used by scheduled task runners.
97#[async_trait::async_trait]
98pub trait ScheduledTaskStore: Clone + Send + Sync + 'static {
99    /// Store error type.
100    type Error: std::fmt::Display + Send + Sync + 'static;
101
102    /// Ensures one scheduled task is present in the catalog.
103    async fn ensure_scheduled_task(
104        &self,
105        entry: ScheduledTaskCatalogEntry<'_>,
106    ) -> std::result::Result<(), Self::Error>;
107
108    /// Attempts to claim one due scheduled task firing.
109    async fn claim_scheduled_task(
110        &self,
111        request: ScheduledTaskClaimRequest<'_>,
112    ) -> std::result::Result<Option<ScheduledTaskClaim>, Self::Error>;
113
114    /// Renews an owned claim while the task body is still running.
115    ///
116    /// Returning `Ok(false)` means the ownership predicate did not match, so
117    /// the worker must treat the claim as lost and stop renewing. Returning
118    /// `Err(_)` is treated as transient and retried on the next renewal tick.
119    async fn renew_scheduled_task_claim(
120        &self,
121        renewal: ScheduledTaskClaimRenewal<'_>,
122    ) -> std::result::Result<bool, Self::Error>;
123
124    /// Completes a claimed firing and advances the next due timestamp.
125    async fn complete_scheduled_task(
126        &self,
127        completion: ScheduledTaskCompletion,
128    ) -> std::result::Result<bool, Self::Error>;
129}
130
131/// Configuration for one scheduled periodic runtime task worker.
132pub struct ScheduledPeriodicTask<Name, State, Store, IntervalFn, TaskFn, PanicFn, RecordFn> {
133    /// Product task identifier.
134    pub name: Name,
135    /// Product namespace.
136    pub namespace: &'static str,
137    /// Stable task wire name.
138    pub task_name: &'static str,
139    /// Operator-facing display name.
140    pub display_name: &'static str,
141    /// Process-unique runtime owner id.
142    pub owner_id: String,
143    /// Claim TTL used to recover from crashed workers.
144    pub claim_ttl: Duration,
145    /// Reads the latest product-configured interval.
146    pub interval_fn: IntervalFn,
147    /// Optional upper bound for positive jitter.
148    pub jitter_cap: Option<Duration>,
149    /// Shared shutdown token.
150    pub shutdown_token: CancellationToken,
151    /// Product runtime state passed to callbacks.
152    pub state: State,
153    /// Scheduled task store.
154    pub store: Store,
155    /// Product callbacks for execution, panic conversion, and recording.
156    pub hooks: RecordedTaskHooks<TaskFn, PanicFn, RecordFn>,
157}
158
159/// Configuration for a leased group of scheduled runtime tasks.
160///
161/// This is the high-level entrypoint for multi-instance Aster services. Forge
162/// generates the process owner id, supervises the runtime lease, creates the
163/// lease-scoped [`BackgroundTasks`] group, and wires every declared scheduled
164/// task into the shared catalog store. Product code only declares singleton
165/// workers and scheduled task bodies through [`ScheduledRuntimeTaskGroup`].
166#[derive(Clone)]
167pub struct LeasedScheduledRuntimeConfig<
168    Name,
169    Outcome,
170    State,
171    LeaseStore,
172    ScheduleStore,
173    PanicFn,
174    RecordFn,
175> {
176    namespace: &'static str,
177    lease_id: String,
178    lease_store: LeaseStore,
179    schedule_store: ScheduleStore,
180    claim_ttl: Duration,
181    lease_ttl: Duration,
182    lease_renew_interval: Duration,
183    lease_standby_retry_interval: Duration,
184    state: State,
185    panic_outcome: PanicFn,
186    record_outcome: RecordFn,
187    _name: PhantomData<fn() -> Name>,
188    _outcome: PhantomData<fn() -> Outcome>,
189}
190
191impl<Name, Outcome, State, LeaseStore, ScheduleStore, PanicFn, RecordFn>
192    LeasedScheduledRuntimeConfig<Name, Outcome, State, LeaseStore, ScheduleStore, PanicFn, RecordFn>
193{
194    /// Creates configuration for one leased scheduled runtime task group.
195    pub fn new<RecordFut>(
196        namespace: &'static str,
197        lease_id: impl Into<String>,
198        lease_store: LeaseStore,
199        schedule_store: ScheduleStore,
200        state: State,
201        panic_outcome: PanicFn,
202        record_outcome: RecordFn,
203    ) -> Self
204    where
205        PanicFn: Fn(String) -> Outcome,
206        RecordFn:
207            Fn(State, Name, ScheduledTaskClaim, DateTime<Utc>, DateTime<Utc>, Outcome) -> RecordFut,
208        RecordFut: Future<Output = ()> + Send + 'static,
209    {
210        Self {
211            namespace,
212            lease_id: lease_id.into(),
213            lease_store,
214            schedule_store,
215            claim_ttl: Duration::from_secs(120),
216            lease_ttl: aster_forge_runtime::DEFAULT_RUNTIME_LEASE_TTL,
217            lease_renew_interval: Duration::from_secs(10),
218            lease_standby_retry_interval: aster_forge_runtime::DEFAULT_RUNTIME_LEASE_RETRY_INTERVAL,
219            state,
220            panic_outcome,
221            record_outcome,
222            _name: PhantomData,
223            _outcome: PhantomData,
224        }
225    }
226
227    /// Sets the scheduled task claim TTL.
228    pub const fn claim_ttl(mut self, claim_ttl: Duration) -> Self {
229        self.claim_ttl = claim_ttl;
230        self
231    }
232
233    /// Sets the runtime lease TTL.
234    pub const fn lease_ttl(mut self, lease_ttl: Duration) -> Self {
235        self.lease_ttl = lease_ttl;
236        self
237    }
238
239    /// Sets the runtime lease renewal interval for the active owner.
240    pub const fn lease_renew_interval(mut self, lease_renew_interval: Duration) -> Self {
241        self.lease_renew_interval = lease_renew_interval;
242        self
243    }
244
245    /// Sets the standby retry interval while another process owns the lease.
246    pub const fn lease_standby_retry_interval(
247        mut self,
248        lease_standby_retry_interval: Duration,
249    ) -> Self {
250        self.lease_standby_retry_interval = lease_standby_retry_interval;
251        self
252    }
253
254    /// Runs this configured leased scheduled runtime group until shutdown.
255    ///
256    /// Prefer this method at product entrypoints because it keeps the call
257    /// shaped like a component declaration: configure shared resources once,
258    /// then declare workers and scheduled tasks in the closure.
259    pub async fn run<ConfigureFn>(self, shutdown_token: CancellationToken, configure: ConfigureFn)
260    where
261        Name: RegisteredRuntimeTaskKind + Send + Sync + 'static,
262        State: Clone + Send + Sync + 'static,
263        LeaseStore: aster_forge_runtime::RuntimeLeaseStore,
264        ScheduleStore: ScheduledTaskStore,
265        ConfigureFn: for<'a> FnMut(
266                &mut ScheduledRuntimeTaskGroup<
267                    'a,
268                    Name,
269                    State,
270                    ScheduleStore,
271                    PanicFn,
272                    RecordFn,
273                    Outcome,
274                >,
275            ) + Send
276            + 'static,
277        PanicFn: Clone + Fn(String) -> Outcome + Send + Sync + 'static,
278        RecordFn: Clone + Send + Sync + 'static,
279        Outcome: Send + 'static,
280    {
281        run_leased_scheduled_runtime_tasks(self, shutdown_token, configure).await;
282    }
283
284    fn into_parts(
285        self,
286    ) -> LeasedScheduledRuntimeParts<State, LeaseStore, ScheduleStore, PanicFn, RecordFn> {
287        let owner_id = aster_forge_runtime::new_runtime_lease_owner_id();
288        let lease_config =
289            aster_forge_runtime::RuntimeLeaseConfig::new(self.lease_id, owner_id.clone())
290                .ttl(self.lease_ttl)
291                .renew_interval(self.lease_renew_interval)
292                .standby_retry_interval(self.lease_standby_retry_interval);
293        LeasedScheduledRuntimeParts {
294            namespace: self.namespace,
295            owner_id,
296            lease_store: self.lease_store,
297            schedule_store: self.schedule_store,
298            claim_ttl: self.claim_ttl,
299            state: self.state,
300            panic_outcome: self.panic_outcome,
301            record_outcome: self.record_outcome,
302            lease_config,
303        }
304    }
305}
306
307struct LeasedScheduledRuntimeParts<State, LeaseStore, ScheduleStore, PanicFn, RecordFn> {
308    namespace: &'static str,
309    owner_id: String,
310    lease_store: LeaseStore,
311    schedule_store: ScheduleStore,
312    claim_ttl: Duration,
313    state: State,
314    panic_outcome: PanicFn,
315    record_outcome: RecordFn,
316    lease_config: aster_forge_runtime::RuntimeLeaseConfig,
317}
318
319/// Lease-scoped task group used by product registration closures.
320///
321/// A value of this type exists only while Forge is building the worker group
322/// for one lease acquisition. Use [`Self::worker`] for singleton workers that
323/// should run only on the active owner, and [`Self::scheduled`] for tasks that
324/// should additionally coordinate each firing through the scheduled task
325/// catalog.
326pub struct ScheduledRuntimeTaskGroup<'a, Name, State, Store, PanicFn, RecordFn, Outcome> {
327    tasks: &'a mut BackgroundTasks,
328    namespace: &'static str,
329    owner_id: String,
330    claim_ttl: Duration,
331    shutdown_token: CancellationToken,
332    state: State,
333    store: Store,
334    panic_outcome: PanicFn,
335    record_outcome: RecordFn,
336    _name: std::marker::PhantomData<Name>,
337    _outcome: std::marker::PhantomData<Outcome>,
338}
339
340impl<'a, Name, State, Store, PanicFn, RecordFn, Outcome>
341    ScheduledRuntimeTaskGroup<'a, Name, State, Store, PanicFn, RecordFn, Outcome>
342where
343    State: Clone + Send + Sync + 'static,
344{
345    /// Spawns one lease-scoped singleton worker into this group.
346    pub fn worker<WorkerFn, WorkerFut>(&mut self, worker: WorkerFn)
347    where
348        WorkerFn: FnOnce(CancellationToken, State) -> WorkerFut,
349        WorkerFut: Future<Output = ()> + Send + 'static,
350    {
351        self.tasks
352            .push(worker(self.shutdown_token.clone(), self.state.clone()));
353    }
354
355    /// Returns a clone of the lease-scoped shutdown token.
356    pub fn shutdown_token(&self) -> CancellationToken {
357        self.shutdown_token.clone()
358    }
359
360    /// Returns a clone of the product runtime state.
361    pub fn state(&self) -> State {
362        self.state.clone()
363    }
364}
365
366impl<'a, Name, State, Store, PanicFn, RecordFn, Outcome>
367    ScheduledRuntimeTaskGroup<'a, Name, State, Store, PanicFn, RecordFn, Outcome>
368where
369    Name: RegisteredRuntimeTaskKind + Send + Sync + 'static,
370    State: Clone + Send + Sync + 'static,
371    Store: ScheduledTaskStore,
372    PanicFn: Clone + Fn(String) -> Outcome + Send + Sync + 'static,
373    RecordFn: Clone + Send + Sync + 'static,
374    Outcome: Send + 'static,
375{
376    /// Registers one scheduled runtime task in the lease-scoped worker group.
377    pub fn scheduled<IntervalFn, TaskFn, TaskFut, RecordFut>(
378        &mut self,
379        name: Name,
380        interval_fn: IntervalFn,
381        jitter_cap: Option<Duration>,
382        task_fn: TaskFn,
383    ) where
384        IntervalFn: Fn(&State) -> Duration + Send + Sync + 'static,
385        TaskFn: Fn(State) -> TaskFut + Send + Sync + 'static,
386        TaskFut: Future<Output = Outcome> + Send + 'static,
387        RecordFn:
388            Fn(State, Name, ScheduledTaskClaim, DateTime<Utc>, DateTime<Utc>, Outcome) -> RecordFut,
389        RecordFut: Future<Output = ()> + Send + 'static,
390    {
391        self.tasks
392            .push(run_scheduled_periodic_task(ScheduledPeriodicTask {
393                name,
394                namespace: self.namespace,
395                task_name: name.as_str(),
396                display_name: name.display_name(),
397                owner_id: self.owner_id.clone(),
398                claim_ttl: self.claim_ttl,
399                interval_fn,
400                jitter_cap,
401                shutdown_token: self.shutdown_token.clone(),
402                state: self.state.clone(),
403                store: self.store.clone(),
404                hooks: RecordedTaskHooks::new(
405                    task_fn,
406                    self.panic_outcome.clone(),
407                    self.record_outcome.clone(),
408                ),
409            }));
410    }
411}
412
413/// Runs a lease-supervised scheduled runtime task group until shutdown.
414///
415/// Forge owns the lifecycle glue: owner id generation, runtime lease
416/// supervision, lease-scoped shutdown token creation, scheduled task catalog
417/// registration, and graceful worker shutdown. Product code supplies the
418/// runtime config and a closure that declares workers and scheduled tasks.
419async fn run_leased_scheduled_runtime_tasks<
420    Name,
421    Outcome,
422    State,
423    LeaseStore,
424    ScheduleStore,
425    ConfigureFn,
426    PanicFn,
427    RecordFn,
428>(
429    config: LeasedScheduledRuntimeConfig<
430        Name,
431        Outcome,
432        State,
433        LeaseStore,
434        ScheduleStore,
435        PanicFn,
436        RecordFn,
437    >,
438    shutdown_token: CancellationToken,
439    mut configure: ConfigureFn,
440) where
441    Name: RegisteredRuntimeTaskKind + Send + Sync + 'static,
442    State: Clone + Send + Sync + 'static,
443    LeaseStore: aster_forge_runtime::RuntimeLeaseStore,
444    ScheduleStore: ScheduledTaskStore,
445    ConfigureFn: for<'a> FnMut(
446            &mut ScheduledRuntimeTaskGroup<
447                'a,
448                Name,
449                State,
450                ScheduleStore,
451                PanicFn,
452                RecordFn,
453                Outcome,
454            >,
455        ) + Send
456        + 'static,
457    PanicFn: Clone + Fn(String) -> Outcome + Send + Sync + 'static,
458    RecordFn: Clone + Send + Sync + 'static,
459    Outcome: Send + 'static,
460{
461    let parts = config.into_parts();
462    let LeasedScheduledRuntimeParts {
463        namespace,
464        owner_id,
465        lease_store,
466        schedule_store,
467        claim_ttl,
468        state,
469        panic_outcome,
470        record_outcome,
471        lease_config,
472    } = parts;
473
474    aster_forge_runtime::run_runtime_lease_supervisor(
475        lease_store,
476        lease_config,
477        shutdown_token,
478        move |leased_shutdown_token| {
479            let mut tasks = BackgroundTasks::with_shutdown_token(leased_shutdown_token);
480            let group_shutdown_token = tasks.shutdown_token();
481            let mut group = ScheduledRuntimeTaskGroup {
482                tasks: &mut tasks,
483                namespace,
484                owner_id: owner_id.clone(),
485                claim_ttl,
486                shutdown_token: group_shutdown_token,
487                state: state.clone(),
488                store: schedule_store.clone(),
489                panic_outcome: panic_outcome.clone(),
490                record_outcome: record_outcome.clone(),
491                _name: std::marker::PhantomData,
492                _outcome: std::marker::PhantomData,
493            };
494            configure(&mut group);
495            tasks
496        },
497        |background_tasks| async move {
498            background_tasks.shutdown().await;
499        },
500    )
501    .await;
502}
503
504/// Runs a scheduled periodic task until shutdown.
505///
506/// Unlike [`crate::run_periodic_task`], this runner first claims a due catalog row. If the row is
507/// not due, or another process owns a fresh claim, the worker skips that iteration. Successful and
508/// failed task outcomes both complete the claim and advance `next_run_at`; crashes and process
509/// exits before completion are recovered by claim expiry.
510///
511/// While the task body runs, a renewal loop extends the claim at
512/// [`scheduled_claim_renew_interval`] ticks, so a task that outlives `claim_ttl` is not reclaimed
513/// and executed twice by another runtime. Renewal failures never abort the task body: a lost
514/// claim only stops the renewal loop, and completion still guards on ownership.
515pub async fn run_scheduled_periodic_task<
516    Name,
517    State,
518    Store,
519    IntervalFn,
520    TaskFn,
521    TaskFut,
522    PanicFn,
523    RecordFn,
524    RecordFut,
525    Outcome,
526>(
527    task: ScheduledPeriodicTask<Name, State, Store, IntervalFn, TaskFn, PanicFn, RecordFn>,
528) where
529    Name: Copy + Send + 'static,
530    State: Clone + Send + Sync + 'static,
531    Store: ScheduledTaskStore,
532    IntervalFn: Fn(&State) -> Duration + Send + Sync + 'static,
533    TaskFn: Fn(State) -> TaskFut + Send + Sync + 'static,
534    TaskFut: Future<Output = Outcome> + Send + 'static,
535    PanicFn: Fn(String) -> Outcome + Send + Sync + 'static,
536    RecordFn: Fn(State, Name, ScheduledTaskClaim, DateTime<Utc>, DateTime<Utc>, Outcome) -> RecordFut
537        + Send
538        + Sync
539        + 'static,
540    RecordFut: Future<Output = ()> + Send + 'static,
541    Outcome: Send + 'static,
542{
543    if task.shutdown_token.is_cancelled() {
544        return;
545    }
546
547    run_scheduled_periodic_iteration(&task)
548        .instrument(tracing::info_span!("bg_task", task.name = task.task_name))
549        .await;
550
551    loop {
552        let sleep_duration =
553            periodic_sleep_duration((task.interval_fn)(&task.state), task.jitter_cap);
554        tokio::select! {
555            biased;
556            _ = task.shutdown_token.cancelled() => break,
557            _ = tokio::time::sleep(sleep_duration) => {}
558        }
559
560        if task.shutdown_token.is_cancelled() {
561            break;
562        }
563
564        run_scheduled_periodic_iteration(&task)
565            .instrument(tracing::info_span!("bg_task", task.name = task.task_name))
566            .await;
567    }
568}
569
570/// Derives the claim renewal tick from the claim TTL.
571///
572/// Renewing three times per TTL window means two consecutive missed ticks still
573/// leave one renewal before expiry. The floor keeps `tokio::time::interval`
574/// away from a zero period for pathological TTLs.
575pub fn scheduled_claim_renew_interval(claim_ttl: Duration) -> Duration {
576    (claim_ttl / 3).max(Duration::from_millis(10))
577}
578
579/// Renews one owned scheduled task claim until stopped or the claim is lost.
580///
581/// Mirrors the background task heartbeat loop: `Ok(false)` means the ownership
582/// predicate no longer matches (another runtime reclaimed the firing), so the
583/// loop stops; `Err(_)` is logged and retried on the next tick. The loop never
584/// aborts the running task body — completion still guards on ownership.
585pub async fn run_scheduled_claim_renewal_loop<Store>(
586    store: Store,
587    claim: ScheduledTaskClaim,
588    claim_ttl: Duration,
589    interval: Duration,
590    stop_token: CancellationToken,
591) where
592    Store: ScheduledTaskStore,
593{
594    let mut ticker = tokio::time::interval(interval);
595    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
596    ticker.tick().await;
597
598    loop {
599        tokio::select! {
600            _ = stop_token.cancelled() => return,
601            _ = ticker.tick() => {
602                let renewal = ScheduledTaskClaimRenewal {
603                    claim: &claim,
604                    now: Utc::now(),
605                    claim_ttl,
606                };
607                let result = tokio::select! {
608                    _ = stop_token.cancelled() => return,
609                    result = store.renew_scheduled_task_claim(renewal) => result,
610                };
611
612                match result {
613                    Ok(true) => {}
614                    Ok(false) => {
615                        tracing::warn!(
616                            task.name = claim.task_name,
617                            "scheduled task claim lost; stopping claim renewal"
618                        );
619                        return;
620                    }
621                    Err(error) => {
622                        tracing::warn!(
623                            task.name = claim.task_name,
624                            error = %error,
625                            "scheduled task claim renewal failed; retrying next tick"
626                        );
627                    }
628                }
629            }
630        }
631    }
632}
633
634async fn run_scheduled_periodic_iteration<
635    Name,
636    State,
637    Store,
638    IntervalFn,
639    TaskFn,
640    TaskFut,
641    PanicFn,
642    RecordFn,
643    RecordFut,
644    Outcome,
645>(
646    task: &ScheduledPeriodicTask<Name, State, Store, IntervalFn, TaskFn, PanicFn, RecordFn>,
647) where
648    Name: Copy + Send + 'static,
649    State: Clone + Send + Sync + 'static,
650    Store: ScheduledTaskStore,
651    IntervalFn: Fn(&State) -> Duration + Send + Sync + 'static,
652    TaskFn: Fn(State) -> TaskFut + Send + Sync + 'static,
653    TaskFut: Future<Output = Outcome> + Send + 'static,
654    PanicFn: Fn(String) -> Outcome + Send + Sync + 'static,
655    RecordFn: Fn(State, Name, ScheduledTaskClaim, DateTime<Utc>, DateTime<Utc>, Outcome) -> RecordFut
656        + Send
657        + Sync
658        + 'static,
659    RecordFut: Future<Output = ()> + Send + 'static,
660    Outcome: Send + 'static,
661{
662    let now = Utc::now();
663    if let Err(error) = task
664        .store
665        .ensure_scheduled_task(ScheduledTaskCatalogEntry {
666            namespace: task.namespace,
667            task_name: task.task_name,
668            display_name: task.display_name,
669            first_run_at: now,
670        })
671        .await
672    {
673        tracing::warn!(
674            task.name = task.task_name,
675            error = %error,
676            "failed to ensure scheduled task catalog row"
677        );
678        return;
679    }
680
681    let claim = match task
682        .store
683        .claim_scheduled_task(ScheduledTaskClaimRequest {
684            namespace: task.namespace,
685            task_name: task.task_name,
686            owner_id: &task.owner_id,
687            now,
688            claim_ttl: task.claim_ttl,
689        })
690        .await
691    {
692        Ok(Some(claim)) => claim,
693        Ok(None) => return,
694        Err(error) => {
695            tracing::warn!(
696                task.name = task.task_name,
697                error = %error,
698                "failed to claim scheduled task"
699            );
700            return;
701        }
702    };
703
704    let renewal_stop = task.shutdown_token.child_token();
705    let renewal_handle = tokio::spawn(run_scheduled_claim_renewal_loop(
706        task.store.clone(),
707        claim.clone(),
708        task.claim_ttl,
709        scheduled_claim_renew_interval(task.claim_ttl),
710        renewal_stop.clone(),
711    ));
712
713    let started_at = Utc::now();
714    let outcome = match std::panic::AssertUnwindSafe((task.hooks.task_fn)(task.state.clone()))
715        .catch_unwind()
716        .await
717    {
718        Ok(outcome) => outcome,
719        Err(panic) => {
720            let panic_message = panic_payload_message(&panic);
721            tracing::error!(
722                task.name = task.task_name,
723                "scheduled task panicked: {panic_message}"
724            );
725            (task.hooks.panic_outcome)(panic_message)
726        }
727    };
728    let finished_at = Utc::now();
729
730    renewal_stop.cancel();
731    if let Err(error) = renewal_handle.await {
732        tracing::warn!(
733            task.name = task.task_name,
734            error = %error,
735            "scheduled task claim renewal worker stopped unexpectedly"
736        );
737    }
738
739    let record_result = std::panic::AssertUnwindSafe((task.hooks.record_outcome)(
740        task.state.clone(),
741        task.name,
742        claim.clone(),
743        started_at,
744        finished_at,
745        outcome,
746    ))
747    .catch_unwind()
748    .await;
749    if let Err(panic) = record_result {
750        let panic_message = panic_payload_message(&panic);
751        tracing::error!(
752            task.name = task.task_name,
753            "scheduled task outcome recorder panicked: {panic_message}"
754        );
755        return;
756    }
757
758    let Some(next_run_at) = next_scheduled_run_at(finished_at, (task.interval_fn)(&task.state))
759    else {
760        tracing::warn!(
761            task.name = task.task_name,
762            "scheduled task interval overflowed while computing next run"
763        );
764        return;
765    };
766
767    match task
768        .store
769        .complete_scheduled_task(ScheduledTaskCompletion {
770            claim,
771            finished_at,
772            next_run_at,
773        })
774        .await
775    {
776        Ok(true) => {}
777        Ok(false) => {
778            tracing::warn!(
779                task.name = task.task_name,
780                "scheduled task claim was not completed because ownership changed"
781            );
782        }
783        Err(error) => {
784            tracing::warn!(
785                task.name = task.task_name,
786                error = %error,
787                "failed to complete scheduled task claim"
788            );
789        }
790    }
791}
792
793/// Computes the next run timestamp after a completed scheduled task firing.
794pub fn next_scheduled_run_at(
795    finished_at: DateTime<Utc>,
796    interval: Duration,
797) -> Option<DateTime<Utc>> {
798    let interval = chrono::Duration::from_std(interval).ok()?;
799    finished_at.checked_add_signed(interval)
800}
801
802#[cfg(test)]
803mod tests {
804    use std::sync::Arc;
805    use std::sync::Mutex;
806    use std::sync::atomic::{AtomicUsize, Ordering};
807
808    use async_trait::async_trait;
809    use chrono::{TimeZone, Utc};
810    use tokio_util::sync::CancellationToken;
811
812    use super::{
813        LeasedScheduledRuntimeConfig, ScheduledPeriodicTask, ScheduledRuntimeTaskGroup,
814        ScheduledTaskCatalogEntry, ScheduledTaskClaim, ScheduledTaskClaimRequest,
815        ScheduledTaskCompletion, ScheduledTaskStore, next_scheduled_run_at,
816        run_scheduled_claim_renewal_loop, run_scheduled_periodic_task,
817        scheduled_claim_renew_interval,
818    };
819    use crate::{RecordedTaskHooks, RegisteredRuntimeTaskKind};
820
821    #[derive(Clone)]
822    struct MemoryScheduleStore {
823        calls: Arc<AtomicUsize>,
824        completions: Arc<AtomicUsize>,
825        renewals: Arc<AtomicUsize>,
826        renewal_script: Arc<Mutex<std::collections::VecDeque<Result<bool, String>>>>,
827    }
828
829    fn test_interval(_: &()) -> std::time::Duration {
830        std::time::Duration::from_secs(60)
831    }
832
833    #[derive(Clone)]
834    struct AlwaysAcquireLeaseStore {
835        acquired: Arc<AtomicUsize>,
836        released: Arc<AtomicUsize>,
837    }
838
839    impl AlwaysAcquireLeaseStore {
840        fn new() -> Self {
841            Self {
842                acquired: Arc::new(AtomicUsize::new(0)),
843                released: Arc::new(AtomicUsize::new(0)),
844            }
845        }
846    }
847
848    #[async_trait]
849    impl aster_forge_runtime::RuntimeLeaseStore for AlwaysAcquireLeaseStore {
850        type Error = String;
851
852        async fn try_acquire(
853            &self,
854            _claim: aster_forge_runtime::RuntimeLeaseClaim<'_>,
855        ) -> Result<aster_forge_runtime::RuntimeLeaseAcquire, Self::Error> {
856            self.acquired.fetch_add(1, Ordering::SeqCst);
857            Ok(aster_forge_runtime::RuntimeLeaseAcquire::Acquired)
858        }
859
860        async fn renew(
861            &self,
862            _lease_id: &str,
863            _owner_id: &str,
864            _now: chrono::DateTime<Utc>,
865            _expires_at: chrono::DateTime<Utc>,
866        ) -> Result<bool, Self::Error> {
867            Ok(true)
868        }
869
870        async fn release(&self, _lease_id: &str, _owner_id: &str) -> Result<(), Self::Error> {
871            self.released.fetch_add(1, Ordering::SeqCst);
872            Ok(())
873        }
874    }
875
876    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
877    enum TestRuntimeTask {
878        Cleanup,
879    }
880
881    impl RegisteredRuntimeTaskKind for TestRuntimeTask {
882        fn as_str(self) -> &'static str {
883            "cleanup"
884        }
885
886        fn display_name(self) -> &'static str {
887            "Cleanup"
888        }
889
890        fn from_wire_value(value: &str) -> Option<Self> {
891            (value == "cleanup").then_some(Self::Cleanup)
892        }
893    }
894
895    #[async_trait]
896    impl ScheduledTaskStore for MemoryScheduleStore {
897        type Error = String;
898
899        async fn ensure_scheduled_task(
900            &self,
901            entry: ScheduledTaskCatalogEntry<'_>,
902        ) -> Result<(), Self::Error> {
903            assert_eq!(entry.namespace, "aster_test");
904            assert_eq!(entry.task_name, "cleanup");
905            Ok(())
906        }
907
908        async fn claim_scheduled_task(
909            &self,
910            request: ScheduledTaskClaimRequest<'_>,
911        ) -> Result<Option<ScheduledTaskClaim>, Self::Error> {
912            if self.calls.fetch_add(1, Ordering::SeqCst) > 0 {
913                return Ok(None);
914            }
915            Ok(Some(ScheduledTaskClaim {
916                task_id: "aster_test:cleanup".to_string(),
917                namespace: request.namespace.to_string(),
918                task_name: request.task_name.to_string(),
919                owner_id: request.owner_id.to_string(),
920                scheduled_at: request.now,
921                claimed_at: request.now,
922                claim_expires_at: request.now,
923            }))
924        }
925
926        async fn renew_scheduled_task_claim(
927            &self,
928            renewal: super::ScheduledTaskClaimRenewal<'_>,
929        ) -> Result<bool, Self::Error> {
930            assert_eq!(renewal.claim.task_name, "cleanup");
931            assert!(!renewal.claim_ttl.is_zero());
932            self.renewals.fetch_add(1, Ordering::SeqCst);
933            let scripted = self
934                .renewal_script
935                .lock()
936                .expect("renewal script should lock")
937                .pop_front();
938            scripted.unwrap_or(Ok(true))
939        }
940
941        async fn complete_scheduled_task(
942            &self,
943            completion: ScheduledTaskCompletion,
944        ) -> Result<bool, Self::Error> {
945            assert_eq!(completion.claim.task_name, "cleanup");
946            assert!(completion.next_run_at >= completion.finished_at);
947            self.completions.fetch_add(1, Ordering::SeqCst);
948            Ok(true)
949        }
950    }
951
952    fn memory_store() -> MemoryScheduleStore {
953        MemoryScheduleStore {
954            calls: Arc::new(AtomicUsize::new(0)),
955            completions: Arc::new(AtomicUsize::new(0)),
956            renewals: Arc::new(AtomicUsize::new(0)),
957            renewal_script: Arc::new(Mutex::new(std::collections::VecDeque::new())),
958        }
959    }
960
961    #[tokio::test]
962    async fn scheduled_periodic_task_claims_records_and_completes_one_due_run() {
963        let shutdown = CancellationToken::new();
964        let ran = Arc::new(AtomicUsize::new(0));
965        let recorded = Arc::new(AtomicUsize::new(0));
966        let store = memory_store();
967        let completions = store.completions.clone();
968        let ran_for_task = ran.clone();
969        let recorded_for_hook = recorded.clone();
970        let shutdown_for_hook = shutdown.clone();
971
972        run_scheduled_periodic_task(ScheduledPeriodicTask {
973            name: "cleanup",
974            namespace: "aster_test",
975            task_name: "cleanup",
976            display_name: "Cleanup",
977            owner_id: "runtime-a".to_string(),
978            claim_ttl: std::time::Duration::from_secs(30),
979            interval_fn: test_interval,
980            jitter_cap: None,
981            shutdown_token: shutdown.clone(),
982            state: (),
983            store,
984            hooks: RecordedTaskHooks::new(
985                move |()| {
986                    let ran = ran_for_task.clone();
987                    async move {
988                        ran.fetch_add(1, Ordering::SeqCst);
989                        "ok"
990                    }
991                },
992                |_| "panic",
993                move |(),
994                      _name: &str,
995                      claim: ScheduledTaskClaim,
996                      _started_at,
997                      _finished_at,
998                      outcome| {
999                    let recorded = recorded_for_hook.clone();
1000                    let shutdown = shutdown_for_hook.clone();
1001                    async move {
1002                        assert_eq!(claim.task_name, "cleanup");
1003                        assert_eq!(outcome, "ok");
1004                        recorded.fetch_add(1, Ordering::SeqCst);
1005                        shutdown.cancel();
1006                    }
1007                },
1008            ),
1009        })
1010        .await;
1011
1012        assert_eq!(ran.load(Ordering::SeqCst), 1);
1013        assert_eq!(recorded.load(Ordering::SeqCst), 1);
1014        assert_eq!(completions.load(Ordering::SeqCst), 1);
1015    }
1016
1017    #[tokio::test]
1018    async fn scheduled_periodic_task_records_panic_outcome_and_completes_claim() {
1019        let shutdown = CancellationToken::new();
1020        let store = memory_store();
1021        let completions = store.completions.clone();
1022        let recorded = Arc::new(Mutex::new(Vec::new()));
1023        let recorded_for_hook = recorded.clone();
1024        let shutdown_for_hook = shutdown.clone();
1025
1026        run_scheduled_periodic_task(ScheduledPeriodicTask {
1027            name: "cleanup",
1028            namespace: "aster_test",
1029            task_name: "cleanup",
1030            display_name: "Cleanup",
1031            owner_id: "runtime-a".to_string(),
1032            claim_ttl: std::time::Duration::from_secs(30),
1033            interval_fn: test_interval,
1034            jitter_cap: None,
1035            shutdown_token: shutdown.clone(),
1036            state: (),
1037            store,
1038            hooks: RecordedTaskHooks::new(
1039                move |()| async move {
1040                    panic!("scheduled body failed");
1041                    #[allow(unreachable_code)]
1042                    "ok".to_string()
1043                },
1044                |message| format!("panic:{message}"),
1045                move |(), _name, _claim, _started_at, _finished_at, outcome| {
1046                    let recorded = recorded_for_hook.clone();
1047                    let shutdown = shutdown_for_hook.clone();
1048                    async move {
1049                        recorded
1050                            .lock()
1051                            .expect("recorded outcomes should lock")
1052                            .push(outcome);
1053                        shutdown.cancel();
1054                    }
1055                },
1056            ),
1057        })
1058        .await;
1059
1060        assert_eq!(
1061            recorded
1062                .lock()
1063                .expect("recorded outcomes should lock")
1064                .as_slice(),
1065            ["panic:scheduled body failed"]
1066        );
1067        assert_eq!(completions.load(Ordering::SeqCst), 1);
1068    }
1069
1070    #[tokio::test]
1071    async fn scheduled_periodic_task_does_not_complete_claim_when_recorder_panics() {
1072        let shutdown = CancellationToken::new();
1073        let store = memory_store();
1074        let completions = store.completions.clone();
1075
1076        run_scheduled_periodic_task(ScheduledPeriodicTask {
1077            name: "cleanup",
1078            namespace: "aster_test",
1079            task_name: "cleanup",
1080            display_name: "Cleanup",
1081            owner_id: "runtime-a".to_string(),
1082            claim_ttl: std::time::Duration::from_secs(30),
1083            interval_fn: test_interval,
1084            jitter_cap: None,
1085            shutdown_token: shutdown.clone(),
1086            state: (),
1087            store,
1088            hooks: RecordedTaskHooks::new(
1089                move |()| {
1090                    let shutdown = shutdown.clone();
1091                    async move {
1092                        shutdown.cancel();
1093                        "ok"
1094                    }
1095                },
1096                |_| "panic",
1097                move |(), _name, _claim, _started_at, _finished_at, _outcome| async move {
1098                    panic!("record failed");
1099                },
1100            ),
1101        })
1102        .await;
1103
1104        assert_eq!(completions.load(Ordering::SeqCst), 0);
1105    }
1106
1107    #[tokio::test]
1108    async fn leased_scheduled_runtime_group_runs_worker_and_scheduled_task() {
1109        let lease_store = AlwaysAcquireLeaseStore::new();
1110        let acquired = lease_store.acquired.clone();
1111        let released = lease_store.released.clone();
1112        let schedule_store = memory_store();
1113        let completions = schedule_store.completions.clone();
1114        let worker_runs = Arc::new(AtomicUsize::new(0));
1115        let scheduled_runs = Arc::new(AtomicUsize::new(0));
1116        let recorded_runs = Arc::new(AtomicUsize::new(0));
1117        let shutdown = CancellationToken::new();
1118        let config = LeasedScheduledRuntimeConfig::new(
1119            "aster_test",
1120            "aster_test.background",
1121            lease_store,
1122            schedule_store,
1123            (),
1124            |_| "panic",
1125            {
1126                let recorded_runs = recorded_runs.clone();
1127                let shutdown = shutdown.clone();
1128                move |(), _name, claim: ScheduledTaskClaim, _started_at, _finished_at, outcome| {
1129                    let recorded_runs = recorded_runs.clone();
1130                    let shutdown = shutdown.clone();
1131                    async move {
1132                        assert_eq!(claim.task_name, "cleanup");
1133                        assert_eq!(outcome, "ok");
1134                        recorded_runs.fetch_add(1, Ordering::SeqCst);
1135                        shutdown.cancel();
1136                    }
1137                }
1138            },
1139        )
1140        .claim_ttl(std::time::Duration::from_secs(30))
1141        .lease_ttl(std::time::Duration::from_secs(30))
1142        .lease_renew_interval(std::time::Duration::from_secs(10))
1143        .lease_standby_retry_interval(std::time::Duration::from_secs(5));
1144        let worker_runs_for_group = worker_runs.clone();
1145        let scheduled_runs_for_group = scheduled_runs.clone();
1146
1147        config
1148            .run(
1149                shutdown.clone(),
1150                move |group: &mut ScheduledRuntimeTaskGroup<
1151                    '_,
1152                    TestRuntimeTask,
1153                    (),
1154                    _,
1155                    _,
1156                    _,
1157                    &'static str,
1158                >| {
1159                    let worker_runs = worker_runs_for_group.clone();
1160                    group.worker(move |shutdown_token, ()| async move {
1161                        worker_runs.fetch_add(1, Ordering::SeqCst);
1162                        shutdown_token.cancelled().await;
1163                    });
1164                    let scheduled_runs = scheduled_runs_for_group.clone();
1165                    group.scheduled(TestRuntimeTask::Cleanup, test_interval, None, move |()| {
1166                        let scheduled_runs = scheduled_runs.clone();
1167                        async move {
1168                            scheduled_runs.fetch_add(1, Ordering::SeqCst);
1169                            "ok"
1170                        }
1171                    });
1172                },
1173            )
1174            .await;
1175
1176        assert_eq!(acquired.load(Ordering::SeqCst), 1);
1177        assert_eq!(released.load(Ordering::SeqCst), 1);
1178        assert_eq!(worker_runs.load(Ordering::SeqCst), 1);
1179        assert_eq!(scheduled_runs.load(Ordering::SeqCst), 1);
1180        assert_eq!(recorded_runs.load(Ordering::SeqCst), 1);
1181        assert_eq!(completions.load(Ordering::SeqCst), 1);
1182    }
1183
1184    #[test]
1185    fn next_scheduled_run_at_adds_interval() {
1186        let finished_at = Utc.with_ymd_and_hms(2026, 6, 26, 1, 2, 3).unwrap();
1187        assert_eq!(
1188            next_scheduled_run_at(finished_at, std::time::Duration::from_secs(60)),
1189            Some(Utc.with_ymd_and_hms(2026, 6, 26, 1, 3, 3).unwrap())
1190        );
1191    }
1192
1193    fn test_claim() -> ScheduledTaskClaim {
1194        let now = Utc.with_ymd_and_hms(2026, 6, 26, 1, 0, 0).unwrap();
1195        ScheduledTaskClaim {
1196            task_id: "aster_test:cleanup".to_string(),
1197            namespace: "aster_test".to_string(),
1198            task_name: "cleanup".to_string(),
1199            owner_id: "runtime-a".to_string(),
1200            scheduled_at: now,
1201            claimed_at: now,
1202            claim_expires_at: now,
1203        }
1204    }
1205
1206    #[test]
1207    fn claim_renew_interval_is_one_third_of_ttl_with_floor() {
1208        assert_eq!(
1209            scheduled_claim_renew_interval(std::time::Duration::from_secs(120)),
1210            std::time::Duration::from_secs(40)
1211        );
1212        assert_eq!(
1213            scheduled_claim_renew_interval(std::time::Duration::from_millis(30)),
1214            std::time::Duration::from_millis(10)
1215        );
1216        // A zero TTL never reaches the store, but the interval must not panic.
1217        assert_eq!(
1218            scheduled_claim_renew_interval(std::time::Duration::ZERO),
1219            std::time::Duration::from_millis(10)
1220        );
1221    }
1222
1223    #[tokio::test]
1224    async fn claim_renewal_loop_renews_until_stopped() {
1225        let store = memory_store();
1226        let renewals = store.renewals.clone();
1227        let stop = CancellationToken::new();
1228
1229        let handle = tokio::spawn(run_scheduled_claim_renewal_loop(
1230            store,
1231            test_claim(),
1232            std::time::Duration::from_secs(30),
1233            std::time::Duration::from_millis(1),
1234            stop.clone(),
1235        ));
1236
1237        while renewals.load(Ordering::SeqCst) == 0 {
1238            tokio::task::yield_now().await;
1239        }
1240        stop.cancel();
1241        handle.await.expect("renewal loop should stop cleanly");
1242        assert!(renewals.load(Ordering::SeqCst) >= 1);
1243    }
1244
1245    #[tokio::test]
1246    async fn claim_renewal_loop_stops_when_claim_is_lost() {
1247        let store = memory_store();
1248        store
1249            .renewal_script
1250            .lock()
1251            .expect("renewal script should lock")
1252            .push_back(Ok(false));
1253        let renewals = store.renewals.clone();
1254
1255        // Ok(false) means the ownership predicate no longer matches: the loop
1256        // must stop on its own instead of hammering a reclaimed row.
1257        run_scheduled_claim_renewal_loop(
1258            store,
1259            test_claim(),
1260            std::time::Duration::from_secs(30),
1261            std::time::Duration::from_millis(1),
1262            CancellationToken::new(),
1263        )
1264        .await;
1265
1266        assert_eq!(renewals.load(Ordering::SeqCst), 1);
1267    }
1268
1269    #[tokio::test]
1270    async fn claim_renewal_loop_retries_after_transient_errors() {
1271        let store = memory_store();
1272        {
1273            let mut script = store
1274                .renewal_script
1275                .lock()
1276                .expect("renewal script should lock");
1277            script.push_back(Err("database temporarily unavailable".to_string()));
1278            script.push_back(Err("database temporarily unavailable".to_string()));
1279        }
1280        let renewals = store.renewals.clone();
1281        let stop = CancellationToken::new();
1282
1283        let handle = tokio::spawn(run_scheduled_claim_renewal_loop(
1284            store,
1285            test_claim(),
1286            std::time::Duration::from_secs(30),
1287            std::time::Duration::from_millis(1),
1288            stop.clone(),
1289        ));
1290
1291        while renewals.load(Ordering::SeqCst) < 3 {
1292            tokio::task::yield_now().await;
1293        }
1294        stop.cancel();
1295        handle.await.expect("renewal loop should stop cleanly");
1296        assert!(renewals.load(Ordering::SeqCst) >= 3);
1297    }
1298
1299    #[tokio::test]
1300    async fn scheduled_periodic_task_renews_claim_while_task_body_runs() {
1301        let shutdown = CancellationToken::new();
1302        let store = memory_store();
1303        let renewals = store.renewals.clone();
1304        let completions = store.completions.clone();
1305        let renewals_for_task = renewals.clone();
1306        let shutdown_for_hook = shutdown.clone();
1307
1308        run_scheduled_periodic_task(ScheduledPeriodicTask {
1309            name: "cleanup",
1310            namespace: "aster_test",
1311            task_name: "cleanup",
1312            display_name: "Cleanup",
1313            owner_id: "runtime-a".to_string(),
1314            // 30ms TTL derives a 10ms renewal interval, so a body that waits for
1315            // two renewals provably outlives the original claim window.
1316            claim_ttl: std::time::Duration::from_millis(30),
1317            interval_fn: test_interval,
1318            jitter_cap: None,
1319            shutdown_token: shutdown.clone(),
1320            state: (),
1321            store,
1322            hooks: RecordedTaskHooks::new(
1323                move |()| {
1324                    let renewals = renewals_for_task.clone();
1325                    async move {
1326                        tokio::time::timeout(std::time::Duration::from_secs(5), async move {
1327                            while renewals.load(Ordering::SeqCst) < 2 {
1328                                tokio::task::yield_now().await;
1329                            }
1330                        })
1331                        .await
1332                        .expect("claim should be renewed while the task body runs");
1333                        "ok"
1334                    }
1335                },
1336                |_| "panic",
1337                move |(), _name, _claim, _started_at, _finished_at, _outcome| {
1338                    let shutdown = shutdown_for_hook.clone();
1339                    async move {
1340                        shutdown.cancel();
1341                    }
1342                },
1343            ),
1344        })
1345        .await;
1346
1347        assert_eq!(completions.load(Ordering::SeqCst), 1);
1348        let renewed = renewals.load(Ordering::SeqCst);
1349        assert!(renewed >= 2);
1350        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1351        assert_eq!(
1352            renewals.load(Ordering::SeqCst),
1353            renewed,
1354            "renewal loop must stop before the claim is completed"
1355        );
1356    }
1357}
aster_forge_tasks/schedule.rs

aster_forge_tasks/
schedule.rs
