datafusion_dist/
cluster.rs

1use std::{
2    collections::HashMap,
3    fmt::{Debug, Display},
4};
5
6use serde::{Deserialize, Serialize};
7
8use crate::{DistError, DistResult};
9
10#[async_trait::async_trait]
11pub trait DistCluster: Debug + Send + Sync {
12    // Send heartbeat
13    async fn heartbeat(&self, node_id: NodeId, state: NodeState) -> DistResult<()>;
14    // Get alive nodes
15    async fn alive_nodes(&self) -> DistResult<HashMap<NodeId, NodeState>>;
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
19pub struct NodeId {
20    pub host: String,
21    pub port: u16,
22}
23
24impl Display for NodeId {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        write!(f, "{}:{}", self.host, self.port)
27    }
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
31pub enum NodeStatus {
32    #[default]
33    Available,
34    Terminating,
35}
36
37impl Display for NodeStatus {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        match self {
40            NodeStatus::Available => write!(f, "Available"),
41            NodeStatus::Terminating => write!(f, "Terminating"),
42        }
43    }
44}
45
46impl std::str::FromStr for NodeStatus {
47    type Err = DistError;
48
49    fn from_str(s: &str) -> Result<Self, Self::Err> {
50        match s {
51            "Available" => Ok(NodeStatus::Available),
52            "Terminating" => Ok(NodeStatus::Terminating),
53            _ => Err(DistError::internal(format!("Unknown NodeStatus: {s}"))),
54        }
55    }
56}
57
58#[derive(Debug, Clone, Default, Serialize, Deserialize)]
59pub struct NodeState {
60    pub status: NodeStatus,
61    pub total_memory: u64,
62    pub used_memory: u64,
63    pub free_memory: u64,
64    pub available_memory: u64,
65    pub global_cpu_usage: f32,
66    pub num_running_tasks: u32,
67}
datafusion_dist/cluster.rs

datafusion_dist/
cluster.rs
