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, 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    /// The job version this allocation was created from.
35    #[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}