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, 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    /// 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
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
213/// Concrete allocation lifecycle phase.
214pub type AllocationPhase =
215    WorkloadPhase<AllocWarmProgress, AllocExecuteDetail, AllocContractDetail, AllocTerminalDetail>;
216
217// ── Node-Level Detail Types ────────────────────────────────────
218
219#[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    /// WireGuard mesh tunnel established.
225    #[serde(default)]
226    pub wireguard_tunnel_up: bool,
227    /// Number of mesh peers connected.
228    #[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
251/// Concrete node lifecycle phase.
252pub type NodePhase =
253    WorkloadPhase<NodeWarmProgress, NodeExecuteDetail, NodeContractDetail, NodeTerminalDetail>;
254
255// ── Desired/Observed State (for Raft replication) ──────────────
256
257/// What the scheduler declares an allocation should be.
258#[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/// What a node observes an allocation to actually be.
270#[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
279// ── Migration from legacy enums ────────────────────────────────
280
281use 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")); // restart
359        assert!(is_valid_transition("warming", "terminal")); // fast-fail
360
361        assert!(!is_valid_transition("initial", "executing")); // skip warm
362        assert!(!is_valid_transition("executing", "warming")); // backward
363        assert!(!is_valid_transition("terminal", "initial")); // dead is dead
364        assert!(!is_valid_transition("initial", "contracting")); // nothing to contract
365    }
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}