Skip to main content

termivar_scanner/distributed/
worker.rs

1use std::cmp::Ordering;
2use std::collections::BTreeSet;
3
4use super::coordinator::{next_counter, next_u32, CoordinatorState};
5use super::lease::{ensure_queue_entry, TaskLease};
6use super::limits::{MAX_WORKER_CAPACITY, MAX_WORKER_TAGS, UTILIZATION_BASIS_POINTS};
7use super::model::{validate_identifier, TaskStatus};
8use super::DistributedError;
9
10/// Observational worker metadata tags. Ordering is stable and non-randomized;
11/// the coordinator does not use these tags for eligibility or affinity.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
13pub enum WorkerTag {
14    Linux,
15    Windows,
16    Gpu,
17    Internal,
18    External,
19}
20
21impl WorkerTag {
22    pub fn as_str(self) -> &'static str {
23        match self {
24            Self::Linux => "linux",
25            Self::Windows => "windows",
26            Self::Gpu => "gpu",
27            Self::Internal => "internal",
28            Self::External => "external",
29        }
30    }
31}
32
33/// Worker eligibility state.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum WorkerStatus {
36    Healthy,
37    Busy,
38    Degraded,
39    Offline,
40}
41
42impl WorkerStatus {
43    pub fn as_str(self) -> &'static str {
44        match self {
45            Self::Healthy => "healthy",
46            Self::Busy => "busy",
47            Self::Degraded => "degraded",
48            Self::Offline => "offline",
49        }
50    }
51}
52
53/// Caller-provided worker admission record.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct WorkerSpec {
56    pub worker_id: String,
57    pub capacity: u32,
58    pub cpu_basis_points: u16,
59    pub memory_basis_points: u16,
60    pub network_basis_points: u16,
61    pub tags: BTreeSet<WorkerTag>,
62}
63
64impl WorkerSpec {
65    pub fn new(worker_id: impl Into<String>, capacity: u32) -> Self {
66        Self {
67            worker_id: worker_id.into(),
68            capacity,
69            cpu_basis_points: 0,
70            memory_basis_points: 0,
71            network_basis_points: 0,
72            tags: BTreeSet::new(),
73        }
74    }
75}
76
77/// Caller-supplied heartbeat and resource observation.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub struct WorkerObservation {
80    pub status: WorkerStatus,
81    pub cpu_basis_points: u16,
82    pub memory_basis_points: u16,
83    pub network_basis_points: u16,
84}
85
86impl Default for WorkerObservation {
87    fn default() -> Self {
88        Self {
89            status: WorkerStatus::Healthy,
90            cpu_basis_points: 0,
91            memory_basis_points: 0,
92            network_basis_points: 0,
93        }
94    }
95}
96
97/// Coordinator-owned worker snapshot.
98#[derive(Debug, Clone, PartialEq, Eq)]
99pub struct WorkerNode {
100    pub(super) worker_id: String,
101    pub(super) generation: u64,
102    pub(super) status: WorkerStatus,
103    pub(super) capacity: u32,
104    pub(super) current_tasks: u32,
105    pub(super) completed_tasks: u64,
106    pub(super) last_heartbeat: u64,
107    pub(super) cpu_basis_points: u16,
108    pub(super) memory_basis_points: u16,
109    pub(super) network_basis_points: u16,
110    pub(super) tags: BTreeSet<WorkerTag>,
111}
112
113impl WorkerNode {
114    pub fn worker_id(&self) -> &str {
115        &self.worker_id
116    }
117
118    pub fn generation(&self) -> u64 {
119        self.generation
120    }
121
122    pub fn status(&self) -> WorkerStatus {
123        self.status
124    }
125
126    pub fn capacity(&self) -> u32 {
127        self.capacity
128    }
129
130    pub fn current_tasks(&self) -> u32 {
131        self.current_tasks
132    }
133
134    pub fn completed_tasks(&self) -> u64 {
135        self.completed_tasks
136    }
137
138    pub fn last_heartbeat(&self) -> u64 {
139        self.last_heartbeat
140    }
141
142    pub fn cpu_basis_points(&self) -> u16 {
143        self.cpu_basis_points
144    }
145
146    pub fn memory_basis_points(&self) -> u16 {
147        self.memory_basis_points
148    }
149
150    pub fn network_basis_points(&self) -> u16 {
151        self.network_basis_points
152    }
153
154    pub fn tags(&self) -> &BTreeSet<WorkerTag> {
155        &self.tags
156    }
157
158    pub fn effective_capacity(&self) -> u32 {
159        if self.status != WorkerStatus::Healthy {
160            return 0;
161        }
162        let max_load = self
163            .cpu_basis_points
164            .max(self.memory_basis_points)
165            .max(self.network_basis_points) as u64;
166        let free = u64::from(UTILIZATION_BASIS_POINTS) - max_load;
167        let scaled = u64::from(self.capacity) * free;
168        let divisor = u64::from(UTILIZATION_BASIS_POINTS);
169        scaled.div_ceil(divisor) as u32
170    }
171
172    pub fn available_slots(&self) -> u32 {
173        self.effective_capacity().saturating_sub(self.current_tasks)
174    }
175
176    fn is_eligible(&self, now_secs: u64, heartbeat_timeout_secs: u64) -> bool {
177        self.status == WorkerStatus::Healthy
178            && now_secs.saturating_sub(self.last_heartbeat) <= heartbeat_timeout_secs
179            && self.available_slots() > 0
180    }
181
182    fn selection_key(&self, now_secs: u64) -> (u32, u64, u16, u16, u16) {
183        (
184            self.available_slots(),
185            u64::MAX - now_secs.saturating_sub(self.last_heartbeat),
186            UTILIZATION_BASIS_POINTS - self.cpu_basis_points,
187            UTILIZATION_BASIS_POINTS - self.memory_basis_points,
188            UTILIZATION_BASIS_POINTS - self.network_basis_points,
189        )
190    }
191}
192
193pub(super) fn validate_worker_spec(spec: &WorkerSpec) -> Result<(), DistributedError> {
194    validate_identifier(&spec.worker_id, "worker_id").map_err(|_| {
195        DistributedError::InvalidWorker {
196            reason: "worker_id is invalid",
197        }
198    })?;
199    if spec.capacity == 0 {
200        return Err(DistributedError::InvalidWorker {
201            reason: "capacity is zero",
202        });
203    }
204    if spec.capacity > MAX_WORKER_CAPACITY {
205        return Err(DistributedError::InvalidWorker {
206            reason: "capacity exceeds absolute maximum",
207        });
208    }
209    if spec.cpu_basis_points > UTILIZATION_BASIS_POINTS
210        || spec.memory_basis_points > UTILIZATION_BASIS_POINTS
211        || spec.network_basis_points > UTILIZATION_BASIS_POINTS
212    {
213        return Err(DistributedError::InvalidWorker {
214            reason: "utilization exceeds 10000 basis points",
215        });
216    }
217    if spec.tags.len() > MAX_WORKER_TAGS {
218        return Err(DistributedError::InvalidWorker {
219            reason: "too many worker tags",
220        });
221    }
222    Ok(())
223}
224
225pub(super) fn ensure_observation_time(
226    state: &CoordinatorState,
227    now_secs: u64,
228) -> Result<(), DistributedError> {
229    if now_secs < state.logical_time {
230        return Err(DistributedError::LogicalTimeRegression {
231            current: state.logical_time,
232            proposed: now_secs,
233        });
234    }
235    Ok(())
236}
237
238type WorkerSelectionKey = (u32, u64, u16, u16, u16);
239
240pub(super) fn best_worker_id(state: &CoordinatorState, now_secs: u64) -> Option<String> {
241    let mut best: Option<(&str, WorkerSelectionKey)> = None;
242    for (worker_id, worker) in &state.workers {
243        if !worker.is_eligible(now_secs, state.limits.heartbeat_timeout_secs) {
244            continue;
245        }
246        let key = worker.selection_key(now_secs);
247        let replace = match best {
248            None => true,
249            Some((best_id, best_key)) => match key.cmp(&best_key) {
250                Ordering::Greater => true,
251                Ordering::Equal => worker_id.as_str() < best_id,
252                Ordering::Less => false,
253            },
254        };
255        if replace {
256            best = Some((worker_id, key));
257        }
258    }
259    best.map(|(worker_id, _)| worker_id.to_owned())
260}
261
262pub(super) fn prepare_assignment(
263    state: &CoordinatorState,
264    task_id: &str,
265    worker_id: &str,
266    now_secs: u64,
267    lease_ttl_secs: u64,
268) -> Result<TaskLease, DistributedError> {
269    if lease_ttl_secs == 0 {
270        return Err(DistributedError::InvalidLimit {
271            name: "lease_ttl_secs",
272        });
273    }
274    if lease_ttl_secs > state.limits.max_lease_ttl_secs {
275        return Err(DistributedError::InvalidLimitRelationship {
276            reason: "lease_ttl_secs exceeds configured maximum",
277        });
278    }
279    let entry = state
280        .tasks
281        .get(task_id)
282        .ok_or_else(|| DistributedError::TaskNotFound {
283            task_id: task_id.to_owned(),
284        })?;
285    if entry.task.status != TaskStatus::Queued {
286        return Err(DistributedError::TaskNotQueued {
287            task_id: task_id.to_owned(),
288            status: entry.task.status,
289        });
290    }
291    ensure_queue_entry(state, entry)?;
292    let worker = state
293        .workers
294        .get(worker_id)
295        .ok_or_else(|| DistributedError::WorkerNotFound {
296            worker_id: worker_id.to_owned(),
297        })?;
298    if worker.status != WorkerStatus::Healthy
299        || now_secs.saturating_sub(worker.last_heartbeat) > state.limits.heartbeat_timeout_secs
300    {
301        return Err(DistributedError::WorkerUnavailable {
302            worker_id: worker_id.to_owned(),
303        });
304    }
305    if worker.available_slots() == 0 {
306        return Err(DistributedError::WorkerAtCapacity {
307            worker_id: worker_id.to_owned(),
308        });
309    }
310    next_u32(worker.current_tasks, "worker_current_tasks")?;
311    let attempt = next_u32(entry.task.attempt, "attempt")?;
312    next_counter(entry.task.record_version, "record_version")?;
313    let expires_at =
314        now_secs
315            .checked_add(lease_ttl_secs)
316            .ok_or(DistributedError::CounterExhausted {
317                counter: "lease_deadline",
318            })?;
319    next_counter(state.next_lease_id, "lease_id")?;
320    Ok(TaskLease {
321        task_id: task_id.to_owned(),
322        worker_id: worker_id.to_owned(),
323        worker_generation: worker.generation,
324        task_generation: entry.task.task_generation,
325        attempt,
326        lease_id: state.next_lease_id,
327        acquired_at: now_secs,
328        expires_at,
329    })
330}
331
332pub(super) fn apply_assignment(
333    state: &mut CoordinatorState,
334    lease: &TaskLease,
335) -> Result<(), DistributedError> {
336    let entry = state
337        .tasks
338        .get(&lease.task_id)
339        .ok_or(DistributedError::StateInvariant {
340            reason: "assignment task disappeared",
341        })?;
342    let queue_key = entry
343        .queue_key
344        .clone()
345        .ok_or(DistributedError::StateInvariant {
346            reason: "queued task has no queue key",
347        })?;
348    let record_version = next_counter(entry.task.record_version, "record_version")?;
349    if !state.queue.remove(&queue_key) {
350        return Err(DistributedError::StateInvariant {
351            reason: "assignment queue key disappeared",
352        });
353    }
354    let entry = state
355        .tasks
356        .get_mut(&lease.task_id)
357        .ok_or(DistributedError::StateInvariant {
358            reason: "assignment task disappeared during mutation",
359        })?;
360    entry.queue_key = None;
361    entry.task.status = TaskStatus::Leased;
362    entry.task.assigned_to = Some(lease.worker_id.clone());
363    entry.task.attempt = lease.attempt;
364    entry.task.lease = Some(lease.clone());
365    entry.task.record_version = record_version;
366    let worker =
367        state
368            .workers
369            .get_mut(&lease.worker_id)
370            .ok_or(DistributedError::StateInvariant {
371                reason: "assignment worker disappeared",
372            })?;
373    worker.current_tasks = next_u32(worker.current_tasks, "worker_current_tasks")?;
374    state.next_lease_id = next_counter(state.next_lease_id, "lease_id")?;
375    Ok(())
376}
377
378pub(super) fn preflight_worker_release(
379    state: &CoordinatorState,
380    lease: &TaskLease,
381) -> Result<(), DistributedError> {
382    let worker = state
383        .workers
384        .get(&lease.worker_id)
385        .ok_or(DistributedError::StateInvariant {
386            reason: "lease references a missing worker",
387        })?;
388    if worker.generation != lease.worker_generation || worker.current_tasks == 0 {
389        return Err(DistributedError::StateInvariant {
390            reason: "worker generation or active count disagrees with lease",
391        });
392    }
393    Ok(())
394}
395
396pub(super) fn release_worker(
397    state: &mut CoordinatorState,
398    lease: &TaskLease,
399) -> Result<(), DistributedError> {
400    let worker =
401        state
402            .workers
403            .get_mut(&lease.worker_id)
404            .ok_or(DistributedError::StateInvariant {
405                reason: "lease worker disappeared during release",
406            })?;
407    if worker.generation != lease.worker_generation || worker.current_tasks == 0 {
408        return Err(DistributedError::StateInvariant {
409            reason: "worker cannot release lease exactly once",
410        });
411    }
412    worker.current_tasks -= 1;
413    Ok(())
414}