termivar_scanner/distributed/
lease.rs1use 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#[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#[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#[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#[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#[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}