Skip to main content

tatara_core/domain/
lifecycle.rs

1//! Universal workload lifecycle — the distributed state machine.
2//!
3//! Every workload (task, allocation, job, node) follows:
4//!   Initial → Warming → Executing → Contracting → Terminal
5//!
6//! This applies at every level with type-specific detail structs.
7//! The generic `WorkloadPhase<W, E, C, T>` is parameterized by
8//! detail types for each phase.
9
10use chrono::{DateTime, Utc};
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13use uuid::Uuid;
14
15// ── Generic Phase Enum ─────────────────────────────────────────
16
17/// The universal workload lifecycle phase.
18///
19/// State transitions form a strict DAG:
20/// ```text
21///   Initial → Warming → Executing → Contracting → Terminal
22///                     ↘ Contracting (warm failed)
23///                     ↘ Terminal (fast-fail, no cleanup needed)
24///                                  → Initial (successful contraction, restart)
25/// ```
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27#[serde(tag = "phase", rename_all = "snake_case")]
28pub enum WorkloadPhase<W, E, C, T> {
29    /// Defined but not active. Zero resources allocated.
30    Initial,
31    /// Preparing to execute. Resources being acquired.
32    Warming(W),
33    /// Active and serving.
34    Executing(E),
35    /// Gracefully winding down.
36    Contracting(C),
37    /// Final state. Will not transition again (unless recycled to Initial).
38    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
65/// Validate that a phase transition is legal.
66pub 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// ── Shared Types ───────────────────────────────────────────────
80
81/// Why a workload is contracting.
82#[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/// Final outcome of a workload.
94#[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/// Health status (used across levels).
104#[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/// Desired phase for an allocation (what the scheduler wants).
119#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
120#[serde(rename_all = "snake_case")]
121pub enum DesiredPhase {
122    Active,
123    Stopped { reason: ContractReason },
124}
125
126// ── Task-Level Detail Types ────────────────────────────────────
127
128#[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
174/// Concrete task lifecycle phase.
175pub type TaskPhase =
176    WorkloadPhase<TaskWarmProgress, TaskExecuteDetail, TaskContractDetail, TaskTerminalDetail>;
177
178// ── Allocation-Level Detail Types ──────────────────────────────
179
180#[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    /// Network identity assigned in the routing table.
186    #[serde(default)]
187    pub network_identity_assigned: bool,
188    /// Endpoint registered in the networking plane.
189    #[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
225/// Concrete allocation lifecycle phase.
226pub type AllocationPhase =
227    WorkloadPhase<AllocWarmProgress, AllocExecuteDetail, AllocContractDetail, AllocTerminalDetail>;
228
229// ── Node-Level Detail Types ────────────────────────────────────
230
231#[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    /// WireGuard mesh tunnel established.
237    #[serde(default)]
238    pub wireguard_tunnel_up: bool,
239    /// Number of mesh peers connected.
240    #[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
275/// Concrete node lifecycle phase.
276pub type NodePhase =
277    WorkloadPhase<NodeWarmProgress, NodeExecuteDetail, NodeContractDetail, NodeTerminalDetail>;
278
279// ── Desired/Observed State (for Raft replication) ──────────────
280
281/// What the scheduler declares an allocation should be.
282#[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/// What a node observes an allocation to actually be.
294#[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
303// ── Migration from legacy enums ────────────────────────────────
304
305use 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")); // restart
383        assert!(is_valid_transition("warming", "terminal")); // fast-fail
384
385        assert!(!is_valid_transition("initial", "executing")); // skip warm
386        assert!(!is_valid_transition("executing", "warming")); // backward
387        assert!(!is_valid_transition("terminal", "initial")); // dead is dead
388        assert!(!is_valid_transition("initial", "contracting")); // nothing to contract
389    }
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}