termivar_scanner/distributed/
queue.rs

1use std::sync::{Arc, Mutex};
2
3use super::coordinator::{next_counter, CoordinatorState};
4use super::lease::{CancellationProof, CompletionReceipt, FailureProof, ScanTask};
5use super::limits::{validate_limits, DistributedLimits};
6use super::model::{
7    validate_task_command_id, validate_task_spec, StateSnapshot, TaskPriority, TaskSpec,
8    TaskStatus, Transition,
9};
10use super::{lock_state, DistributedError};
11
12#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
13pub(super) struct QueueKey {
14    inverted_priority: u8,
15    enqueue_ordinal: u64,
16    pub(super) task_id: String,
17}
18
19impl QueueKey {
20    pub(super) fn new(priority: TaskPriority, enqueue_ordinal: u64, task_id: String) -> Self {
21        Self {
22            inverted_priority: u8::MAX - priority as u8,
23            enqueue_ordinal,
24            task_id,
25        }
26    }
27}
28
29#[derive(Debug, Clone)]
30pub(super) struct TaskEntry {
31    pub(super) task: ScanTask,
32    pub(super) queue_key: Option<QueueKey>,
33    pub(super) completion: Option<CompletionReceipt>,
34    pub(super) cancellation: Option<CancellationProof>,
35    pub(super) failure: Option<FailureProof>,
36}
37
38/// Cloneable task facade sharing the pool's single state lock.
39#[derive(Clone)]
40pub struct TaskQueue {
41    pub(super) state: Arc<Mutex<CoordinatorState>>,
42}
43
44impl TaskQueue {
45    pub fn new() -> Self {
46        Self {
47            state: Arc::new(Mutex::new(CoordinatorState::new(
48                DistributedLimits::default(),
49            ))),
50        }
51    }
52
53    pub fn with_limits(limits: DistributedLimits) -> Result<Self, DistributedError> {
54        validate_limits(limits)?;
55        Ok(Self {
56            state: Arc::new(Mutex::new(CoordinatorState::new(limits))),
57        })
58    }
59
60    /// Admit one fresh task at an exact coordinator revision.
61    pub fn enqueue(
62        &self,
63        expected_revision: u64,
64        now_secs: u64,
65        spec: TaskSpec,
66    ) -> Result<Transition<ScanTask>, DistributedError> {
67        validate_task_spec(&spec)?;
68        let mut state = lock_state(&self.state)?;
69        let revision = state.preflight_command(expected_revision, now_secs)?;
70        if state.tasks.contains_key(&spec.task_id) {
71            return Err(DistributedError::TaskAlreadyExists {
72                task_id: spec.task_id,
73            });
74        }
75        if state.tasks.len() >= state.limits.max_task_records {
76            return Err(DistributedError::TaskRecordCapacityReached {
77                limit: state.limits.max_task_records,
78            });
79        }
80        if state.active_tasks >= state.limits.max_active_tasks {
81            return Err(DistributedError::ActiveTaskCapacityReached {
82                limit: state.limits.max_active_tasks,
83            });
84        }
85        if state.queue.len() >= state.limits.max_queued_tasks {
86            return Err(DistributedError::QueuedTaskCapacityReached {
87                limit: state.limits.max_queued_tasks,
88            });
89        }
90        if state
91            .terminal_tasks
92            .checked_add(state.active_tasks)
93            .and_then(|value| value.checked_add(1))
94            .is_none_or(|value| value > state.limits.max_terminal_tasks)
95        {
96            return Err(DistributedError::TerminalCapacityReserved {
97                limit: state.limits.max_terminal_tasks,
98            });
99        }
100        let next_ordinal = next_counter(state.next_enqueue_ordinal, "enqueue_ordinal")?;
101        let task = ScanTask {
102            task_id: spec.task_id,
103            scan_id: spec.scan_id,
104            target_ref: spec.target_ref,
105            phases: spec.phases,
106            status: TaskStatus::Queued,
107            priority: spec.priority,
108            created_at: now_secs,
109            started_at: None,
110            completed_at: None,
111            retry_count: 0,
112            attempt: 0,
113            task_generation: 0,
114            record_version: 1,
115            assigned_to: None,
116            lease: None,
117        };
118        let key = QueueKey::new(
119            task.priority,
120            state.next_enqueue_ordinal,
121            task.task_id.clone(),
122        );
123        state.next_enqueue_ordinal = next_ordinal;
124        state.queue.insert(key.clone());
125        state.tasks.insert(
126            task.task_id.clone(),
127            TaskEntry {
128                task: task.clone(),
129                queue_key: Some(key),
130                completion: None,
131                cancellation: None,
132                failure: None,
133            },
134        );
135        state.active_tasks += 1;
136        state.commit(revision, now_secs);
137        Ok(Transition {
138            revision,
139            value: task,
140        })
141    }
142
143    /// Peek without consuming or mutating the task record.
144    pub fn peek_next(&self) -> Result<Option<ScanTask>, DistributedError> {
145        let state = lock_state(&self.state)?;
146        let Some(task_id) = state.next_task_id() else {
147            return Ok(None);
148        };
149        state
150            .tasks
151            .get(&task_id)
152            .map(|entry| Some(entry.task.clone()))
153            .ok_or(DistributedError::StateInvariant {
154                reason: "queue key references a missing task",
155            })
156    }
157
158    pub fn get_task(&self, task_id: &str) -> Result<Option<ScanTask>, DistributedError> {
159        validate_task_command_id(task_id)?;
160        Ok(lock_state(&self.state)?
161            .tasks
162            .get(task_id)
163            .map(|entry| entry.task.clone()))
164    }
165
166    /// Return tasks in stable task-ID order.
167    pub fn tasks(&self) -> Result<Vec<ScanTask>, DistributedError> {
168        Ok(lock_state(&self.state)?
169            .tasks
170            .values()
171            .map(|entry| entry.task.clone())
172            .collect())
173    }
174
175    pub fn snapshot(&self) -> Result<StateSnapshot, DistributedError> {
176        Ok(lock_state(&self.state)?.snapshot())
177    }
178}
179
180impl Default for TaskQueue {
181    fn default() -> Self {
182        Self::new()
183    }
184}
termivar_scanner/distributed/queue.rs

termivar_scanner/distributed/
queue.rs
