tatara_core/domain/
allocation.rs1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use uuid::Uuid;
5
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
7#[serde(rename_all = "snake_case")]
8pub enum AllocationState {
9 Pending,
10 Running,
11 Complete,
12 Failed,
13 Lost,
14}
15
16#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
17#[serde(rename_all = "snake_case")]
18pub enum TaskRunState {
19 #[default]
20 Pending,
21 Running,
22 Dead,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct Allocation {
27 pub id: Uuid,
28 pub job_id: String,
29 pub group_name: String,
30 pub node_id: String,
31 pub state: AllocationState,
32 pub created_at: DateTime<Utc>,
33 pub task_states: HashMap<String, TaskState>,
34 #[serde(default)]
36 pub job_version: u64,
37}
38
39#[derive(Debug, Clone, Default, Serialize, Deserialize)]
40pub struct TaskState {
41 pub state: TaskRunState,
42 pub pid: Option<u32>,
43 pub exit_code: Option<i32>,
44 pub started_at: Option<DateTime<Utc>>,
45 pub finished_at: Option<DateTime<Utc>>,
46 pub restarts: u32,
47}
48
49impl Allocation {
50 pub fn new(
51 job_id: String,
52 group_name: String,
53 node_id: String,
54 task_names: Vec<String>,
55 ) -> Self {
56 let task_states = task_names
57 .into_iter()
58 .map(|name| (name, TaskState::new()))
59 .collect();
60
61 Self {
62 id: Uuid::new_v4(),
63 job_id,
64 group_name,
65 node_id,
66 state: AllocationState::Pending,
67 created_at: Utc::now(),
68 task_states,
69 job_version: 0,
70 }
71 }
72
73 pub fn with_job_version(mut self, version: u64) -> Self {
74 self.job_version = version;
75 self
76 }
77
78 pub fn is_terminal(&self) -> bool {
79 matches!(
80 self.state,
81 AllocationState::Complete | AllocationState::Failed | AllocationState::Lost
82 )
83 }
84}
85
86impl TaskState {
87 pub fn new() -> Self {
88 Self::default()
89 }
90}