tatara_engine/domain/
store_adapter.rs

1//! Adapter between ClusterStore (Raft-backed) and the Evaluator/Scheduler.
2//!
3//! The Evaluator was originally written against StateStore (local JSON file).
4//! This adapter wraps ClusterStore to provide the same read interface,
5//! ensuring the scheduler reads from Raft-replicated state instead of
6//! divergent local state.
7
8use anyhow::Result;
9use std::collections::HashMap;
10use std::sync::Arc;
11use uuid::Uuid;
12
13use crate::cluster::store::ClusterStore;
14use tatara_core::cluster::types::NodeMeta;
15use tatara_core::domain::allocation::Allocation;
16use tatara_core::domain::job::{Job, JobStatus};
17use tatara_core::domain::node::{Node, NodeStatus};
18
19/// Wraps ClusterStore to provide StateStore-compatible reads for the Evaluator.
20pub struct ClusterStoreAdapter {
21    store: Arc<ClusterStore>,
22}
23
24impl ClusterStoreAdapter {
25    pub fn new(store: Arc<ClusterStore>) -> Self {
26        Self { store }
27    }
28
29    /// List all jobs (from Raft state, eventually consistent).
30    pub async fn list_jobs(&self) -> Vec<Job> {
31        self.store.list_jobs().await
32    }
33
34    /// List all nodes as `Node` type (converted from `NodeMeta`).
35    pub async fn list_nodes(&self) -> Vec<Node> {
36        self.store
37            .list_nodes()
38            .await
39            .into_iter()
40            .map(node_meta_to_node)
41            .collect()
42    }
43
44    /// Get the current scheduling generation.
45    pub async fn scheduling_generation(&self) -> u64 {
46        let state = self.store.state().await;
47        state.scheduling_generation
48    }
49
50    /// Submit a job through Raft.
51    pub async fn put_job(&self, job: Job) -> Result<()> {
52        self.store.put_job(job).await?;
53        Ok(())
54    }
55
56    /// Update job status through Raft.
57    pub async fn update_job_status(&self, job_id: &str, status: JobStatus) -> Result<()> {
58        self.store.update_job_status(job_id, status).await?;
59        Ok(())
60    }
61
62    /// Submit an allocation through Raft.
63    pub async fn put_allocation(&self, alloc: Allocation) -> Result<()> {
64        self.store.put_allocation(alloc).await?;
65        Ok(())
66    }
67
68    /// Check if this node is the Raft leader.
69    pub async fn is_leader(&self) -> bool {
70        self.store.is_leader().await
71    }
72}
73
74/// Convert NodeMeta (cluster type) to Node (domain type) for the Evaluator.
75fn node_meta_to_node(meta: NodeMeta) -> Node {
76    let mut attributes = HashMap::new();
77    attributes.insert("os".to_string(), meta.os.clone());
78    attributes.insert("arch".to_string(), meta.arch.clone());
79    attributes.insert("hostname".to_string(), meta.hostname.clone());
80
81    Node {
82        id: format!("{}", meta.node_id),
83        address: meta.http_addr.clone(),
84        status: if meta.eligible {
85            NodeStatus::Ready
86        } else {
87            NodeStatus::Draining
88        },
89        eligible: meta.eligible,
90        total_resources: meta.total_resources,
91        available_resources: meta.available_resources,
92        attributes,
93        drivers: meta.drivers,
94        last_heartbeat: meta.joined_at,
95        allocations: Vec::new(),
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use chrono::Utc;
103    use tatara_core::cluster::types::NodeRoles;
104    use tatara_core::domain::job::Resources;
105
106    #[test]
107    fn test_node_meta_conversion() {
108        let meta = NodeMeta {
109            node_id: 42,
110            hostname: "test-host".to_string(),
111            http_addr: "127.0.0.1:4646".to_string(),
112            gossip_addr: "127.0.0.1:5679".to_string(),
113            raft_addr: "127.0.0.1:4649".to_string(),
114            os: "darwin".to_string(),
115            arch: "aarch64".to_string(),
116            roles: NodeRoles::default(),
117            drivers: vec![],
118            total_resources: Resources {
119                cpu_mhz: 4000,
120                memory_mb: 8192,
121            },
122            available_resources: Resources {
123                cpu_mhz: 3000,
124                memory_mb: 6144,
125            },
126            allocations_running: 2,
127            joined_at: Utc::now(),
128            version: "0.2.0".to_string(),
129            eligible: true,
130            wireguard_pubkey: None,
131            tunnel_address: None,
132        };
133
134        let node = node_meta_to_node(meta);
135        assert_eq!(node.id, "42");
136        assert_eq!(node.status, NodeStatus::Ready);
137        assert!(node.eligible);
138        assert_eq!(node.attributes["arch"], "aarch64");
139        assert_eq!(node.total_resources.cpu_mhz, 4000);
140    }
141}
tatara_engine/domain/store_adapter.rs

tatara_engine/domain/
store_adapter.rs
