Skip to main content

termivar_scanner/distributed/
lease.rs

1use super::coordinator::CoordinatorState;
2use super::model::{TaskPriority, TaskStatus};
3use super::queue::TaskEntry;
4use super::worker::{preflight_worker_release, release_worker};
5use super::DistributedError;
6
7/// Exact logical lease fence for one worker generation and task attempt.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct TaskLease {
10    pub(super) task_id: String,
11    pub(super) worker_id: String,
12    pub(super) worker_generation: u64,
13    pub(super) task_generation: u64,
14    pub(super) attempt: u32,
15    pub(super) lease_id: u64,
16    pub(super) acquired_at: u64,
17    pub(super) expires_at: u64,
18}
19
20impl TaskLease {
21    pub fn task_id(&self) -> &str {
22        &self.task_id
23    }
24    pub fn worker_id(&self) -> &str {
25        &self.worker_id
26    }
27    pub fn worker_generation(&self) -> u64 {
28        self.worker_generation
29    }
30    pub fn task_generation(&self) -> u64 {
31        self.task_generation
32    }
33    pub fn attempt(&self) -> u32 {
34        self.attempt
35    }
36    pub fn lease_id(&self) -> u64 {
37        self.lease_id
38    }
39    pub fn acquired_at(&self) -> u64 {
40        self.acquired_at
41    }
42    pub fn expires_at(&self) -> u64 {
43        self.expires_at
44    }
45}
46
47/// Optimistic ownership fence for a queued task.
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct QueuedTaskFence {
50    task_id: String,
51    task_generation: u64,
52    record_version: u64,
53}
54
55impl QueuedTaskFence {
56    pub fn task_id(&self) -> &str {
57        &self.task_id
58    }
59    pub fn task_generation(&self) -> u64 {
60        self.task_generation
61    }
62    pub fn record_version(&self) -> u64 {
63        self.record_version
64    }
65}
66
67/// Exact logical cancellation fence for queued or leased work.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub enum TaskOwnership {
70    Queued(QueuedTaskFence),
71    Leased(TaskLease),
72}
73
74impl TaskOwnership {
75    pub(super) fn task_id(&self) -> &str {
76        match self {
77            Self::Queued(fence) => &fence.task_id,
78            Self::Leased(lease) => &lease.task_id,
79        }
80    }
81}
82
83/// Versioned task snapshot.
84#[derive(Clone, PartialEq, Eq)]
85pub struct ScanTask {
86    pub(super) task_id: String,
87    pub(super) scan_id: String,
88    pub(super) target_ref: String,
89    pub(super) phases: Vec<u8>,
90    pub(super) status: TaskStatus,
91    pub(super) priority: TaskPriority,
92    pub(super) created_at: u64,
93    pub(super) started_at: Option<u64>,
94    pub(super) completed_at: Option<u64>,
95    pub(super) retry_count: u32,
96    pub(super) attempt: u32,
97    pub(super) task_generation: u64,
98    pub(super) record_version: u64,
99    pub(super) assigned_to: Option<String>,
100    pub(super) lease: Option<TaskLease>,
101}
102
103impl std::fmt::Debug for ScanTask {
104    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        formatter
106            .debug_struct("ScanTask")
107            .field("task_id", &self.task_id)
108            .field("scan_id", &self.scan_id)
109            .field("target_ref", &"<opaque>")
110            .field("phases", &self.phases)
111            .field("status", &self.status)
112            .field("priority", &self.priority)
113            .field("created_at", &self.created_at)
114            .field("started_at", &self.started_at)
115            .field("completed_at", &self.completed_at)
116            .field("retry_count", &self.retry_count)
117            .field("attempt", &self.attempt)
118            .field("task_generation", &self.task_generation)
119            .field("record_version", &self.record_version)
120            .field("assigned_to", &self.assigned_to)
121            .field("lease", &self.lease)
122            .finish()
123    }
124}
125
126impl ScanTask {
127    pub fn task_id(&self) -> &str {
128        &self.task_id
129    }
130
131    pub fn scan_id(&self) -> &str {
132        &self.scan_id
133    }
134
135    pub fn target_ref(&self) -> &str {
136        &self.target_ref
137    }
138
139    pub fn phases(&self) -> &[u8] {
140        &self.phases
141    }
142
143    pub fn status(&self) -> TaskStatus {
144        self.status
145    }
146
147    pub fn priority(&self) -> TaskPriority {
148        self.priority
149    }
150
151    pub fn created_at(&self) -> u64 {
152        self.created_at
153    }
154
155    pub fn started_at(&self) -> Option<u64> {
156        self.started_at
157    }
158
159    pub fn completed_at(&self) -> Option<u64> {
160        self.completed_at
161    }
162
163    pub fn retry_count(&self) -> u32 {
164        self.retry_count
165    }
166
167    pub fn attempt(&self) -> u32 {
168        self.attempt
169    }
170
171    pub fn task_generation(&self) -> u64 {
172        self.task_generation
173    }
174
175    pub fn record_version(&self) -> u64 {
176        self.record_version
177    }
178
179    pub fn assigned_to(&self) -> Option<&str> {
180        self.assigned_to.as_deref()
181    }
182
183    pub fn lease(&self) -> Option<&TaskLease> {
184        self.lease.as_ref()
185    }
186
187    pub fn ownership(&self) -> Option<TaskOwnership> {
188        match self.status {
189            TaskStatus::Queued => Some(TaskOwnership::Queued(QueuedTaskFence {
190                task_id: self.task_id.clone(),
191                task_generation: self.task_generation,
192                record_version: self.record_version,
193            })),
194            TaskStatus::Leased | TaskStatus::Running => {
195                self.lease.clone().map(TaskOwnership::Leased)
196            },
197            _ => None,
198        }
199    }
200
201    pub fn age_secs(&self, now_secs: u64) -> u64 {
202        now_secs.saturating_sub(self.created_at)
203    }
204}
205
206#[derive(Debug, Clone, PartialEq, Eq)]
207pub(super) enum CancellationProof {
208    Queued(QueuedTaskFence),
209    Leased(TaskLease),
210}
211
212#[derive(Debug, Clone, PartialEq, Eq)]
213pub(super) struct FailureProof {
214    pub(super) lease: TaskLease,
215    pub(super) outcome: FailureOutcome,
216}
217
218impl CancellationProof {
219    pub(super) fn matches(&self, ownership: &TaskOwnership) -> bool {
220        matches!(
221            (self, ownership),
222            (Self::Queued(left), TaskOwnership::Queued(right)) if left == right
223        ) || matches!(
224            (self, ownership),
225            (Self::Leased(left), TaskOwnership::Leased(right)) if left == right
226        )
227    }
228}
229
230/// Exact receipt for a terminal completion. Fields are intentionally private.
231#[derive(Debug, Clone, PartialEq, Eq)]
232pub struct CompletionReceipt {
233    pub(super) task_id: String,
234    pub(super) task_generation: u64,
235    pub(super) attempt: u32,
236    pub(super) lease_id: u64,
237    pub(super) worker_id: String,
238    pub(super) worker_generation: u64,
239    pub(super) record_version: u64,
240}
241
242impl CompletionReceipt {
243    pub fn task_id(&self) -> &str {
244        &self.task_id
245    }
246    pub fn task_generation(&self) -> u64 {
247        self.task_generation
248    }
249    pub fn attempt(&self) -> u32 {
250        self.attempt
251    }
252    pub fn lease_id(&self) -> u64 {
253        self.lease_id
254    }
255    pub fn worker_id(&self) -> &str {
256        &self.worker_id
257    }
258    pub fn worker_generation(&self) -> u64 {
259        self.worker_generation
260    }
261    pub fn record_version(&self) -> u64 {
262        self.record_version
263    }
264
265    pub(super) fn matches_lease(&self, lease: &TaskLease) -> bool {
266        self.task_id == lease.task_id
267            && self.task_generation == lease.task_generation
268            && self.attempt == lease.attempt
269            && self.lease_id == lease.lease_id
270            && self.worker_id == lease.worker_id
271            && self.worker_generation == lease.worker_generation
272    }
273}
274
275#[derive(Debug, Clone, PartialEq, Eq)]
276pub enum StartOutcome {
277    Started { record_version: u64 },
278    AlreadyRunning { record_version: u64 },
279}
280
281#[derive(Debug, Clone, PartialEq, Eq)]
282pub enum CompletionOutcome {
283    Completed(CompletionReceipt),
284    AlreadyCompleted(CompletionReceipt),
285}
286
287impl CompletionOutcome {
288    pub fn receipt(&self) -> &CompletionReceipt {
289        match self {
290            Self::Completed(receipt) | Self::AlreadyCompleted(receipt) => receipt,
291        }
292    }
293}
294
295#[derive(Debug, Clone, PartialEq, Eq)]
296pub enum CancellationOutcome {
297    Cancelled { record_version: u64 },
298    AlreadyCancelled { record_version: u64 },
299}
300
301#[derive(Debug, Clone, PartialEq, Eq)]
302pub enum FailureOutcome {
303    Requeued {
304        task_generation: u64,
305        retry_count: u32,
306        record_version: u64,
307    },
308    RetryExhausted {
309        retry_count: u32,
310        record_version: u64,
311    },
312}
313
314pub(super) fn ensure_queue_entry(
315    state: &CoordinatorState,
316    entry: &TaskEntry,
317) -> Result<(), DistributedError> {
318    match entry.queue_key.as_ref() {
319        Some(key) if state.queue.contains(key) && key.task_id == entry.task.task_id => Ok(()),
320        _ => Err(DistributedError::StateInvariant {
321            reason: "queued task does not have exactly one queue key",
322        }),
323    }
324}
325
326pub(super) fn ensure_current_lease(
327    entry: &TaskEntry,
328    lease: &TaskLease,
329    now_secs: u64,
330) -> Result<(), DistributedError> {
331    if entry.task.lease.as_ref() != Some(lease)
332        || entry.task.task_generation != lease.task_generation
333        || entry.task.attempt != lease.attempt
334        || entry.task.assigned_to.as_deref() != Some(lease.worker_id.as_str())
335    {
336        return Err(DistributedError::StaleOwnership {
337            task_id: lease.task_id.clone(),
338        });
339    }
340    if now_secs >= lease.expires_at {
341        return Err(DistributedError::LeaseExpired {
342            task_id: lease.task_id.clone(),
343        });
344    }
345    Ok(())
346}
347
348pub(super) fn validate_cancellation_ownership(
349    entry: &TaskEntry,
350    ownership: &TaskOwnership,
351    now_secs: u64,
352) -> Result<(), DistributedError> {
353    match ownership {
354        TaskOwnership::Queued(fence) => {
355            if entry.task.status != TaskStatus::Queued
356                || entry.task.task_id != fence.task_id
357                || entry.task.task_generation != fence.task_generation
358                || entry.task.record_version != fence.record_version
359            {
360                return Err(DistributedError::StaleOwnership {
361                    task_id: fence.task_id.clone(),
362                });
363            }
364        },
365        TaskOwnership::Leased(lease) => ensure_current_lease(entry, lease, now_secs)?,
366    }
367    Ok(())
368}
369
370pub(super) fn terminalize_queued_task(
371    state: &mut CoordinatorState,
372    task_id: &str,
373    terminal_status: TaskStatus,
374    now_secs: u64,
375    record_version: u64,
376) -> Result<(), DistributedError> {
377    let entry = state
378        .tasks
379        .get(task_id)
380        .ok_or(DistributedError::StateInvariant {
381            reason: "queued terminal task disappeared",
382        })?;
383    if entry.task.status != TaskStatus::Queued || terminal_status.is_active() {
384        return Err(DistributedError::StateInvariant {
385            reason: "invalid queued terminalization",
386        });
387    }
388    let key = entry
389        .queue_key
390        .clone()
391        .ok_or(DistributedError::StateInvariant {
392            reason: "queued terminal task has no queue key",
393        })?;
394    if !state.queue.remove(&key) {
395        return Err(DistributedError::StateInvariant {
396            reason: "queued terminal key disappeared",
397        });
398    }
399    let entry = state
400        .tasks
401        .get_mut(task_id)
402        .ok_or(DistributedError::StateInvariant {
403            reason: "queued terminal task disappeared during mutation",
404        })?;
405    entry.queue_key = None;
406    entry.task.status = terminal_status;
407    entry.task.completed_at = Some(now_secs);
408    entry.task.record_version = record_version;
409    state.active_tasks -= 1;
410    state.terminal_tasks += 1;
411    Ok(())
412}
413
414pub(super) fn terminalize_leased_task(
415    state: &mut CoordinatorState,
416    lease: &TaskLease,
417    terminal_status: TaskStatus,
418    now_secs: u64,
419    record_version: u64,
420) -> Result<(), DistributedError> {
421    if terminal_status.is_active() {
422        return Err(DistributedError::StateInvariant {
423            reason: "leased terminalization received active status",
424        });
425    }
426    preflight_worker_release(state, lease)?;
427    release_worker(state, lease)?;
428    let entry = state
429        .tasks
430        .get_mut(&lease.task_id)
431        .ok_or(DistributedError::StateInvariant {
432            reason: "leased terminal task disappeared",
433        })?;
434    entry.task.status = terminal_status;
435    entry.task.assigned_to = None;
436    entry.task.lease = None;
437    entry.task.completed_at = Some(now_secs);
438    entry.task.record_version = record_version;
439    state.active_tasks -= 1;
440    state.terminal_tasks += 1;
441    Ok(())
442}