Skip to main content

tatara_core/domain/
allocation.rs

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