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.
136pub async fn run_claimed_task_batch_with_store<
137    Store,
138    Task,
139    Kind,
140    SortKey,
141    SortFn,
142    LeaseExpiresFn,
143    RetryDelayFn,
144>(
145    store: Store,
146    claimed_tasks: Vec<(Task, TaskLease)>,
147    sort_key: SortFn,
148    shutdown_token: CancellationToken,
149    config: ClaimedTaskExecutionConfig<LeaseExpiresFn, RetryDelayFn>,
150) -> std::result::Result<crate::DispatchStats, Store::Error>
151where
152    Store: ClaimedTaskExecutionStore<Task, Kind> + 'static,
153    Task: ExecutableTaskRecord<Kind> + Clone + Send + Sync + 'static,
154    Kind: Copy + std::fmt::Display + Send + Sync + 'static,
155    SortKey: Ord,
156    SortFn: FnMut(&(Task, TaskLease)) -> SortKey,
157    LeaseExpiresFn: Fn(DateTime<Utc>) -> DateTime<Utc> + Copy + Send + Sync + 'static,
158    RetryDelayFn: Fn(i32) -> i64 + Copy + Send + Sync + 'static,
159{
160    run_claimed_task_batch(claimed_tasks, sort_key, |(task, lease)| {
161        let store = store.clone();
162        let shutdown_token = shutdown_token.clone();
163        async move { process_claimed_task(store, task, lease, shutdown_token, config).await }
164    })
165    .await
166}
167
168/// Runs one claimed task through heartbeat, processing, retry, and failure handling.
169pub async fn process_claimed_task<Store, Task, Kind, LeaseExpiresFn, RetryDelayFn>(
170    store: Store,
171    task: Task,
172    lease: TaskLease,
173    shutdown_token: CancellationToken,
174    config: ClaimedTaskExecutionConfig<LeaseExpiresFn, RetryDelayFn>,
175) -> std::result::Result<TaskDispatchOutcome, Store::Error>
176where
177    Store: ClaimedTaskExecutionStore<Task, Kind> + 'static,
178    Task: ExecutableTaskRecord<Kind> + Send + Sync + 'static,
179    Kind: Copy + std::fmt::Display + Send + Sync + 'static,
180    LeaseExpiresFn: Fn(DateTime<Utc>) -> DateTime<Utc> + Copy + Send + Sync + 'static,
181    RetryDelayFn: Fn(i32) -> i64 + Copy + Send + Sync + 'static,
182{
183    let context = TaskExecutionContext::new(lease, config.renewal_timeout, shutdown_token);
184    let lease_guard = context.lease_guard().clone();
185    let heartbeat_stop = CancellationToken::new();
186    let heartbeat_handle = spawn_task_heartbeat_with_interval(
187        store.clone(),
188        lease_guard.clone(),
189        heartbeat_stop.clone(),
190        config.heartbeat_interval,
191        config.lease_expires_at,
192    );
193    let heartbeat_cancel_guard = heartbeat_stop.clone().drop_guard();
194
195    let task_result = match context.ensure_active() {
196        Ok(()) => store.process_task(&task, context).await,
197        Err(error) => Err(Store::Error::from(error)),
198    };
199    drop(heartbeat_cancel_guard);
200    stop_task_heartbeat(heartbeat_stop, heartbeat_handle).await;
201
202    match task_result {
203        Ok(()) => {
204            store.record_task_transition(&task, "succeeded");
205            Ok(TaskDispatchOutcome::succeeded())
206        }
207        Err(error)
208            if store.is_lease_lost_error(&error)
209                || store.is_lease_renewal_timed_out_error(&error)
210                || store.is_worker_shutdown_requested_error(&error) =>
211        {
212            if store.is_worker_shutdown_requested_error(&error)
213                && store.release_task_for_shutdown(&task, lease).await?
214            {
215                store.wake_dispatcher();
216            }
217            tracing::info!(
218                task_id = task.id(),
219                processing_token = lease.processing_token,
220                "background task worker stopped before completion; skipping stale completion"
221            );
222            Ok(TaskDispatchOutcome::default())
223        }
224        Err(error) => {
225            let attempt_count = task.attempt_count().saturating_add(1);
226            let storage_error = store.storage_error(&error);
227            let display_error = store.display_error(&storage_error);
228            let failed_steps_json = store.failed_steps_json(&task, &display_error).await;
229            let retry_class = store.retry_class(&task, &error);
230            let should_auto_retry =
231                retry_class.should_auto_retry() && attempt_count < task.max_attempts();
232
233            if should_auto_retry {
234                let retry_at = Utc::now()
235                    + chrono::Duration::seconds((config.retry_delay_secs)(attempt_count));
236                let retried = store
237                    .mark_task_retry(
238                        &task,
239                        lease,
240                        TaskRetryUpdate {
241                            attempt_count,
242                            retry_at,
243                            storage_error: &storage_error,
244                            display_error: &display_error,
245                            failed_steps_json: failed_steps_json.as_deref(),
246                        },
247                    )
248                    .await?;
249                if !retried {
250                    tracing::info!(
251                        task_id = task.id(),
252                        processing_token = lease.processing_token,
253                        "background task lease moved before retry state update; ignoring stale worker"
254                    );
255                    return Ok(TaskDispatchOutcome::default());
256                }
257
258                tracing::warn!(
259                    task_id = task.id(),
260                    kind = %task.kind(),
261                    attempt_count,
262                    retry_at = %retry_at,
263                    error = %display_error,
264                    "background task failed; scheduled retry"
265                );
266                store.wake_dispatcher();
267                store.record_task_transition(&task, "retry");
268                Ok(TaskDispatchOutcome::retried())
269            } else {
270                let finished_at = Utc::now();
271                let failed = store
272                    .mark_task_failed(
273                        &task,
274                        lease,
275                        TaskPermanentFailure {
276                            attempt_count,
277                            storage_error: &storage_error,
278                            display_error: &display_error,
279                            failed_steps_json: failed_steps_json.as_deref(),
280                            failure_can_retry: retry_class.can_manual_retry(),
281                            finished_at,
282                        },
283                    )
284                    .await?;
285                if !failed {
286                    tracing::info!(
287                        task_id = task.id(),
288                        processing_token = lease.processing_token,
289                        "background task lease moved before failure state update; ignoring stale worker"
290                    );
291                    return Ok(TaskDispatchOutcome::default());
292                }
293
294                tracing::warn!(
295                    task_id = task.id(),
296                    kind = %task.kind(),
297                    attempt_count,
298                    error = %display_error,
299                    "background task permanently failed"
300                );
301                store.record_task_transition(&task, "failed");
302                Ok(TaskDispatchOutcome::failed())
303            }
304        }
305    }
306}
307
308/// Converts an operation into a boxed future.
309pub fn boxed_task_future<'a, T, Error, Fut>(future: Fut) -> crate::TaskProcessFuture<'a, Error>
310where
311    Fut: Future<Output = std::result::Result<T, Error>> + Send + 'a,
312    Error: Send + 'a,
313{
314    Box::pin(async move {
315        future.await?;
316        Ok(())
317    })
318}
319
320#[cfg(test)]
321mod tests {
322    use std::collections::VecDeque;
323    use std::sync::{Arc, Mutex};
324    use std::time::Duration;
325
326    use chrono::{DateTime, Utc};
327    use tokio_util::sync::CancellationToken;
328
329    use super::{
330        ClaimedTaskExecutionConfig, ClaimedTaskExecutionStore, ExecutableTaskRecord,
331        TaskPermanentFailure, TaskRetryUpdate, process_claimed_task,
332        run_claimed_task_batch_with_store,
333    };
334    use crate::{
335        DispatchStats, TaskCoreError, TaskExecutionContext, TaskHeartbeatStore, TaskLease,
336        TaskRecord, TaskRetryClass, default_task_retry_delay_secs, task_lease_expires_at,
337    };
338
339    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
340    enum TestKind {
341        Example,
342    }
343
344    impl std::fmt::Display for TestKind {
345        fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
346            formatter.write_str("example")
347        }
348    }
349
350    #[derive(Debug, Clone)]
351    struct TestTask {
352        id: i64,
353        attempt_count: i32,
354        max_attempts: i32,
355        order: i32,
356    }
357
358    impl TaskRecord<TestKind> for TestTask {
359        fn id(&self) -> i64 {
360            self.id
361        }
362
363        fn kind(&self) -> TestKind {
364            TestKind::Example
365        }
366
367        fn payload_json(&self) -> &str {
368            "{}"
369        }
370
371        fn result_json(&self) -> Option<&str> {
372            None
373        }
374    }
375
376    impl ExecutableTaskRecord<TestKind> for TestTask {
377        fn attempt_count(&self) -> i32 {
378            self.attempt_count
379        }
380
381        fn max_attempts(&self) -> i32 {
382            self.max_attempts
383        }
384    }
385
386    #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
387    enum TestError {
388        #[error("{0}")]
389        Core(#[from] TaskCoreError),
390        #[error("{0}")]
391        Business(String),
392    }
393
394    #[derive(Debug, Default)]
395    struct StoreState {
396        process_results: VecDeque<std::result::Result<(), TestError>>,
397        failed: usize,
398        retried: usize,
399        released: usize,
400        wakes: usize,
401        transitions: Vec<&'static str>,
402        processed_order: Vec<i64>,
403        last_failed_steps: Option<String>,
404    }
405
406    #[derive(Clone)]
407    struct TestStore {
408        state: Arc<Mutex<StoreState>>,
409    }
410
411    impl TestStore {
412        fn with_results(results: Vec<std::result::Result<(), TestError>>) -> Self {
413            Self {
414                state: Arc::new(Mutex::new(StoreState {
415                    process_results: results.into(),
416                    ..StoreState::default()
417                })),
418            }
419        }
420
421        fn state(&self) -> std::sync::MutexGuard<'_, StoreState> {
422            self.state.lock().expect("store state should lock")
423        }
424    }
425
426    #[async_trait::async_trait]
427    impl TaskHeartbeatStore for TestStore {
428        type Error = TestError;
429
430        async fn touch_task_heartbeat(
431            &self,
432            _lease: TaskLease,
433            _now: DateTime<Utc>,
434            _lease_expires_at: DateTime<Utc>,
435        ) -> std::result::Result<bool, Self::Error> {
436            Ok(true)
437        }
438    }
439
440    #[async_trait::async_trait]
441    impl ClaimedTaskExecutionStore<TestTask, TestKind> for TestStore {
442        async fn process_task(
443            &self,
444            task: &TestTask,
445            _context: TaskExecutionContext,
446        ) -> std::result::Result<(), Self::Error> {
447            let mut state = self.state();
448            state.processed_order.push(task.id);
449            state.process_results.pop_front().unwrap_or(Ok(()))
450        }
451
452        fn is_lease_lost_error(&self, error: &Self::Error) -> bool {
453            matches!(error, TestError::Core(error) if error.is_task_lease_lost())
454        }
455
456        fn is_lease_renewal_timed_out_error(&self, error: &Self::Error) -> bool {
457            matches!(error, TestError::Core(error) if error.is_task_lease_renewal_timed_out())
458        }
459
460        fn is_worker_shutdown_requested_error(&self, error: &Self::Error) -> bool {
461            matches!(error, TestError::Core(error) if error.is_task_worker_shutdown_requested())
462        }
463
464        fn retry_class(&self, _task: &TestTask, error: &Self::Error) -> TaskRetryClass {
465            match error {
466                TestError::Business(message) if message == "never" => TaskRetryClass::Never,
467                _ => TaskRetryClass::Auto,
468            }
469        }
470
471        fn storage_error(&self, error: &Self::Error) -> String {
472            error.to_string()
473        }
474
475        fn display_error(&self, storage_error: &str) -> String {
476            storage_error.to_string()
477        }
478
479        async fn failed_steps_json(&self, _task: &TestTask, display_error: &str) -> Option<String> {
480            Some(format!("failed:{display_error}"))
481        }
482
483        async fn mark_task_failed(
484            &self,
485            _task: &TestTask,
486            _lease: TaskLease,
487            failure: TaskPermanentFailure<'_>,
488        ) -> std::result::Result<bool, Self::Error> {
489            let mut state = self.state();
490            state.failed += 1;
491            state.last_failed_steps = failure.failed_steps_json.map(str::to_string);
492            Ok(true)
493        }
494
495        async fn mark_task_retry(
496            &self,
497            _task: &TestTask,
498            _lease: TaskLease,
499            _retry: TaskRetryUpdate<'_>,
500        ) -> std::result::Result<bool, Self::Error> {
501            self.state().retried += 1;
502            Ok(true)
503        }
504
505        async fn release_task_for_shutdown(
506            &self,
507            _task: &TestTask,
508            _lease: TaskLease,
509        ) -> std::result::Result<bool, Self::Error> {
510            self.state().released += 1;
511            Ok(true)
512        }
513
514        fn record_task_transition(&self, _task: &TestTask, status: &'static str) {
515            self.state().transitions.push(status);
516        }
517
518        fn wake_dispatcher(&self) {
519            self.state().wakes += 1;
520        }
521    }
522
523    type TestExecutionConfig =
524        ClaimedTaskExecutionConfig<fn(DateTime<Utc>) -> DateTime<Utc>, fn(i32) -> i64>;
525
526    fn config() -> TestExecutionConfig {
527        ClaimedTaskExecutionConfig {
528            renewal_timeout: Duration::from_secs(60),
529            heartbeat_interval: Duration::from_secs(60),
530            lease_expires_at: |now| task_lease_expires_at(now, 60),
531            retry_delay_secs: default_task_retry_delay_secs,
532        }
533    }
534
535    fn task(id: i64, attempt_count: i32, max_attempts: i32) -> TestTask {
536        TestTask {
537            id,
538            attempt_count,
539            max_attempts,
540            order: i32::try_from(id).expect("test id should fit in i32"),
541        }
542    }
543
544    #[tokio::test]
545    async fn process_claimed_task_records_success() {
546        let store = TestStore::with_results(vec![Ok(())]);
547
548        let outcome = process_claimed_task(
549            store.clone(),
550            task(7, 0, 3),
551            TaskLease::new(7, 2),
552            CancellationToken::new(),
553            config(),
554        )
555        .await
556        .expect("task should succeed");
557
558        assert_eq!(outcome.succeeded, 1);
559        assert_eq!(store.state().transitions, vec!["succeeded"]);
560    }
561
562    #[tokio::test]
563    async fn process_claimed_task_retries_auto_failure_with_budget() {
564        let store = TestStore::with_results(vec![Err(TestError::Business("retry".to_string()))]);
565
566        let outcome = process_claimed_task(
567            store.clone(),
568            task(7, 0, 3),
569            TaskLease::new(7, 2),
570            CancellationToken::new(),
571            config(),
572        )
573        .await
574        .expect("task should retry");
575
576        assert_eq!(outcome.retried, 1);
577        let state = store.state();
578        assert_eq!(state.retried, 1);
579        assert_eq!(state.wakes, 1);
580        assert_eq!(state.transitions, vec!["retry"]);
581    }
582
583    #[tokio::test]
584    async fn process_claimed_task_fails_when_retry_budget_is_exhausted() {
585        let store = TestStore::with_results(vec![Err(TestError::Business("retry".to_string()))]);
586
587        let outcome = process_claimed_task(
588            store.clone(),
589            task(7, 2, 3),
590            TaskLease::new(7, 2),
591            CancellationToken::new(),
592            config(),
593        )
594        .await
595        .expect("task should fail permanently");
596
597        assert_eq!(outcome.failed, 1);
598        let state = store.state();
599        assert_eq!(state.failed, 1);
600        assert_eq!(state.last_failed_steps.as_deref(), Some("failed:retry"));
601        assert_eq!(state.transitions, vec!["failed"]);
602    }
603
604    #[tokio::test]
605    async fn process_claimed_task_releases_shutdown_without_failure() {
606        let lease = TaskLease::new(7, 2);
607        let store = TestStore::with_results(vec![Err(TestError::Core(
608            TaskCoreError::WorkerShutdownRequested {
609                task_id: lease.task_id,
610                processing_token: lease.processing_token,
611            },
612        ))]);
613
614        let outcome = process_claimed_task(
615            store.clone(),
616            task(7, 0, 3),
617            lease,
618            CancellationToken::new(),
619            config(),
620        )
621        .await
622        .expect("shutdown should release");
623
624        assert_eq!(outcome, crate::TaskDispatchOutcome::default());
625        let state = store.state();
626        assert_eq!(state.released, 1);
627        assert_eq!(state.wakes, 1);
628        assert_eq!(state.failed, 0);
629        assert_eq!(state.retried, 0);
630        assert!(state.transitions.is_empty());
631    }
632
633    #[tokio::test]
634    async fn batch_runner_sorts_and_aggregates_claimed_tasks() {
635        let store = TestStore::with_results(vec![Ok(()), Ok(())]);
636        let claimed = vec![
637            (task(2, 0, 3), TaskLease::new(2, 1)),
638            (task(1, 0, 3), TaskLease::new(1, 1)),
639        ];
640
641        let stats = run_claimed_task_batch_with_store(
642            store.clone(),
643            claimed,
644            |(task, _)| task.order,
645            CancellationToken::new(),
646            config(),
647        )
648        .await
649        .expect("batch should succeed");
650
651        assert_eq!(
652            stats,
653            DispatchStats {
654                claimed: 0,
655                succeeded: 2,
656                retried: 0,
657                failed: 0,
658            }
659        );
660        assert_eq!(store.state().processed_order, vec![1, 2]);
661    }
662}
aster_forge_tasks/execution.rs

aster_forge_tasks/
execution.rs
