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#[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#[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#[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#[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#[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}