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

aster_forge_tasks/
lease.rs
