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.
88///
89/// # Errors
90///
91/// Returns an error when the heartbeat loses its lease or the store operation fails.
92pub fn evaluate_heartbeat_result<Error>(
93    lease_guard: &TaskLeaseGuard,
94    result: std::result::Result<bool, Error>,
95) -> std::result::Result<(), Error>
96where
97    Error: From<TaskCoreError> + std::fmt::Display,
98{
99    let lease = lease_guard.lease();
100    match result {
101        Ok(true) => {
102            lease_guard.record_renewed();
103            Ok(())
104        }
105        Ok(false) => {
106            tracing::info!(
107                task_id = lease.task_id,
108                processing_token = lease.processing_token,
109                "background task lease lost; stopping outdated worker"
110            );
111            Err(lease_guard.mark_lost().into())
112        }
113        Err(error) => {
114            tracing::warn!(
115                task_id = lease.task_id,
116                processing_token = lease.processing_token,
117                error = %error,
118                "background task heartbeat update failed; continuing and retrying next heartbeat"
119            );
120            lease_guard.ensure_active().map_err(Error::from)
121        }
122    }
123}
124
125/// Stops and awaits a heartbeat worker.
126pub async fn stop_task_heartbeat(stop_token: CancellationToken, heartbeat_handle: JoinHandle<()>) {
127    stop_token.cancel();
128    if let Err(error) = heartbeat_handle.await {
129        tracing::warn!(error = %error, "background task heartbeat worker stopped unexpectedly");
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use std::sync::Arc;
136    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
137    use std::time::Duration;
138
139    use chrono::{DateTime, Utc};
140    use tokio_util::sync::CancellationToken;
141
142    use super::{
143        TaskHeartbeatStore, evaluate_heartbeat_result, run_task_heartbeat_loop,
144        spawn_task_heartbeat_with_interval, stop_task_heartbeat,
145    };
146    use crate::{TaskCoreError, TaskLease, TaskLeaseGuard, task_lease_expires_at};
147
148    struct TestHeartbeatStore {
149        touches: Arc<AtomicUsize>,
150        should_match: Arc<AtomicBool>,
151    }
152
153    #[async_trait::async_trait]
154    impl TaskHeartbeatStore for TestHeartbeatStore {
155        type Error = TaskCoreError;
156
157        async fn touch_task_heartbeat(
158            &self,
159            _lease: TaskLease,
160            _now: DateTime<Utc>,
161            _lease_expires_at: DateTime<Utc>,
162        ) -> std::result::Result<bool, Self::Error> {
163            self.touches.fetch_add(1, Ordering::SeqCst);
164            Ok(self.should_match.load(Ordering::SeqCst))
165        }
166    }
167
168    #[test]
169    fn heartbeat_result_records_successful_renewal() {
170        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
171
172        evaluate_heartbeat_result::<TaskCoreError>(&guard, Ok(true))
173            .expect("heartbeat should renew");
174
175        guard.ensure_active().expect("lease should remain active");
176    }
177
178    #[test]
179    fn heartbeat_result_marks_lost_on_false_update() {
180        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
181
182        let error = evaluate_heartbeat_result::<TaskCoreError>(&guard, Ok(false))
183            .expect_err("false update should lose lease");
184
185        assert!(error.is_task_lease_lost());
186        assert!(
187            guard
188                .ensure_active()
189                .is_err_and(|error| error.is_task_lease_lost())
190        );
191    }
192
193    #[test]
194    fn heartbeat_result_keeps_retrying_transient_error_before_timeout() {
195        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
196
197        evaluate_heartbeat_result(
198            &guard,
199            Err(TaskCoreError::codec("database temporarily unavailable")),
200        )
201        .expect("transient error should keep lease alive before timeout");
202    }
203
204    #[test]
205    fn heartbeat_result_stops_after_renewal_timeout() {
206        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::ZERO);
207
208        let error = evaluate_heartbeat_result(
209            &guard,
210            Err(TaskCoreError::codec("database temporarily unavailable")),
211        )
212        .expect_err("timed-out lease should stop after transient error");
213
214        assert!(error.is_task_lease_renewal_timed_out());
215    }
216
217    #[tokio::test]
218    async fn heartbeat_loop_runs_until_stopped() {
219        let touches = Arc::new(AtomicUsize::new(0));
220        let store = TestHeartbeatStore {
221            touches: touches.clone(),
222            should_match: Arc::new(AtomicBool::new(true)),
223        };
224        let stop_token = CancellationToken::new();
225        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
226
227        let handle = tokio::spawn(run_task_heartbeat_loop(
228            store,
229            guard,
230            stop_token.clone(),
231            Duration::from_millis(1),
232            |now| task_lease_expires_at(now, 60),
233        ));
234
235        while touches.load(Ordering::SeqCst) == 0 {
236            tokio::task::yield_now().await;
237        }
238
239        stop_task_heartbeat(stop_token, handle).await;
240        assert!(touches.load(Ordering::SeqCst) >= 1);
241    }
242
243    #[tokio::test]
244    async fn spawned_heartbeat_worker_can_be_stopped() {
245        let touches = Arc::new(AtomicUsize::new(0));
246        let store = TestHeartbeatStore {
247            touches: touches.clone(),
248            should_match: Arc::new(AtomicBool::new(true)),
249        };
250        let stop_token = CancellationToken::new();
251        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
252
253        let handle = spawn_task_heartbeat_with_interval(
254            store,
255            guard,
256            stop_token.clone(),
257            Duration::from_millis(1),
258            |now| task_lease_expires_at(now, 60),
259        );
260
261        while touches.load(Ordering::SeqCst) == 0 {
262            tokio::task::yield_now().await;
263        }
264
265        stop_task_heartbeat(stop_token, handle).await;
266        assert!(touches.load(Ordering::SeqCst) >= 1);
267    }
268
269    #[tokio::test]
270    async fn heartbeat_loop_stops_when_persisted_lease_is_lost() {
271        let touches = Arc::new(AtomicUsize::new(0));
272        let store = TestHeartbeatStore {
273            touches: touches.clone(),
274            should_match: Arc::new(AtomicBool::new(false)),
275        };
276        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
277
278        run_task_heartbeat_loop(
279            store,
280            guard,
281            CancellationToken::new(),
282            Duration::from_millis(1),
283            |now| task_lease_expires_at(now, 60),
284        )
285        .await;
286
287        assert_eq!(touches.load(Ordering::SeqCst), 1);
288    }
289}
aster_forge_tasks/heartbeat.rs

aster_forge_tasks/
heartbeat.rs
