1use chrono::{DateTime, Utc};
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13use uuid::Uuid;
14
15#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27#[serde(tag = "phase", rename_all = "snake_case")]
28pub enum WorkloadPhase<W, E, C, T> {
29 Initial,
31 Warming(W),
33 Executing(E),
35 Contracting(C),
37 Terminal(T),
39}
40
41impl<W, E, C, T> WorkloadPhase<W, E, C, T> {
42 pub fn phase_name(&self) -> &'static str {
43 match self {
44 Self::Initial => "initial",
45 Self::Warming(_) => "warming",
46 Self::Executing(_) => "executing",
47 Self::Contracting(_) => "contracting",
48 Self::Terminal(_) => "terminal",
49 }
50 }
51
52 pub fn is_terminal(&self) -> bool {
53 matches!(self, Self::Terminal(_))
54 }
55
56 pub fn is_active(&self) -> bool {
57 matches!(self, Self::Warming(_) | Self::Executing(_))
58 }
59
60 pub fn is_initial(&self) -> bool {
61 matches!(self, Self::Initial)
62 }
63}
64
65pub fn is_valid_transition(from: &str, to: &str) -> bool {
67 matches!(
68 (from, to),
69 ("initial", "warming")
70 | ("warming", "executing")
71 | ("warming", "contracting")
72 | ("warming", "terminal")
73 | ("executing", "contracting")
74 | ("contracting", "terminal")
75 | ("contracting", "initial")
76 )
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
83#[serde(rename_all = "snake_case")]
84pub enum ContractReason {
85 Stopped,
86 Superseded { new_version: u64 },
87 NodeDrain,
88 ScaleDown,
89 HealthFailure,
90 ResourcePressure,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
95#[serde(rename_all = "snake_case")]
96pub enum Outcome {
97 Success,
98 Failed,
99 Lost,
100 Cancelled,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
105#[serde(rename_all = "snake_case")]
106pub enum HealthStatus {
107 #[default]
108 Unknown,
109 Passing,
110 Warning {
111 message: String,
112 },
113 Critical {
114 message: String,
115 },
116}
117
118#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
120#[serde(rename_all = "snake_case")]
121pub enum DesiredPhase {
122 Active,
123 Stopped { reason: ContractReason },
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
129pub struct TaskWarmProgress {
130 pub fetch_progress: f64,
131 pub deps_resolved: bool,
132 pub port_allocated: bool,
133 pub warmup_checks_passed: u32,
134 pub warmup_checks_required: u32,
135}
136
137impl Default for TaskWarmProgress {
138 fn default() -> Self {
139 Self {
140 fetch_progress: 0.0,
141 deps_resolved: false,
142 port_allocated: false,
143 warmup_checks_passed: 0,
144 warmup_checks_required: 1,
145 }
146 }
147}
148
149#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
150pub struct TaskExecuteDetail {
151 pub pid: Option<u32>,
152 pub container_id: Option<String>,
153 pub health: HealthStatus,
154 pub started_at: DateTime<Utc>,
155 pub health_check_epoch: u64,
156}
157
158#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
159pub struct TaskContractDetail {
160 pub reason: ContractReason,
161 pub signal_sent_at: Option<DateTime<Utc>>,
162 pub grace_period_secs: u64,
163 pub drain_connections: bool,
164}
165
166#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
167pub struct TaskTerminalDetail {
168 pub outcome: Outcome,
169 pub exit_code: Option<i32>,
170 pub finished_at: DateTime<Utc>,
171 pub restarts: u32,
172}
173
174pub type TaskPhase =
176 WorkloadPhase<TaskWarmProgress, TaskExecuteDetail, TaskContractDetail, TaskTerminalDetail>;
177
178#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
181pub struct AllocWarmProgress {
182 pub secrets_resolved: bool,
183 pub volumes_mounted: bool,
184 pub task_progress: HashMap<String, TaskWarmProgress>,
185 #[serde(default)]
187 pub network_identity_assigned: bool,
188 #[serde(default)]
190 pub endpoint_registered: bool,
191}
192
193impl Default for AllocWarmProgress {
194 fn default() -> Self {
195 Self {
196 secrets_resolved: false,
197 volumes_mounted: false,
198 task_progress: HashMap::new(),
199 network_identity_assigned: false,
200 endpoint_registered: false,
201 }
202 }
203}
204
205#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
206pub struct AllocExecuteDetail {
207 pub registered_in_catalog: bool,
208 pub health: HealthStatus,
209 pub task_states: HashMap<String, TaskPhase>,
210}
211
212#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
213pub struct AllocContractDetail {
214 pub reason: ContractReason,
215 pub deregistered_from_catalog: bool,
216 pub task_states: HashMap<String, TaskPhase>,
217}
218
219#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
220pub struct AllocTerminalDetail {
221 pub outcome: Outcome,
222 pub finished_at: DateTime<Utc>,
223}
224
225pub type AllocationPhase =
227 WorkloadPhase<AllocWarmProgress, AllocExecuteDetail, AllocContractDetail, AllocTerminalDetail>;
228
229#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
232pub struct NodeWarmProgress {
233 pub raft_joined: bool,
234 pub gossip_converged: bool,
235 pub drivers_ready: Vec<String>,
236 #[serde(default)]
238 pub wireguard_tunnel_up: bool,
239 #[serde(default)]
241 pub mesh_peers_connected: u32,
242}
243
244impl Default for NodeWarmProgress {
245 fn default() -> Self {
246 Self {
247 raft_joined: false,
248 gossip_converged: false,
249 drivers_ready: Vec::new(),
250 wireguard_tunnel_up: false,
251 mesh_peers_connected: 0,
252 }
253 }
254}
255
256#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
257pub struct NodeExecuteDetail {
258 pub eligible: bool,
259 pub allocation_count: u32,
260 pub last_heartbeat: DateTime<Utc>,
261}
262
263#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
264pub struct NodeContractDetail {
265 pub reason: ContractReason,
266 pub draining_allocations: Vec<Uuid>,
267}
268
269#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
270pub struct NodeTerminalDetail {
271 pub departed_at: DateTime<Utc>,
272 pub reason: ContractReason,
273}
274
275pub type NodePhase =
277 WorkloadPhase<NodeWarmProgress, NodeExecuteDetail, NodeContractDetail, NodeTerminalDetail>;
278
279#[derive(Debug, Clone, Serialize, Deserialize)]
283pub struct DesiredAllocationState {
284 pub alloc_id: Uuid,
285 pub job_id: String,
286 pub group_name: String,
287 pub node_id: String,
288 pub job_version: u64,
289 pub desired_phase: DesiredPhase,
290 pub generation: u64,
291}
292
293#[derive(Debug, Clone, Serialize, Deserialize)]
295pub struct ObservedAllocationState {
296 pub alloc_id: Uuid,
297 pub node_id: String,
298 pub phase: AllocationPhase,
299 pub observed_at: DateTime<Utc>,
300 pub observation_seq: u64,
301}
302
303use super::allocation::{AllocationState, TaskRunState};
306
307impl From<AllocationState> for AllocationPhase {
308 fn from(state: AllocationState) -> Self {
309 match state {
310 AllocationState::Pending => AllocationPhase::Initial,
311 AllocationState::Running => AllocationPhase::Executing(AllocExecuteDetail {
312 registered_in_catalog: false,
313 health: HealthStatus::Unknown,
314 task_states: HashMap::new(),
315 }),
316 AllocationState::Complete => AllocationPhase::Terminal(AllocTerminalDetail {
317 outcome: Outcome::Success,
318 finished_at: Utc::now(),
319 }),
320 AllocationState::Failed => AllocationPhase::Terminal(AllocTerminalDetail {
321 outcome: Outcome::Failed,
322 finished_at: Utc::now(),
323 }),
324 AllocationState::Lost => AllocationPhase::Terminal(AllocTerminalDetail {
325 outcome: Outcome::Lost,
326 finished_at: Utc::now(),
327 }),
328 }
329 }
330}
331
332impl From<TaskRunState> for TaskPhase {
333 fn from(state: TaskRunState) -> Self {
334 match state {
335 TaskRunState::Pending => TaskPhase::Initial,
336 TaskRunState::Running => TaskPhase::Executing(TaskExecuteDetail {
337 pid: None,
338 container_id: None,
339 health: HealthStatus::Unknown,
340 started_at: Utc::now(),
341 health_check_epoch: 0,
342 }),
343 TaskRunState::Dead => TaskPhase::Terminal(TaskTerminalDetail {
344 outcome: Outcome::Failed,
345 exit_code: None,
346 finished_at: Utc::now(),
347 restarts: 0,
348 }),
349 }
350 }
351}
352
353#[cfg(test)]
354mod tests {
355 use super::*;
356
357 #[test]
358 fn test_phase_names() {
359 let phase: TaskPhase = TaskPhase::Initial;
360 assert_eq!(phase.phase_name(), "initial");
361
362 let phase: TaskPhase = TaskPhase::Warming(TaskWarmProgress::default());
363 assert_eq!(phase.phase_name(), "warming");
364 assert!(phase.is_active());
365
366 let phase: TaskPhase = TaskPhase::Terminal(TaskTerminalDetail {
367 outcome: Outcome::Success,
368 exit_code: Some(0),
369 finished_at: Utc::now(),
370 restarts: 0,
371 });
372 assert!(phase.is_terminal());
373 assert!(!phase.is_active());
374 }
375
376 #[test]
377 fn test_valid_transitions() {
378 assert!(is_valid_transition("initial", "warming"));
379 assert!(is_valid_transition("warming", "executing"));
380 assert!(is_valid_transition("executing", "contracting"));
381 assert!(is_valid_transition("contracting", "terminal"));
382 assert!(is_valid_transition("contracting", "initial")); assert!(is_valid_transition("warming", "terminal")); assert!(!is_valid_transition("initial", "executing")); assert!(!is_valid_transition("executing", "warming")); assert!(!is_valid_transition("terminal", "initial")); assert!(!is_valid_transition("initial", "contracting")); }
390
391 #[test]
392 fn test_allocation_phase_from_legacy() {
393 let phase: AllocationPhase = AllocationState::Pending.into();
394 assert!(phase.is_initial());
395
396 let phase: AllocationPhase = AllocationState::Running.into();
397 assert_eq!(phase.phase_name(), "executing");
398
399 let phase: AllocationPhase = AllocationState::Failed.into();
400 assert!(phase.is_terminal());
401 }
402
403 #[test]
404 fn test_task_phase_from_legacy() {
405 let phase: TaskPhase = TaskRunState::Pending.into();
406 assert!(phase.is_initial());
407
408 let phase: TaskPhase = TaskRunState::Running.into();
409 assert!(phase.is_active());
410
411 let phase: TaskPhase = TaskRunState::Dead.into();
412 assert!(phase.is_terminal());
413 }
414
415 #[test]
416 fn test_serde_roundtrip() {
417 let phase = AllocationPhase::Warming(AllocWarmProgress {
418 secrets_resolved: true,
419 volumes_mounted: false,
420 network_identity_assigned: true,
421 endpoint_registered: false,
422 task_progress: HashMap::from([(
423 "web".to_string(),
424 TaskWarmProgress {
425 fetch_progress: 0.75,
426 deps_resolved: true,
427 port_allocated: true,
428 warmup_checks_passed: 2,
429 warmup_checks_required: 3,
430 },
431 )]),
432 });
433
434 let json = serde_json::to_string(&phase).unwrap();
435 let back: AllocationPhase = serde_json::from_str(&json).unwrap();
436 assert_eq!(phase, back);
437 }
438
439 #[test]
440 fn test_desired_phase_serde() {
441 let desired = DesiredPhase::Stopped {
442 reason: ContractReason::ScaleDown,
443 };
444 let json = serde_json::to_string(&desired).unwrap();
445 let back: DesiredPhase = serde_json::from_str(&json).unwrap();
446 assert_eq!(desired, back);
447 }
448}