aster_forge_tasks/
lease.rs

1//! Processing lease guards for background task workers.
2//!
3//! A processing lease protects persisted task state from stale workers. Product crates still own
4//! the database columns and compare-and-swap updates, but Forge owns the in-memory guard used by
5//! task code and heartbeat loops to decide whether the current worker may keep writing progress.
6
7use std::sync::{Arc, Mutex, MutexGuard};
8use std::time::{Duration, Instant};
9
10use tokio_util::sync::CancellationToken;
11
12use crate::{Result, TaskCoreError};
13
14/// Persisted processing lease assigned when a task is claimed.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct TaskLease {
17    /// Persisted task identifier.
18    pub task_id: i64,
19    /// Processing token assigned by the successful claim.
20    pub processing_token: i64,
21}
22
23impl TaskLease {
24    /// Creates a task processing lease.
25    #[must_use]
26    pub const fn new(task_id: i64, processing_token: i64) -> Self {
27        Self {
28            task_id,
29            processing_token,
30        }
31    }
32}
33
34/// Shared in-memory lease guard observed by task code and heartbeat code.
35#[derive(Debug, Clone)]
36pub struct TaskLeaseGuard {
37    lease: TaskLease,
38    renewal_timeout: Duration,
39    shutdown_token: Option<CancellationToken>,
40    state: Arc<Mutex<TaskLeaseGuardState>>,
41}
42
43#[derive(Debug)]
44struct TaskLeaseGuardState {
45    last_renewed_at: Instant,
46    termination: Option<TaskLeaseTermination>,
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50enum TaskLeaseTermination {
51    Lost,
52    RenewalTimedOut,
53    ShutdownRequested,
54}
55
56impl TaskLeaseGuard {
57    /// Creates a guard with the provided renewal timeout.
58    #[must_use]
59    pub fn new(lease: TaskLease, renewal_timeout: Duration) -> Self {
60        Self {
61            lease,
62            renewal_timeout,
63            shutdown_token: None,
64            state: Arc::new(Mutex::new(TaskLeaseGuardState {
65                last_renewed_at: Instant::now(),
66                termination: None,
67            })),
68        }
69    }
70
71    /// Creates a guard that also observes worker shutdown.
72    #[must_use]
73    pub fn with_shutdown_token(
74        lease: TaskLease,
75        renewal_timeout: Duration,
76        shutdown_token: CancellationToken,
77    ) -> Self {
78        Self {
79            shutdown_token: Some(shutdown_token),
80            ..Self::new(lease, renewal_timeout)
81        }
82    }
83
84    /// Returns the persisted lease represented by this guard.
85    #[must_use]
86    pub const fn lease(&self) -> TaskLease {
87        self.lease
88    }
89
90    /// Records a successful persistent lease renewal.
91    pub fn record_renewed(&self) {
92        let mut state = self.lock_state();
93        if state.termination.is_none() {
94            state.last_renewed_at = Instant::now();
95        }
96    }
97
98    /// Marks the lease as lost and returns the corresponding error.
99    #[must_use]
100    pub fn mark_lost(&self) -> TaskCoreError {
101        let mut state = self.lock_state();
102        state.termination = Some(TaskLeaseTermination::Lost);
103        task_lease_lost(self.lease)
104    }
105
106    /// Marks worker shutdown and returns the corresponding error.
107    #[must_use]
108    pub fn mark_shutdown_requested(&self) -> TaskCoreError {
109        let mut state = self.lock_state();
110        state.termination = Some(TaskLeaseTermination::ShutdownRequested);
111        task_worker_shutdown_requested(self.lease)
112    }
113
114    /// Returns success only while the lease is still safe for writes.
115    ///
116    /// # Errors
117    ///
118    /// Returns [`TaskCoreError`] when the lease has been lost or cancellation was requested.
119    pub fn ensure_active(&self) -> Result<()> {
120        let mut state = self.lock_state();
121        match state.termination {
122            Some(TaskLeaseTermination::Lost) => return Err(task_lease_lost(self.lease)),
123            Some(TaskLeaseTermination::RenewalTimedOut) => {
124                return Err(task_lease_renewal_timed_out(self.lease));
125            }
126            Some(TaskLeaseTermination::ShutdownRequested) => {
127                return Err(task_worker_shutdown_requested(self.lease));
128            }
129            None => {}
130        }
131        if self
132            .shutdown_token
133            .as_ref()
134            .is_some_and(CancellationToken::is_cancelled)
135        {
136            state.termination = Some(TaskLeaseTermination::ShutdownRequested);
137            return Err(task_worker_shutdown_requested(self.lease));
138        }
139        if state.last_renewed_at.elapsed() >= self.renewal_timeout {
140            state.termination = Some(TaskLeaseTermination::RenewalTimedOut);
141            return Err(task_lease_renewal_timed_out(self.lease));
142        }
143        Ok(())
144    }
145
146    fn lock_state(&self) -> MutexGuard<'_, TaskLeaseGuardState> {
147        match self.state.lock() {
148            Ok(guard) => guard,
149            Err(poisoned) => poisoned.into_inner(),
150        }
151    }
152}
153
154/// Task execution context passed to product task implementations.
155#[derive(Debug, Clone)]
156pub struct TaskExecutionContext {
157    lease_guard: TaskLeaseGuard,
158    shutdown_token: CancellationToken,
159}
160
161impl TaskExecutionContext {
162    /// Creates a task execution context.
163    #[must_use]
164    pub fn new(
165        lease: TaskLease,
166        renewal_timeout: Duration,
167        shutdown_token: CancellationToken,
168    ) -> Self {
169        Self {
170            lease_guard: TaskLeaseGuard::with_shutdown_token(
171                lease,
172                renewal_timeout,
173                shutdown_token.clone(),
174            ),
175            shutdown_token,
176        }
177    }
178
179    /// Returns the lease guard used by progress and heartbeat updates.
180    #[must_use]
181    pub const fn lease_guard(&self) -> &TaskLeaseGuard {
182        &self.lease_guard
183    }
184
185    /// Returns success only while the worker should continue task execution.
186    ///
187    /// # Errors
188    ///
189    /// Returns [`TaskCoreError`] when the lease has been lost or shutdown was requested.
190    pub fn ensure_active(&self) -> Result<()> {
191        self.lease_guard.ensure_active()
192    }
193
194    /// Sleeps until `duration` elapses or shutdown is requested.
195    ///
196    /// # Errors
197    ///
198    /// Returns [`TaskCoreError`] when shutdown wins before the requested sleep completes.
199    pub async fn sleep_or_shutdown(&self, duration: Duration) -> Result<()> {
200        self.lease_guard.ensure_active()?;
201
202        tokio::select! {
203            biased;
204            () = self.shutdown_token.cancelled() => Err(self.lease_guard.mark_shutdown_requested()),
205            () = tokio::time::sleep(duration) => Ok(()),
206        }
207    }
208
209    /// Waits for shutdown and then returns the shutdown-requested lease error.
210    ///
211    /// # Errors
212    ///
213    /// Returns [`TaskCoreError`] when shutdown or lease loss has been observed.
214    pub async fn shutdown_requested(&self) -> Result<()> {
215        self.shutdown_token.cancelled().await;
216        Err(self.lease_guard.mark_shutdown_requested())
217    }
218}
219
220/// Creates the error used when a worker loses its persisted lease.
221pub const fn task_lease_lost(lease: TaskLease) -> TaskCoreError {
222    TaskCoreError::LeaseLost {
223        task_id: lease.task_id,
224        processing_token: lease.processing_token,
225    }
226}
227
228/// Creates the error used when a worker exceeds its renewal timeout.
229pub const fn task_lease_renewal_timed_out(lease: TaskLease) -> TaskCoreError {
230    TaskCoreError::LeaseRenewalTimedOut {
231        task_id: lease.task_id,
232        processing_token: lease.processing_token,
233    }
234}
235
236/// Creates the error used when a worker observes cooperative shutdown.
237pub const fn task_worker_shutdown_requested(lease: TaskLease) -> TaskCoreError {
238    TaskCoreError::WorkerShutdownRequested {
239        task_id: lease.task_id,
240        processing_token: lease.processing_token,
241    }
242}
243
244/// Returns the persisted lease expiry timestamp for a claim or heartbeat update.
245#[must_use]
246pub fn task_lease_expires_at(
247    now: chrono::DateTime<chrono::Utc>,
248    processing_stale_secs: i64,
249) -> chrono::DateTime<chrono::Utc> {
250    now + chrono::Duration::seconds(processing_stale_secs.max(1))
251}
252
253/// Returns the in-memory renewal timeout used to stop unsafe workers.
254#[must_use]
255pub fn task_lease_renewal_timeout(processing_stale_secs: i64, heartbeat_secs: u64) -> Duration {
256    let stale_secs = i64_to_u64_saturating(processing_stale_secs.max(1));
257    let heartbeat_secs = heartbeat_secs.max(1);
258    Duration::from_secs(stale_secs.saturating_sub(heartbeat_secs).max(1))
259}
260
261fn i64_to_u64_saturating(value: i64) -> u64 {
262    u64::try_from(value).unwrap_or(u64::MAX)
263}
264
265#[cfg(test)]
266mod tests {
267    use std::time::Duration;
268
269    use chrono::Utc;
270    use tokio_util::sync::CancellationToken;
271
272    use super::{
273        TaskExecutionContext, TaskLease, TaskLeaseGuard, task_lease_expires_at,
274        task_lease_renewal_timeout,
275    };
276
277    #[test]
278    fn lease_guard_reports_lost_lease_after_mark_lost() {
279        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::from_mins(1));
280
281        let error = guard.mark_lost();
282
283        assert!(error.is_task_lease_lost());
284        assert!(
285            guard
286                .ensure_active()
287                .is_err_and(|error| error.is_task_lease_lost())
288        );
289    }
290
291    #[test]
292    fn lease_guard_reports_renewal_timeout() {
293        let guard = TaskLeaseGuard::new(TaskLease::new(7, 2), Duration::ZERO);
294
295        let error = guard.ensure_active().expect_err("lease should time out");
296
297        assert!(error.is_task_lease_renewal_timed_out());
298    }
299
300    #[test]
301    fn lease_guard_observes_shutdown_token() {
302        let shutdown_token = CancellationToken::new();
303        let guard = TaskLeaseGuard::with_shutdown_token(
304            TaskLease::new(7, 2),
305            Duration::from_mins(1),
306            shutdown_token.clone(),
307        );
308
309        shutdown_token.cancel();
310
311        assert!(
312            guard
313                .ensure_active()
314                .is_err_and(|error| error.is_task_worker_shutdown_requested())
315        );
316    }
317
318    #[tokio::test]
319    async fn execution_context_sleep_returns_on_shutdown() {
320        let shutdown_token = CancellationToken::new();
321        let context = TaskExecutionContext::new(
322            TaskLease::new(7, 2),
323            Duration::from_mins(1),
324            shutdown_token.clone(),
325        );
326
327        shutdown_token.cancel();
328        let error = context
329            .sleep_or_shutdown(Duration::from_mins(1))
330            .await
331            .expect_err("sleep should stop for shutdown");
332
333        assert!(error.is_task_worker_shutdown_requested());
334    }
335
336    #[test]
337    fn lease_timing_helpers_match_yggdrasil_and_drive_policy() {
338        let now = Utc::now();
339
340        assert_eq!(
341            task_lease_expires_at(now, 60),
342            now + chrono::Duration::seconds(60)
343        );
344        assert_eq!(
345            task_lease_expires_at(now, 0),
346            now + chrono::Duration::seconds(1)
347        );
348        assert_eq!(task_lease_renewal_timeout(60, 10), Duration::from_secs(50));
349        assert_eq!(task_lease_renewal_timeout(1, 10), Duration::from_secs(1));
350    }
351}
aster_forge_tasks/lease.rs

aster_forge_tasks/
lease.rs
