meta_srv/key/
node_scheduling.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::str::FromStr;
16
17use snafu::OptionExt;
18
19use crate::error;
20use crate::error::Result;
21
22/// Key prefix of datanode scheduling states.
23const DATANODE_SCHEDULING_PREFIX: &str = "__node_scheduling/datanode/";
24/// Key prefix of flownode scheduling states.
25const FLOWNODE_SCHEDULING_PREFIX: &str = "__node_scheduling/flownode/";
26
27/// The node roles whose scheduling state is tracked.
28///
29/// A subset of [`common_meta::cluster::Role`]: frontends and metasrv nodes never receive
30/// regions, so they are out of scope.
31#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)]
32pub enum NodeSchedulingRole {
33    Datanode,
34    Flownode,
35}
36
37impl NodeSchedulingRole {
38    /// Returns the key prefix of this role, e.g. `__node_scheduling/datanode/`.
39    fn key_prefix(&self) -> &'static str {
40        match self {
41            NodeSchedulingRole::Datanode => DATANODE_SCHEDULING_PREFIX,
42            NodeSchedulingRole::Flownode => FLOWNODE_SCHEDULING_PREFIX,
43        }
44    }
45}
46
47/// Key of a node's scheduling state: `__node_scheduling/{role}/{node_id}`,
48/// for example `__node_scheduling/datanode/7`.
49#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)]
50pub struct NodeSchedulingKey {
51    pub role: NodeSchedulingRole,
52    pub node_id: u64,
53}
54
55impl NodeSchedulingKey {
56    pub fn new(role: NodeSchedulingRole, node_id: u64) -> NodeSchedulingKey {
57        NodeSchedulingKey { role, node_id }
58    }
59
60    /// Returns the prefix covering every node of `role`.
61    pub fn prefix_key(role: NodeSchedulingRole) -> Vec<u8> {
62        role.key_prefix().as_bytes().to_vec()
63    }
64}
65
66impl From<NodeSchedulingKey> for Vec<u8> {
67    fn from(key: NodeSchedulingKey) -> Vec<u8> {
68        format!("{}{}", key.role.key_prefix(), key.node_id).into_bytes()
69    }
70}
71
72impl FromStr for NodeSchedulingKey {
73    type Err = error::Error;
74
75    fn from_str(key: &str) -> Result<NodeSchedulingKey> {
76        for role in [NodeSchedulingRole::Datanode, NodeSchedulingRole::Flownode] {
77            let Some(node_id) = key.strip_prefix(role.key_prefix()) else {
78                continue;
79            };
80
81            let node_id = node_id
82                .parse::<u64>()
83                .ok()
84                .context(error::InvalidNodeSchedulingKeySnafu { key })?;
85            return Ok(NodeSchedulingKey { role, node_id });
86        }
87
88        error::InvalidNodeSchedulingKeySnafu { key }.fail()
89    }
90}
91
92impl TryFrom<Vec<u8>> for NodeSchedulingKey {
93    type Error = error::Error;
94
95    fn try_from(bytes: Vec<u8>) -> Result<NodeSchedulingKey> {
96        let key = String::from_utf8(bytes).map_err(|err| {
97            error::InvalidNodeSchedulingKeySnafu {
98                key: String::from_utf8_lossy(err.as_bytes()).into_owned(),
99            }
100            .build()
101        })?;
102        key.parse()
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn test_key_round_trip() {
112        for (role, expected) in [
113            (NodeSchedulingRole::Datanode, "__node_scheduling/datanode/7"),
114            (NodeSchedulingRole::Flownode, "__node_scheduling/flownode/7"),
115        ] {
116            let key = NodeSchedulingKey::new(role, 7);
117            let bytes: Vec<u8> = key.into();
118            assert_eq!(expected.as_bytes(), bytes.as_slice());
119
120            let parsed = NodeSchedulingKey::try_from(bytes).unwrap();
121            assert_eq!(key, parsed);
122            assert_eq!(role, parsed.role);
123        }
124    }
125
126    #[test]
127    fn test_key_prefix() {
128        assert_eq!(
129            b"__node_scheduling/datanode/",
130            NodeSchedulingKey::prefix_key(NodeSchedulingRole::Datanode).as_slice()
131        );
132        assert_eq!(
133            b"__node_scheduling/flownode/",
134            NodeSchedulingKey::prefix_key(NodeSchedulingRole::Flownode).as_slice()
135        );
136    }
137
138    #[test]
139    fn test_invalid_key() {
140        for key in [
141            "__node_scheduling/frontend/7",
142            "__node_scheduling/datanode/",
143            "__node_scheduling/datanode/x",
144            "__node_scheduling",
145        ] {
146            let err = key.parse::<NodeSchedulingKey>().unwrap_err();
147            assert!(
148                matches!(err, error::Error::InvalidNodeSchedulingKey { .. }),
149                "unexpected error for {key}: {err:?}"
150            );
151        }
152    }
153}
meta_srv/key/node_scheduling.rs

meta_srv/key/
node_scheduling.rs
