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, Default, 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
193#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
194pub struct AllocExecuteDetail {
195 pub registered_in_catalog: bool,
196 pub health: HealthStatus,
197 pub task_states: HashMap<String, TaskPhase>,
198}
199
200#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
201pub struct AllocContractDetail {
202 pub reason: ContractReason,
203 pub deregistered_from_catalog: bool,
204 pub task_states: HashMap<String, TaskPhase>,
205}
206
207#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
208pub struct AllocTerminalDetail {
209 pub outcome: Outcome,
210 pub finished_at: DateTime<Utc>,
211}
212
213pub type AllocationPhase =
215 WorkloadPhase<AllocWarmProgress, AllocExecuteDetail, AllocContractDetail, AllocTerminalDetail>;
216
217#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
220pub struct NodeWarmProgress {
221 pub raft_joined: bool,
222 pub gossip_converged: bool,
223 pub drivers_ready: Vec<String>,
224 #[serde(default)]
226 pub wireguard_tunnel_up: bool,
227 #[serde(default)]
229 pub mesh_peers_connected: u32,
230}
231
232#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
233pub struct NodeExecuteDetail {
234 pub eligible: bool,
235 pub allocation_count: u32,
236 pub last_heartbeat: DateTime<Utc>,
237}
238
239#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
240pub struct NodeContractDetail {
241 pub reason: ContractReason,
242 pub draining_allocations: Vec<Uuid>,
243}
244
245#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
246pub struct NodeTerminalDetail {
247 pub departed_at: DateTime<Utc>,
248 pub reason: ContractReason,
249}
250
251pub type NodePhase =
253 WorkloadPhase<NodeWarmProgress, NodeExecuteDetail, NodeContractDetail, NodeTerminalDetail>;
254
255#[derive(Debug, Clone, Serialize, Deserialize)]
259pub struct DesiredAllocationState {
260 pub alloc_id: Uuid,
261 pub job_id: String,
262 pub group_name: String,
263 pub node_id: String,
264 pub job_version: u64,
265 pub desired_phase: DesiredPhase,
266 pub generation: u64,
267}
268
269#[derive(Debug, Clone, Serialize, Deserialize)]
271pub struct ObservedAllocationState {
272 pub alloc_id: Uuid,
273 pub node_id: String,
274 pub phase: AllocationPhase,
275 pub observed_at: DateTime<Utc>,
276 pub observation_seq: u64,
277}
278
279use super::allocation::{AllocationState, TaskRunState};
282
283impl From<AllocationState> for AllocationPhase {
284 fn from(state: AllocationState) -> Self {
285 match state {
286 AllocationState::Pending => AllocationPhase::Initial,
287 AllocationState::Running => AllocationPhase::Executing(AllocExecuteDetail {
288 registered_in_catalog: false,
289 health: HealthStatus::Unknown,
290 task_states: HashMap::new(),
291 }),
292 AllocationState::Complete => AllocationPhase::Terminal(AllocTerminalDetail {
293 outcome: Outcome::Success,
294 finished_at: Utc::now(),
295 }),
296 AllocationState::Failed => AllocationPhase::Terminal(AllocTerminalDetail {
297 outcome: Outcome::Failed,
298 finished_at: Utc::now(),
299 }),
300 AllocationState::Lost => AllocationPhase::Terminal(AllocTerminalDetail {
301 outcome: Outcome::Lost,
302 finished_at: Utc::now(),
303 }),
304 }
305 }
306}
307
308impl From<TaskRunState> for TaskPhase {
309 fn from(state: TaskRunState) -> Self {
310 match state {
311 TaskRunState::Pending => TaskPhase::Initial,
312 TaskRunState::Running => TaskPhase::Executing(TaskExecuteDetail {
313 pid: None,
314 container_id: None,
315 health: HealthStatus::Unknown,
316 started_at: Utc::now(),
317 health_check_epoch: 0,
318 }),
319 TaskRunState::Dead => TaskPhase::Terminal(TaskTerminalDetail {
320 outcome: Outcome::Failed,
321 exit_code: None,
322 finished_at: Utc::now(),
323 restarts: 0,
324 }),
325 }
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::*;
332
333 #[test]
334 fn test_phase_names() {
335 let phase: TaskPhase = TaskPhase::Initial;
336 assert_eq!(phase.phase_name(), "initial");
337
338 let phase: TaskPhase = TaskPhase::Warming(TaskWarmProgress::default());
339 assert_eq!(phase.phase_name(), "warming");
340 assert!(phase.is_active());
341
342 let phase: TaskPhase = TaskPhase::Terminal(TaskTerminalDetail {
343 outcome: Outcome::Success,
344 exit_code: Some(0),
345 finished_at: Utc::now(),
346 restarts: 0,
347 });
348 assert!(phase.is_terminal());
349 assert!(!phase.is_active());
350 }
351
352 #[test]
353 fn test_valid_transitions() {
354 assert!(is_valid_transition("initial", "warming"));
355 assert!(is_valid_transition("warming", "executing"));
356 assert!(is_valid_transition("executing", "contracting"));
357 assert!(is_valid_transition("contracting", "terminal"));
358 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")); }
366
367 #[test]
368 fn test_allocation_phase_from_legacy() {
369 let phase: AllocationPhase = AllocationState::Pending.into();
370 assert!(phase.is_initial());
371
372 let phase: AllocationPhase = AllocationState::Running.into();
373 assert_eq!(phase.phase_name(), "executing");
374
375 let phase: AllocationPhase = AllocationState::Failed.into();
376 assert!(phase.is_terminal());
377 }
378
379 #[test]
380 fn test_task_phase_from_legacy() {
381 let phase: TaskPhase = TaskRunState::Pending.into();
382 assert!(phase.is_initial());
383
384 let phase: TaskPhase = TaskRunState::Running.into();
385 assert!(phase.is_active());
386
387 let phase: TaskPhase = TaskRunState::Dead.into();
388 assert!(phase.is_terminal());
389 }
390
391 #[test]
392 fn test_serde_roundtrip() {
393 let phase = AllocationPhase::Warming(AllocWarmProgress {
394 secrets_resolved: true,
395 volumes_mounted: false,
396 network_identity_assigned: true,
397 endpoint_registered: false,
398 task_progress: HashMap::from([(
399 "web".to_string(),
400 TaskWarmProgress {
401 fetch_progress: 0.75,
402 deps_resolved: true,
403 port_allocated: true,
404 warmup_checks_passed: 2,
405 warmup_checks_required: 3,
406 },
407 )]),
408 });
409
410 let json = serde_json::to_string(&phase).unwrap();
411 let back: AllocationPhase = serde_json::from_str(&json).unwrap();
412 assert_eq!(phase, back);
413 }
414
415 #[test]
416 fn test_desired_phase_serde() {
417 let desired = DesiredPhase::Stopped {
418 reason: ContractReason::ScaleDown,
419 };
420 let json = serde_json::to_string(&desired).unwrap();
421 let back: DesiredPhase = serde_json::from_str(&json).unwrap();
422 assert_eq!(desired, back);
423 }
424}