aster_forge_tasks/
execution.rs

1//! Claimed task execution lifecycle helpers.
2//!
3//! Product crates still own task rows, repositories, task bodies, error presentation, retention
4//! policy, metrics labels, and wakeup signals. This module owns the stable execution lifecycle
5//! around those product hooks: start a lease-bound context, run a heartbeat worker, ignore stale
6//! workers after lease-control failures, decide whether an ordinary task error should retry or
7//! permanently fail, and aggregate the resulting dispatch counters. Keeping this flow in Forge
8//! makes product task systems smaller while preserving the product-owned persistence boundary.
9
10use std::future::Future;
11use std::time::Duration;
12
13use chrono::{DateTime, Utc};
14use tokio_util::sync::CancellationToken;
15
16use crate::{
17    TaskDispatchOutcome, TaskExecutionContext, TaskHeartbeatStore, TaskLease, TaskRecord,
18    TaskRetryClass, run_claimed_task_batch, spawn_task_heartbeat_with_interval,
19    stop_task_heartbeat,
20};
21
22/// Minimal read-only view required to execute a claimed task.
23pub trait ExecutableTaskRecord<Kind>: TaskRecord<Kind> {
24    /// Number of attempts already persisted before the current execution.
25    fn attempt_count(&self) -> i32;
26
27    /// Maximum number of attempts allowed for automatic retry.
28    fn max_attempts(&self) -> i32;
29}
30
31/// Failure update passed to product storage after a task exhausts automatic retry.
32pub struct TaskPermanentFailure<'a> {
33    /// Attempt count to persist for the just-finished execution.
34    pub attempt_count: i32,
35    /// Truncated/storable error string.
36    pub storage_error: &'a str,
37    /// Human-facing error string used for logs and step details.
38    pub display_error: &'a str,
39    /// Serialized task steps after marking the active step failed, if available.
40    pub failed_steps_json: Option<&'a str>,
41    /// Whether a later manual retry should be allowed.
42    pub failure_can_retry: bool,
43    /// Timestamp when the failure is recorded.
44    pub finished_at: DateTime<Utc>,
45}
46
47/// Retry update passed to product storage after an automatically retryable failure.
48pub struct TaskRetryUpdate<'a> {
49    /// Attempt count to persist for the just-finished execution.
50    pub attempt_count: i32,
51    /// Timestamp when the next automatic retry should become due.
52    pub retry_at: DateTime<Utc>,
53    /// Truncated/storable error string.
54    pub storage_error: &'a str,
55    /// Human-facing error string used for logs and step details.
56    pub display_error: &'a str,
57    /// Serialized task steps after marking the active step failed, if available.
58    pub failed_steps_json: Option<&'a str>,
59}
60
61/// Product hooks required by the shared claimed-task execution lifecycle.
62#[async_trait::async_trait]
63pub trait ClaimedTaskExecutionStore<Task, Kind>: TaskHeartbeatStore + Clone + Send + Sync {
64    /// Runs the product task body.
65    async fn process_task(
66        &self,
67        task: &Task,
68        context: TaskExecutionContext,
69    ) -> std::result::Result<(), Self::Error>;
70
71    /// Returns whether an error means the current worker has lost its persisted lease.
72    fn is_lease_lost_error(&self, error: &Self::Error) -> bool;
73
74    /// Returns whether an error means heartbeat renewal exceeded the in-memory safety deadline.
75    fn is_lease_renewal_timed_out_error(&self, error: &Self::Error) -> bool;
76
77    /// Returns whether an error means cooperative shutdown should release the current lease.
78    fn is_worker_shutdown_requested_error(&self, error: &Self::Error) -> bool;
79
80    /// Classifies an ordinary task failure for retry behavior.
81    fn retry_class(&self, task: &Task, error: &Self::Error) -> TaskRetryClass;
82
83    /// Converts a task error into the string stored in the task row.
84    fn storage_error(&self, error: &Self::Error) -> String;
85
86    /// Converts the stored error string into a human-facing display string.
87    fn display_error(&self, storage_error: &str) -> String;
88
89    /// Builds serialized task steps for a failed task, if the product has step state.
90    async fn failed_steps_json(&self, task: &Task, display_error: &str) -> Option<String>;
91
92    /// Marks the claimed task permanently failed.
93    async fn mark_task_failed(
94        &self,
95        task: &Task,
96        lease: TaskLease,
97        failure: TaskPermanentFailure<'_>,
98    ) -> std::result::Result<bool, Self::Error>;
99
100    /// Marks the claimed task retryable at a future timestamp.
101    async fn mark_task_retry(
102        &self,
103        task: &Task,
104        lease: TaskLease,
105        retry: TaskRetryUpdate<'_>,
106    ) -> std::result::Result<bool, Self::Error>;
107
108    /// Releases a processing lease during cooperative shutdown.
109    async fn release_task_for_shutdown(
110        &self,
111        task: &Task,
112        lease: TaskLease,
113    ) -> std::result::Result<bool, Self::Error>;
114
115    /// Records a product metric or audit hook after a task transition.
116    fn record_task_transition(&self, task: &Task, status: &'static str);
117
118    /// Wakes the product dispatcher after retry or shutdown release creates runnable work.
119    fn wake_dispatcher(&self);
120}
121
122/// Configuration for the shared claimed-task execution lifecycle.
123#[derive(Debug, Clone, Copy)]
124pub struct ClaimedTaskExecutionConfig<LeaseExpiresFn, RetryDelayFn> {
125    /// In-memory timeout used by task code to stop unsafe stale workers.
126    pub renewal_timeout: Duration,
127    /// Interval used by the heartbeat worker.
128    pub heartbeat_interval: Duration,
129    /// Computes the persisted lease expiry for heartbeat writes.
130    pub lease_expires_at: LeaseExpiresFn,
131    /// Computes automatic retry delay in seconds from the new attempt count.
132    pub retry_delay_secs: RetryDelayFn,
133}
134
135/// Runs a claimed task batch using the shared lifecycle and aggregates counters.
136///
137/// # Errors
138///
139/// Returns a store error when claimed-task processing or outcome persistence fails.
140pub async fn run_claimed_task_batch_with_store<
141    Store,
142    Task,
143    Kind,
144    SortKey,
145    SortFn,
146    LeaseExpiresFn,
147    RetryDelayFn,
148>(
149    store: Store,
150    claimed_tasks: Vec<(Task, TaskLease)>,
151    sort_key: SortFn,
152    shutdown_token: CancellationToken,
153    config: ClaimedTaskExecutionConfig<LeaseExpiresFn, RetryDelayFn>,
154) -> std::result::Result<crate::DispatchStats, Store::Error>
155where
156    Store: ClaimedTaskExecutionStore<Task, Kind> + 'static,
157    Task: ExecutableTaskRecord<Kind> + Clone + Send + Sync + 'static,
158    Kind: Copy + std::fmt::Display + Send + Sync + 'static,
159    SortKey: Ord,
160    SortFn: FnMut(&(Task, TaskLease)) -> SortKey,
161    LeaseExpiresFn: Fn(DateTime<Utc>) -> DateTime<Utc> + Copy + Send + Sync + 'static,
162    RetryDelayFn: Fn(i32) -> i64 + Copy + Send + Sync + 'static,
163{
164    run_claimed_task_batch(claimed_tasks, sort_key, |(task, lease)| {
165        let store = store.clone();
166        let shutdown_token = shutdown_token.clone();
167        async move { process_claimed_task(store, task, lease, shutdown_token, config).await }
168    })
169    .await
170}
171
172/// Runs one claimed task through heartbeat, processing, retry, and failure handling.
173///
174/// # Errors
175///
176/// Returns a store error when heartbeat, processing, or fenced outcome persistence fails.
177#[expect(
178    clippy::too_many_lines,
179    reason = "Heartbeat shutdown, lease fencing, retry, and permanent failure form one claimed-task state transition."
180)]
181pub async fn process_claimed_task<Store, Task, Kind, LeaseExpiresFn, RetryDelayFn>(
182    store: Store,
183    task: Task,
184    lease: TaskLease,
185    shutdown_token: CancellationToken,
186    config: ClaimedTaskExecutionConfig<LeaseExpiresFn, RetryDelayFn>,
187) -> std::result::Result<TaskDispatchOutcome, Store::Error>
188where
189    Store: ClaimedTaskExecutionStore<Task, Kind> + 'static,
190    Task: ExecutableTaskRecord<Kind> + Send + Sync + 'static,
191    Kind: Copy + std::fmt::Display + Send + Sync + 'static,
192    LeaseExpiresFn: Fn(DateTime<Utc>) -> DateTime<Utc> + Copy + Send + Sync + 'static,
193    RetryDelayFn: Fn(i32) -> i64 + Copy + Send + Sync + 'static,
194{
195    let context = TaskExecutionContext::new(lease, config.renewal_timeout, shutdown_token);
196    let lease_guard = context.lease_guard().clone();
197    let heartbeat_stop = CancellationToken::new();
198    let heartbeat_handle = spawn_task_heartbeat_with_interval(
199        store.clone(),
200        lease_guard.clone(),
201        heartbeat_stop.clone(),
202        config.heartbeat_interval,
203        config.lease_expires_at,
204    );
205    let heartbeat_cancel_guard = heartbeat_stop.clone().drop_guard();
206
207    let task_result = match context.ensure_active() {
208        Ok(()) => store.process_task(&task, context).await,
209        Err(error) => Err(Store::Error::from(error)),
210    };
211    drop(heartbeat_cancel_guard);
212    stop_task_heartbeat(heartbeat_stop, heartbeat_handle).await;
213
214    match task_result {
215        Ok(()) => {
216            store.record_task_transition(&task, "succeeded");
217            Ok(TaskDispatchOutcome::succeeded())
218        }
219        Err(error)
220            if store.is_lease_lost_error(&error)
221                || store.is_lease_renewal_timed_out_error(&error)
222                || store.is_worker_shutdown_requested_error(&error) =>
223        {
224            if store.is_worker_shutdown_requested_error(&error)
225                && store.release_task_for_shutdown(&task, lease).await?
226            {
227                store.wake_dispatcher();
228            }
229            tracing::info!(
230                task_id = task.id(),
231                processing_token = lease.processing_token,
232                "background task worker stopped before completion; skipping stale completion"
233            );
234            Ok(TaskDispatchOutcome::default())
235        }
236        Err(error) => {
237            let attempt_count = task.attempt_count().saturating_add(1);
238            let storage_error = store.storage_error(&error);
239            let display_error = store.display_error(&storage_error);
240            let failed_steps_json = store.failed_steps_json(&task, &display_error).await;
241            let retry_class = store.retry_class(&task, &error);
242            let should_auto_retry =
243                retry_class.should_auto_retry() && attempt_count < task.max_attempts();
244
245            if should_auto_retry {
246                let retry_at = Utc::now()
247                    + chrono::Duration::seconds((config.retry_delay_secs)(attempt_count));
248                let retried = store
249                    .mark_task_retry(
250                        &task,
251                        lease,
252                        TaskRetryUpdate {
253                            attempt_count,
254                            retry_at,
255                            storage_error: &storage_error,
256                            display_error: &display_error,
257                            failed_steps_json: failed_steps_json.as_deref(),
258                        },
259                    )
260                    .await?;
261                if !retried {
262                    tracing::info!(
263                        task_id = task.id(),
264                        processing_token = lease.processing_token,
265                        "background task lease moved before retry state update; ignoring stale worker"
266                    );
267                    return Ok(TaskDispatchOutcome::default());
268                }
269
270                tracing::warn!(
271                    task_id = task.id(),
272                    kind = %task.kind(),
273                    attempt_count,
274                    retry_at = %retry_at,
275                    error = %display_error,
276                    "background task failed; scheduled retry"
277                );
278                store.wake_dispatcher();
279                store.record_task_transition(&task, "retry");
280                Ok(TaskDispatchOutcome::retried())
281            } else {
282                let finished_at = Utc::now();
283                let failed = store
284                    .mark_task_failed(
285                        &task,
286                        lease,
287                        TaskPermanentFailure {
288                            attempt_count,
289                            storage_error: &storage_error,
290                            display_error: &display_error,
291                            failed_steps_json: failed_steps_json.as_deref(),
292                            failure_can_retry: retry_class.can_manual_retry(),
293                            finished_at,
294                        },
295                    )
296                    .await?;
297                if !failed {
298                    tracing::info!(
299                        task_id = task.id(),
300                        processing_token = lease.processing_token,
301                        "background task lease moved before failure state update; ignoring stale worker"
302                    );
303                    return Ok(TaskDispatchOutcome::default());
304                }
305
306                tracing::warn!(
307                    task_id = task.id(),
308                    kind = %task.kind(),
309                    attempt_count,
310                    error = %display_error,
311                    "background task permanently failed"
312                );
313                store.record_task_transition(&task, "failed");
314                Ok(TaskDispatchOutcome::failed())
315            }
316        }
317    }
318}
319
320/// Converts an operation into a boxed future.
321pub fn boxed_task_future<'a, T, Error, Fut>(future: Fut) -> crate::TaskProcessFuture<'a, Error>
322where
323    Fut: Future<Output = std::result::Result<T, Error>> + Send + 'a,
324    Error: Send + 'a,
325{
326    Box::pin(async move {
327        future.await?;
328        Ok(())
329    })
330}
331
332#[cfg(test)]
333mod tests {
334    use std::collections::VecDeque;
335    use std::sync::{Arc, Mutex};
336    use std::time::Duration;
337
338    use chrono::{DateTime, Utc};
339    use tokio_util::sync::CancellationToken;
340
341    use super::{
342        ClaimedTaskExecutionConfig, ClaimedTaskExecutionStore, ExecutableTaskRecord,
343        TaskPermanentFailure, TaskRetryUpdate, process_claimed_task,
344        run_claimed_task_batch_with_store,
345    };
346    use crate::{
347        DispatchStats, TaskCoreError, TaskExecutionContext, TaskHeartbeatStore, TaskLease,
348        TaskRecord, TaskRetryClass, default_task_retry_delay_secs, task_lease_expires_at,
349    };
350
351    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
352    enum TestKind {
353        Example,
354    }
355
356    impl std::fmt::Display for TestKind {
357        fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
358            formatter.write_str("example")
359        }
360    }
361
362    #[derive(Debug, Clone)]
363    struct TestTask {
364        id: i64,
365        attempt_count: i32,
366        max_attempts: i32,
367        order: i32,
368    }
369
370    impl TaskRecord<TestKind> for TestTask {
371        fn id(&self) -> i64 {
372            self.id
373        }
374
375        fn kind(&self) -> TestKind {
376            TestKind::Example
377        }
378
379        fn payload_json(&self) -> &'static str {
380            "{}"
381        }
382
383        fn result_json(&self) -> Option<&str> {
384            None
385        }
386    }
387
388    impl ExecutableTaskRecord<TestKind> for TestTask {
389        fn attempt_count(&self) -> i32 {
390            self.attempt_count
391        }
392
393        fn max_attempts(&self) -> i32 {
394            self.max_attempts
395        }
396    }
397
398    #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
399    enum TestError {
400        #[error("{0}")]
401        Core(#[from] TaskCoreError),
402        #[error("{0}")]
403        Business(String),
404    }
405
406    #[derive(Debug, Default)]
407    struct StoreState {
408        process_results: VecDeque<std::result::Result<(), TestError>>,
409        failed: usize,
410        retried: usize,
411        released: usize,
412        wakes: usize,
413        transitions: Vec<&'static str>,
414        processed_order: Vec<i64>,
415        last_failed_steps: Option<String>,
416    }
417
418    #[derive(Clone)]
419    struct TestStore {
420        state: Arc<Mutex<StoreState>>,
421    }
422
423    impl TestStore {
424        fn with_results(results: Vec<std::result::Result<(), TestError>>) -> Self {
425            Self {
426                state: Arc::new(Mutex::new(StoreState {
427                    process_results: results.into(),
428                    ..StoreState::default()
429                })),
430            }
431        }
432
433        fn state(&self) -> std::sync::MutexGuard<'_, StoreState> {
434            self.state.lock().expect("store state should lock")
435        }
436    }
437
438    #[async_trait::async_trait]
439    impl TaskHeartbeatStore for TestStore {
440        type Error = TestError;
441
442        async fn touch_task_heartbeat(
443            &self,
444            _lease: TaskLease,
445            _now: DateTime<Utc>,
446            _lease_expires_at: DateTime<Utc>,
447        ) -> std::result::Result<bool, Self::Error> {
448            Ok(true)
449        }
450    }
451
452    #[async_trait::async_trait]
453    impl ClaimedTaskExecutionStore<TestTask, TestKind> for TestStore {
454        async fn process_task(
455            &self,
456            task: &TestTask,
457            _context: TaskExecutionContext,
458        ) -> std::result::Result<(), Self::Error> {
459            let mut state = self.state();
460            state.processed_order.push(task.id);
461            state.process_results.pop_front().unwrap_or(Ok(()))
462        }
463
464        fn is_lease_lost_error(&self, error: &Self::Error) -> bool {
465            matches!(error, TestError::Core(error) if error.is_task_lease_lost())
466        }
467
468        fn is_lease_renewal_timed_out_error(&self, error: &Self::Error) -> bool {
469            matches!(error, TestError::Core(error) if error.is_task_lease_renewal_timed_out())
470        }
471
472        fn is_worker_shutdown_requested_error(&self, error: &Self::Error) -> bool {
473            matches!(error, TestError::Core(error) if error.is_task_worker_shutdown_requested())
474        }
475
476        fn retry_class(&self, _task: &TestTask, error: &Self::Error) -> TaskRetryClass {
477            match error {
478                TestError::Business(message) if message == "never" => TaskRetryClass::Never,
479                _ => TaskRetryClass::Auto,
480            }
481        }
482
483        fn storage_error(&self, error: &Self::Error) -> String {
484            error.to_string()
485        }
486
487        fn display_error(&self, storage_error: &str) -> String {
488            storage_error.to_string()
489        }
490
491        async fn failed_steps_json(&self, _task: &TestTask, display_error: &str) -> Option<String> {
492            Some(format!("failed:{display_error}"))
493        }
494
495        async fn mark_task_failed(
496            &self,
497            _task: &TestTask,
498            _lease: TaskLease,
499            failure: TaskPermanentFailure<'_>,
500        ) -> std::result::Result<bool, Self::Error> {
501            let mut state = self.state();
502            state.failed += 1;
503            state.last_failed_steps = failure.failed_steps_json.map(str::to_string);
504            Ok(true)
505        }
506
507        async fn mark_task_retry(
508            &self,
509            _task: &TestTask,
510            _lease: TaskLease,
511            _retry: TaskRetryUpdate<'_>,
512        ) -> std::result::Result<bool, Self::Error> {
513            self.state().retried += 1;
514            Ok(true)
515        }
516
517        async fn release_task_for_shutdown(
518            &self,
519            _task: &TestTask,
520            _lease: TaskLease,
521        ) -> std::result::Result<bool, Self::Error> {
522            self.state().released += 1;
523            Ok(true)
524        }
525
526        fn record_task_transition(&self, _task: &TestTask, status: &'static str) {
527            self.state().transitions.push(status);
528        }
529
530        fn wake_dispatcher(&self) {
531            self.state().wakes += 1;
532        }
533    }
534
535    type TestExecutionConfig =
536        ClaimedTaskExecutionConfig<fn(DateTime<Utc>) -> DateTime<Utc>, fn(i32) -> i64>;
537
538    fn config() -> TestExecutionConfig {
539        ClaimedTaskExecutionConfig {
540            renewal_timeout: Duration::from_mins(1),
541            heartbeat_interval: Duration::from_mins(1),
542            lease_expires_at: |now| task_lease_expires_at(now, 60),
543            retry_delay_secs: default_task_retry_delay_secs,
544        }
545    }
546
547    fn task(id: i64, attempt_count: i32, max_attempts: i32) -> TestTask {
548        TestTask {
549            id,
550            attempt_count,
551            max_attempts,
552            order: i32::try_from(id).expect("test id should fit in i32"),
553        }
554    }
555
556    #[tokio::test]
557    async fn process_claimed_task_records_success() {
558        let store = TestStore::with_results(vec![Ok(())]);
559
560        let outcome = process_claimed_task(
561            store.clone(),
562            task(7, 0, 3),
563            TaskLease::new(7, 2),
564            CancellationToken::new(),
565            config(),
566        )
567        .await
568        .expect("task should succeed");
569
570        assert_eq!(outcome.succeeded, 1);
571        assert_eq!(store.state().transitions, vec!["succeeded"]);
572    }
573
574    #[tokio::test]
575    async fn process_claimed_task_retries_auto_failure_with_budget() {
576        let store = TestStore::with_results(vec![Err(TestError::Business("retry".to_string()))]);
577
578        let outcome = process_claimed_task(
579            store.clone(),
580            task(7, 0, 3),
581            TaskLease::new(7, 2),
582            CancellationToken::new(),
583            config(),
584        )
585        .await
586        .expect("task should retry");
587
588        assert_eq!(outcome.retried, 1);
589        let state = store.state();
590        assert_eq!(state.retried, 1);
591        assert_eq!(state.wakes, 1);
592        assert_eq!(state.transitions, vec!["retry"]);
593    }
594
595    #[tokio::test]
596    async fn process_claimed_task_fails_when_retry_budget_is_exhausted() {
597        let store = TestStore::with_results(vec![Err(TestError::Business("retry".to_string()))]);
598
599        let outcome = process_claimed_task(
600            store.clone(),
601            task(7, 2, 3),
602            TaskLease::new(7, 2),
603            CancellationToken::new(),
604            config(),
605        )
606        .await
607        .expect("task should fail permanently");
608
609        assert_eq!(outcome.failed, 1);
610        let state = store.state();
611        assert_eq!(state.failed, 1);
612        assert_eq!(state.last_failed_steps.as_deref(), Some("failed:retry"));
613        assert_eq!(state.transitions, vec!["failed"]);
614    }
615
616    #[tokio::test]
617    async fn process_claimed_task_releases_shutdown_without_failure() {
618        let lease = TaskLease::new(7, 2);
619        let store = TestStore::with_results(vec![Err(TestError::Core(
620            TaskCoreError::WorkerShutdownRequested {
621                task_id: lease.task_id,
622                processing_token: lease.processing_token,
623            },
624        ))]);
625
626        let outcome = process_claimed_task(
627            store.clone(),
628            task(7, 0, 3),
629            lease,
630            CancellationToken::new(),
631            config(),
632        )
633        .await
634        .expect("shutdown should release");
635
636        assert_eq!(outcome, crate::TaskDispatchOutcome::default());
637        let state = store.state();
638        assert_eq!(state.released, 1);
639        assert_eq!(state.wakes, 1);
640        assert_eq!(state.failed, 0);
641        assert_eq!(state.retried, 0);
642        assert!(state.transitions.is_empty());
643    }
644
645    #[tokio::test]
646    async fn batch_runner_sorts_and_aggregates_claimed_tasks() {
647        let store = TestStore::with_results(vec![Ok(()), Ok(())]);
648        let claimed = vec![
649            (task(2, 0, 3), TaskLease::new(2, 1)),
650            (task(1, 0, 3), TaskLease::new(1, 1)),
651        ];
652
653        let stats = run_claimed_task_batch_with_store(
654            store.clone(),
655            claimed,
656            |(task, _)| task.order,
657            CancellationToken::new(),
658            config(),
659        )
660        .await
661        .expect("batch should succeed");
662
663        assert_eq!(
664            stats,
665            DispatchStats {
666                claimed: 0,
667                succeeded: 2,
668                retried: 0,
669                failed: 0,
670            }
671        );
672        assert_eq!(store.state().processed_order, vec![1, 2]);
673    }
674}
aster_forge_tasks/execution.rs

aster_forge_tasks/
execution.rs
