common_meta/
cluster.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::hash::{DefaultHasher, Hash, Hasher};
16use std::str::FromStr;
17
18use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest};
19use common_error::ext::ErrorExt;
20use lazy_static::lazy_static;
21use regex::Regex;
22use serde::{Deserialize, Serialize};
23use snafu::{OptionExt, ResultExt, ensure};
24
25use crate::datanode::RegionStat;
26use crate::error::{
27    DecodeJsonSnafu, EncodeJsonSnafu, Error, FromUtf8Snafu, InvalidNodeInfoKeySnafu,
28    InvalidRoleSnafu, ParseNumSnafu, Result,
29};
30use crate::key::flow::flow_state::FlowStat;
31use crate::peer::Peer;
32
33const CLUSTER_NODE_INFO_PREFIX: &str = "__meta_cluster_node_info";
34
35lazy_static! {
36    static ref CLUSTER_NODE_INFO_PREFIX_PATTERN: Regex = Regex::new(&format!(
37        "^{CLUSTER_NODE_INFO_PREFIX}-([0-9]+)-([0-9]+)-([0-9]+)$"
38    ))
39    .unwrap();
40}
41
42/// [ClusterInfo] provides information about the cluster.
43#[async_trait::async_trait]
44pub trait ClusterInfo {
45    type Error: ErrorExt;
46
47    /// List all nodes by role in the cluster. If `role` is `None`, list all nodes.
48    async fn list_nodes(
49        &self,
50        role: Option<Role>,
51    ) -> std::result::Result<Vec<NodeInfo>, Self::Error>;
52
53    /// List all region stats in the cluster.
54    async fn list_region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error>;
55
56    /// List all flow stats in the cluster.
57    async fn list_flow_stats(&self) -> std::result::Result<Option<FlowStat>, Self::Error>;
58
59    // TODO(jeremy): Other info, like region status, etc.
60}
61
62/// The key of [NodeInfo] in the storage. The format is `__meta_cluster_node_info-0-{role}-{node_id}`.
63#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
64pub struct NodeInfoKey {
65    /// The role of the node. It can be `[Role::Datanode]` or `[Role::Frontend]`.
66    pub role: Role,
67    /// The node id.
68    pub node_id: u64,
69}
70
71impl NodeInfoKey {
72    /// Try to create a `NodeInfoKey` from a "good" heartbeat request. "good" as in every needed
73    /// piece of information is provided and valid.  
74    pub fn new(request: &HeartbeatRequest) -> Option<Self> {
75        let HeartbeatRequest { header, peer, .. } = request;
76        let header = header.as_ref()?;
77        let peer = peer.as_ref()?;
78
79        let role = header.role.try_into().ok()?;
80        let node_id = match role {
81            // Because the Frontend is stateless, it's too easy to neglect choosing a unique id
82            // for it when setting up a cluster. So we calculate its id from its address.
83            Role::Frontend => calculate_node_id(&peer.addr),
84            _ => peer.id,
85        };
86
87        Some(NodeInfoKey { role, node_id })
88    }
89
90    pub fn key_prefix() -> String {
91        format!("{}-0-", CLUSTER_NODE_INFO_PREFIX)
92    }
93
94    pub fn key_prefix_with_role(role: Role) -> String {
95        format!("{}-0-{}-", CLUSTER_NODE_INFO_PREFIX, i32::from(role))
96    }
97}
98
99/// Calculate (by using the DefaultHasher) the node's id from its address.
100fn calculate_node_id(addr: &str) -> u64 {
101    let mut hasher = DefaultHasher::new();
102    addr.hash(&mut hasher);
103    hasher.finish()
104}
105
106/// The information of a node in the cluster.
107#[derive(Debug, Serialize, Deserialize)]
108pub struct NodeInfo {
109    /// The peer information. [node_id, address]
110    pub peer: Peer,
111    /// Last activity time in milliseconds.
112    pub last_activity_ts: i64,
113    /// The status of the node. Different roles have different node status.
114    pub status: NodeStatus,
115    // The node build version
116    pub version: String,
117    // The node build git commit hash
118    pub git_commit: String,
119    // The node star timestamp
120    pub start_time_ms: u64,
121    // The node build cpus
122    #[serde(default)]
123    pub cpus: u32,
124    // The node build memory bytes
125    #[serde(default)]
126    pub memory_bytes: u64,
127    // The node build hostname
128    #[serde(default)]
129    pub hostname: String,
130}
131
132#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
133pub enum Role {
134    Datanode,
135    Frontend,
136    Flownode,
137    Metasrv,
138}
139
140#[derive(Debug, Serialize, Deserialize)]
141pub enum NodeStatus {
142    Datanode(DatanodeStatus),
143    Frontend(FrontendStatus),
144    Flownode(FlownodeStatus),
145    Metasrv(MetasrvStatus),
146    Standalone,
147}
148
149impl NodeStatus {
150    // Get the role name of the node status
151    pub fn role_name(&self) -> &str {
152        match self {
153            NodeStatus::Datanode(_) => "DATANODE",
154            NodeStatus::Frontend(_) => "FRONTEND",
155            NodeStatus::Flownode(_) => "FLOWNODE",
156            NodeStatus::Metasrv(_) => "METASRV",
157            NodeStatus::Standalone => "STANDALONE",
158        }
159    }
160}
161
162/// The status of a datanode.
163#[derive(Debug, Serialize, Deserialize)]
164pub struct DatanodeStatus {
165    /// The read capacity units during this period.
166    pub rcus: i64,
167    /// The write capacity units during this period.
168    pub wcus: i64,
169    /// How many leader regions on this node.
170    pub leader_regions: usize,
171    /// How many follower regions on this node.
172    pub follower_regions: usize,
173    /// The workloads of the datanode.
174    pub workloads: DatanodeWorkloads,
175}
176
177/// The status of a frontend.
178#[derive(Debug, Serialize, Deserialize)]
179pub struct FrontendStatus {}
180
181/// The status of a flownode.
182#[derive(Debug, Serialize, Deserialize)]
183pub struct FlownodeStatus {}
184
185/// The status of a metasrv.
186#[derive(Debug, Serialize, Deserialize)]
187pub struct MetasrvStatus {
188    pub is_leader: bool,
189}
190
191impl FromStr for NodeInfoKey {
192    type Err = Error;
193
194    fn from_str(key: &str) -> Result<Self> {
195        let caps = CLUSTER_NODE_INFO_PREFIX_PATTERN
196            .captures(key)
197            .context(InvalidNodeInfoKeySnafu { key })?;
198        ensure!(caps.len() == 4, InvalidNodeInfoKeySnafu { key });
199
200        let role = caps[2].to_string();
201        let node_id = caps[3].to_string();
202        let role: i32 = role.parse().context(ParseNumSnafu {
203            err_msg: format!("invalid role {role}"),
204        })?;
205        let role = Role::try_from(role)?;
206        let node_id: u64 = node_id.parse().context(ParseNumSnafu {
207            err_msg: format!("invalid node_id: {node_id}"),
208        })?;
209
210        Ok(Self { role, node_id })
211    }
212}
213
214impl TryFrom<Vec<u8>> for NodeInfoKey {
215    type Error = Error;
216
217    fn try_from(bytes: Vec<u8>) -> Result<Self> {
218        String::from_utf8(bytes)
219            .context(FromUtf8Snafu {
220                name: "NodeInfoKey",
221            })
222            .map(|x| x.parse())?
223    }
224}
225
226impl From<&NodeInfoKey> for Vec<u8> {
227    fn from(key: &NodeInfoKey) -> Self {
228        format!(
229            "{}-0-{}-{}",
230            CLUSTER_NODE_INFO_PREFIX,
231            i32::from(key.role),
232            key.node_id
233        )
234        .into_bytes()
235    }
236}
237
238impl FromStr for NodeInfo {
239    type Err = Error;
240
241    fn from_str(value: &str) -> Result<Self> {
242        serde_json::from_str(value).context(DecodeJsonSnafu)
243    }
244}
245
246impl TryFrom<Vec<u8>> for NodeInfo {
247    type Error = Error;
248
249    fn try_from(bytes: Vec<u8>) -> Result<Self> {
250        String::from_utf8(bytes)
251            .context(FromUtf8Snafu { name: "NodeInfo" })
252            .map(|x| x.parse())?
253    }
254}
255
256impl TryFrom<NodeInfo> for Vec<u8> {
257    type Error = Error;
258
259    fn try_from(info: NodeInfo) -> Result<Self> {
260        Ok(serde_json::to_string(&info)
261            .context(EncodeJsonSnafu)?
262            .into_bytes())
263    }
264}
265
266impl From<Role> for i32 {
267    fn from(role: Role) -> Self {
268        match role {
269            Role::Datanode => 0,
270            Role::Frontend => 1,
271            Role::Flownode => 2,
272            Role::Metasrv => 99,
273        }
274    }
275}
276
277impl TryFrom<i32> for Role {
278    type Error = Error;
279
280    fn try_from(role: i32) -> Result<Self> {
281        match role {
282            0 => Ok(Self::Datanode),
283            1 => Ok(Self::Frontend),
284            2 => Ok(Self::Flownode),
285            99 => Ok(Self::Metasrv),
286            _ => InvalidRoleSnafu { role }.fail(),
287        }
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use std::assert_matches::assert_matches;
294
295    use common_workload::DatanodeWorkloadType;
296
297    use super::*;
298    use crate::cluster::Role::{Datanode, Frontend};
299    use crate::cluster::{DatanodeStatus, NodeInfo, NodeInfoKey, NodeStatus};
300    use crate::peer::Peer;
301
302    #[test]
303    fn test_node_info_key_round_trip() {
304        let key = NodeInfoKey {
305            role: Datanode,
306            node_id: 2,
307        };
308
309        let key_bytes: Vec<u8> = (&key).into();
310        let new_key: NodeInfoKey = key_bytes.try_into().unwrap();
311
312        assert_eq!(Datanode, new_key.role);
313        assert_eq!(2, new_key.node_id);
314    }
315
316    #[test]
317    fn test_node_info_round_trip() {
318        let node_info = NodeInfo {
319            peer: Peer {
320                id: 1,
321                addr: "127.0.0.1".to_string(),
322            },
323            last_activity_ts: 123,
324            status: NodeStatus::Datanode(DatanodeStatus {
325                rcus: 1,
326                wcus: 2,
327                leader_regions: 3,
328                follower_regions: 4,
329                workloads: DatanodeWorkloads {
330                    types: vec![DatanodeWorkloadType::Hybrid.to_i32()],
331                },
332            }),
333            version: "".to_string(),
334            git_commit: "".to_string(),
335            start_time_ms: 1,
336            cpus: 0,
337            memory_bytes: 0,
338            hostname: "test_hostname".to_string(),
339        };
340
341        let node_info_bytes: Vec<u8> = node_info.try_into().unwrap();
342        let new_node_info: NodeInfo = node_info_bytes.try_into().unwrap();
343
344        assert_matches!(
345            new_node_info,
346            NodeInfo {
347                peer: Peer { id: 1, .. },
348                last_activity_ts: 123,
349                status: NodeStatus::Datanode(DatanodeStatus {
350                    rcus: 1,
351                    wcus: 2,
352                    leader_regions: 3,
353                    follower_regions: 4,
354                    ..
355                }),
356                start_time_ms: 1,
357                ..
358            }
359        );
360    }
361
362    #[test]
363    fn test_node_info_key_prefix() {
364        let prefix = NodeInfoKey::key_prefix();
365        assert_eq!(prefix, "__meta_cluster_node_info-0-");
366
367        let prefix = NodeInfoKey::key_prefix_with_role(Frontend);
368        assert_eq!(prefix, "__meta_cluster_node_info-0-1-");
369    }
370
371    #[test]
372    fn test_calculate_node_id_from_addr() {
373        // Test empty string
374        assert_eq!(calculate_node_id(""), calculate_node_id(""));
375
376        // Test same addresses return same ids
377        let addr1 = "127.0.0.1:8080";
378        let id1 = calculate_node_id(addr1);
379        let id2 = calculate_node_id(addr1);
380        assert_eq!(id1, id2);
381
382        // Test different addresses return different ids
383        let addr2 = "127.0.0.1:8081";
384        let id3 = calculate_node_id(addr2);
385        assert_ne!(id1, id3);
386
387        // Test long address
388        let long_addr = "very.long.domain.name.example.com:9999";
389        let id4 = calculate_node_id(long_addr);
390        assert!(id4 > 0);
391    }
392}