aster_forge_tasks/
heartbeat.rs

1//! Heartbeat loop for claimed background task leases.
2//!
3//! Product crates own the persistence update that extends a task lease. Forge owns the surrounding
4//! loop: interval scheduling, cooperative stop handling, stale-worker detection, and transient
5//! storage-error handling.
6
7use chrono::{DateTime, Utc};
8use tokio::task::JoinHandle;
9use tokio::time::MissedTickBehavior;
10use tokio_util::sync::CancellationToken;
11
12use crate::{TaskCoreError, TaskLease, TaskLeaseGuard};
13
14/// Product storage adapter used by the generic heartbeat loop.
15#[async_trait::async_trait]
16pub trait TaskHeartbeatStore: Send + Sync {
17    /// Product error type returned by heartbeat storage operations.
18    type Error: From<TaskCoreError> + std::fmt::Display + Send;
19
20    /// Attempts to persist one heartbeat renewal for the given lease.
21    ///
22    /// Returning `Ok(false)` means the conditional update did not match the current task status or
23    /// processing token, so the worker must treat the lease as lost. Returning `Err(_)` is treated
24    /// as a transient storage failure unless the in-memory lease guard has already timed out.
25    async fn touch_task_heartbeat(
26        &self,
27        lease: TaskLease,
28        now: DateTime<Utc>,
29        lease_expires_at: DateTime<Utc>,
30    ) -> std::result::Result<bool, Self::Error>;
31}
32
33/// Runs heartbeat updates until stopped, the persisted lease is lost, or renewal times out.
34pub async fn run_task_heartbeat_loop<Store, LeaseExpiresFn>(
35    store: Store,
36    lease_guard: TaskLeaseGuard,
37    stop_token: CancellationToken,
38    interval: std::time::Duration,
39    lease_expires_at: LeaseExpiresFn,
40) where
41    Store: TaskHeartbeatStore,
42    LeaseExpiresFn: Fn(DateTime<Utc>) -> DateTime<Utc> + Send + Sync,
43{
44    let mut heartbeat = tokio::time::interval(interval);
45    heartbeat.set_missed_tick_behavior(MissedTickBehavior::Delay);
46    heartbeat.tick().await;
47
48    loop {
49        tokio::select! {
50            _ = stop_token.cancelled() => return,
51            _ = heartbeat.tick() => {
52                let now = Utc::now();
53                let result = tokio::select! {
54                    _ = stop_token.cancelled() => return,
55                    result = store.touch_task_heartbeat(
56                        lease_guard.lease(),
57                        now,
58                        lease_expires_at(now),
59                    ) => result,
60                };
61
62                if evaluate_heartbeat_result(&lease_guard, result).is_err() {
63                    return;
64                }
65            }
66        }
67    }
68}
69
70/// Spawns a heartbeat worker with the provided interval.
71pub fn spawn_task_heartbeat_with_interval<Store, LeaseExpiresFn>(
72    store: Store,
73    lease_guard: TaskLeaseGuard,
74    stop_token: CancellationToken,
75    interval: std::time::Duration,
76    lease_expires_at: LeaseExpiresFn,
77) -> JoinHandle<()>
78where
79    Store: TaskHeartbeatStore + 'static,
80    LeaseExpiresFn: Fn(DateTime<Utc>) -> DateTime<Utc> + Send + Sync + 'static,
81{
82    tokio::spawn(async move {
83        run_task_heartbeat_loop(store, lease_guard, stop_token, interval, lease_expires_at).await;
84    })
85}
86
87/// Evaluates one persisted heartbeat result and updates the in-memory lease guard.
88pub fn evaluate_heartbeat_result<Error>(
89    lease_guard: &TaskLeaseGuard,
90    result: std::result::Result<bool, Error>,
91) -> std::result::Result<(), Error>
92where
93    Error: From<TaskCoreError> + std::fmt::Display,
94{
95    let lease = lease_guard.lease();
96    match result {
97        Ok(true) => {
98            lease_guard.record_renewed();
99            Ok(())
100        }
101        Ok(false) => {
102            tracing::info!(
103                task_id = lease.task_id,
104                processing_token = lease.processing_token,
105                "background task lease lost; stopping outdated worker"
106            );
107            Err(lease_guard.mark_lost().into())
108        }
109        Err(error) => {
110            tracing::warn!(
111                task_id = lease.task_id,
112                processing_token = lease.processing_token,
113                error = %error,
114                "background task heartbeat update failed; continuing and retrying next heartbeat"
115            );
116            lease_guard.ensure_active().map_err(Error::from)
117        }
118    }
119}
120
121/// Stops and awaits a heartbeat worker.
122pub async fn stop_task_heartbeat(stop_token: CancellationToken, heartbeat_handle: JoinHandle<()>) {
123    stop_token.cancel();
124    if let Err(error) = heartbeat_handle.await {
125        tracing::warn!(error = %error, "background task heartbeat worker stopped unexpectedly");
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use std::sync::Arc;
132    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
133    use std::time::Duration;
134
135    use chrono::{DateTime, Utc};
136    use tokio_util::sync::CancellationToken;
137
138    use super::{
139        TaskHeartbeatStore, evaluate_heartbeat_result, run_task_heartbeat_loop,
140        spawn_task_heartbeat_with_interval, stop_task_heartbeat,
141    };
142    use crate::{TaskCoreError, TaskLease, TaskLeaseGuard, task_lease_expires_at};
143
144    struct TestHeartbeatStore {
145        touches: Arc<AtomicUsize>,
146        should_match: Arc<AtomicBool>,
147    }
148
149    #[async_trait::async_trait]
150    impl TaskHeartbeatStore for TestHeartbeatStore {
151        type Error = TaskCoreError;
152
153        async fn touch_task_heartbeat(
154            &self,
155            _lease: TaskLease,
156            _now: DateTime<Utc>,
157            _lease_expires_at: DateTime<Utc>,
158        ) -> std::result::Result<bool, Self::Error> {
159            self.touches.fetch_add(1, Ordering::SeqCst);
160            Ok(self.should_match.load(Ordering::SeqCst))
161        }
162    }
163
164    #[test]
165    fn heartbeat_result_records_successful_renewal() {
166        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_secs(60));
167
168        evaluate_heartbeat_result::<TaskCoreError>(&guard, Ok(true))
169            .expect("heartbeat should renew");
170
171        guard.ensure_active().expect("lease should remain active");
172    }
173
174    #[test]
175    fn heartbeat_result_marks_lost_on_false_update() {
176        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_secs(60));
177
178        let error = evaluate_heartbeat_result::<TaskCoreError>(&guard, Ok(false))
179            .expect_err("false update should lose lease");
180
181        assert!(error.is_task_lease_lost());
182        assert!(
183            guard
184                .ensure_active()
185                .is_err_and(|error| error.is_task_lease_lost())
186        );
187    }
188
189    #[test]
190    fn heartbeat_result_keeps_retrying_transient_error_before_timeout() {
191        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_secs(60));
192
193        evaluate_heartbeat_result(
194            &guard,
195            Err(TaskCoreError::codec("database temporarily unavailable")),
196        )
197        .expect("transient error should keep lease alive before timeout");
198    }
199
200    #[test]
201    fn heartbeat_result_stops_after_renewal_timeout() {
202        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::ZERO);
203
204        let error = evaluate_heartbeat_result(
205            &guard,
206            Err(TaskCoreError::codec("database temporarily unavailable")),
207        )
208        .expect_err("timed-out lease should stop after transient error");
209
210        assert!(error.is_task_lease_renewal_timed_out());
211    }
212
213    #[tokio::test]
214    async fn heartbeat_loop_runs_until_stopped() {
215        let touches = Arc::new(AtomicUsize::new(0));
216        let store = TestHeartbeatStore {
217            touches: touches.clone(),
218            should_match: Arc::new(AtomicBool::new(true)),
219        };
220        let stop_token = CancellationToken::new();
221        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_secs(60));
222
223        let handle = tokio::spawn(run_task_heartbeat_loop(
224            store,
225            guard,
226            stop_token.clone(),
227            Duration::from_millis(1),
228            |now| task_lease_expires_at(now, 60),
229        ));
230
231        while touches.load(Ordering::SeqCst) == 0 {
232            tokio::task::yield_now().await;
233        }
234
235        stop_task_heartbeat(stop_token, handle).await;
236        assert!(touches.load(Ordering::SeqCst) >= 1);
237    }
238
239    #[tokio::test]
240    async fn spawned_heartbeat_worker_can_be_stopped() {
241        let touches = Arc::new(AtomicUsize::new(0));
242        let store = TestHeartbeatStore {
243            touches: touches.clone(),
244            should_match: Arc::new(AtomicBool::new(true)),
245        };
246        let stop_token = CancellationToken::new();
247        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_secs(60));
248
249        let handle = spawn_task_heartbeat_with_interval(
250            store,
251            guard,
252            stop_token.clone(),
253            Duration::from_millis(1),
254            |now| task_lease_expires_at(now, 60),
255        );
256
257        while touches.load(Ordering::SeqCst) == 0 {
258            tokio::task::yield_now().await;
259        }
260
261        stop_task_heartbeat(stop_token, handle).await;
262        assert!(touches.load(Ordering::SeqCst) >= 1);
263    }
264
265    #[tokio::test]
266    async fn heartbeat_loop_stops_when_persisted_lease_is_lost() {
267        let touches = Arc::new(AtomicUsize::new(0));
268        let store = TestHeartbeatStore {
269            touches: touches.clone(),
270            should_match: Arc::new(AtomicBool::new(false)),
271        };
272        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_secs(60));
273
274        run_task_heartbeat_loop(
275            store,
276            guard,
277            CancellationToken::new(),
278            Duration::from_millis(1),
279            |now| task_lease_expires_at(now, 60),
280        )
281        .await;
282
283        assert_eq!(touches.load(Ordering::SeqCst), 1);
284    }
285}
aster_forge_tasks/heartbeat.rs

aster_forge_tasks/
heartbeat.rs
