meta_client/client/
util.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 api::v1::meta::heartbeat_request::NodeWorkloads;
16use api::v1::meta::{ErrorCode, ResponseHeader};
17use common_meta::cluster::{NodeInfo, NodeStatus};
18use common_time::util::SystemTimer;
19use tonic::{Code, Status};
20
21pub(crate) fn is_unreachable(status: &Status) -> bool {
22    status.code() == Code::Unavailable || status.code() == Code::DeadlineExceeded
23}
24
25pub(crate) fn is_not_leader(header: &Option<ResponseHeader>) -> bool {
26    let Some(header) = header else {
27        return false;
28    };
29
30    let Some(err) = header.error.as_ref() else {
31        return false;
32    };
33
34    err.code == ErrorCode::NotLeader as i32
35}
36
37fn is_active_node(
38    timer: &impl SystemTimer,
39    last_activity_ts: i64,
40    active_duration: std::time::Duration,
41) -> bool {
42    let now = timer.current_time_millis();
43    let elapsed = now.checked_sub(last_activity_ts).unwrap_or(0) as u64;
44    elapsed < active_duration.as_millis() as u64
45}
46
47pub(crate) fn alive_frontends(
48    timer: &impl SystemTimer,
49    nodes: Vec<NodeInfo>,
50    active_duration: std::time::Duration,
51) -> Vec<NodeInfo> {
52    nodes
53        .into_iter()
54        .filter_map(|node| {
55            if matches!(node.status, NodeStatus::Frontend(_))
56                && is_active_node(timer, node.last_activity_ts, active_duration)
57            {
58                Some(node)
59            } else {
60                None
61            }
62        })
63        .collect()
64}
65
66pub(crate) fn alive_datanodes(
67    timer: &impl SystemTimer,
68    nodes: Vec<NodeInfo>,
69    active_duration: std::time::Duration,
70    filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
71) -> Vec<NodeInfo> {
72    let filter = filter.unwrap_or(|_| true);
73
74    nodes
75        .into_iter()
76        .filter_map(|node| {
77            if let NodeStatus::Datanode(status) = &node.status
78                && is_active_node(timer, node.last_activity_ts, active_duration)
79            {
80                let workloads = NodeWorkloads::Datanode(status.workloads.clone());
81                filter(&workloads).then_some(node)
82            } else {
83                None
84            }
85        })
86        .collect()
87}
88
89pub(crate) fn alive_flownodes(
90    timer: &impl SystemTimer,
91    nodes: Vec<NodeInfo>,
92    active_duration: std::time::Duration,
93    filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
94) -> Vec<NodeInfo> {
95    let filter = filter.unwrap_or(|_| true);
96
97    nodes
98        .into_iter()
99        .filter_map(|node| {
100            if let NodeStatus::Flownode(status) = &node.status
101                && is_active_node(timer, node.last_activity_ts, active_duration)
102            {
103                let workloads = NodeWorkloads::Flownode(status.workloads.clone());
104                filter(&workloads).then_some(node)
105            } else {
106                None
107            }
108        })
109        .collect()
110}
111
112#[cfg(test)]
113mod tests {
114    use std::time::Duration;
115
116    use api::v1::meta::heartbeat_request::NodeWorkloads;
117    use api::v1::meta::{DatanodeWorkloads, FlownodeWorkloads, Peer};
118    use common_meta::cluster::{
119        DatanodeStatus, FlownodeStatus, FrontendStatus, MetasrvStatus, NodeInfo, NodeStatus, Role,
120    };
121    use common_time::util::SystemTimer;
122
123    use super::*;
124
125    struct MockSystemTimer(i64);
126
127    impl MockSystemTimer {
128        fn new(now: i64) -> Self {
129            Self(now)
130        }
131    }
132
133    impl SystemTimer for MockSystemTimer {
134        fn current_time_millis(&self) -> i64 {
135            self.0
136        }
137
138        fn current_time_rfc3339(&self) -> String {
139            "1970-01-01T00:00:00Z".to_string()
140        }
141    }
142
143    fn node_info(role: Role, id: u64, addr: &str, last_activity_ts: i64) -> NodeInfo {
144        let status = match role {
145            Role::Frontend => NodeStatus::Frontend(FrontendStatus::default()),
146            Role::Datanode => NodeStatus::Datanode(DatanodeStatus {
147                rcus: 0,
148                wcus: 0,
149                leader_regions: 0,
150                follower_regions: 0,
151                workloads: DatanodeWorkloads { types: vec![] },
152            }),
153            Role::Flownode => NodeStatus::Flownode(FlownodeStatus {
154                workloads: FlownodeWorkloads { types: vec![] },
155            }),
156            Role::Metasrv => NodeStatus::Metasrv(MetasrvStatus { is_leader: false }),
157        };
158
159        NodeInfo {
160            peer: Peer::new(id, addr),
161            last_activity_ts,
162            status,
163            version: String::new(),
164            git_commit: String::new(),
165            start_time_ms: 0,
166            total_cpu_millicores: 0,
167            total_memory_bytes: 0,
168            cpu_usage_millicores: 0,
169            memory_usage_bytes: 0,
170            hostname: String::new(),
171            env_vars: Default::default(),
172            schedulable: None,
173        }
174    }
175
176    fn ingest_only(workloads: &NodeWorkloads) -> bool {
177        matches!(
178            workloads,
179            NodeWorkloads::Datanode(DatanodeWorkloads { types }) if types.as_slice() == [1]
180        )
181    }
182
183    fn empty_flownode_workloads(workloads: &NodeWorkloads) -> bool {
184        matches!(
185            workloads,
186            NodeWorkloads::Flownode(FlownodeWorkloads { types }) if types.is_empty()
187        )
188    }
189
190    #[test]
191    fn test_alive_frontends_filters_by_activity_and_role() {
192        let timer = MockSystemTimer::new(100);
193        let peers = alive_frontends(
194            &timer,
195            vec![
196                node_info(Role::Frontend, 1, "127.0.0.1:3001", 95),
197                node_info(Role::Frontend, 2, "127.0.0.1:3002", 89),
198                node_info(Role::Datanode, 3, "127.0.0.1:4001", 99),
199            ],
200            Duration::from_millis(10),
201        );
202
203        assert_eq!(
204            vec![1],
205            peers
206                .into_iter()
207                .map(|node| node.peer.id)
208                .collect::<Vec<_>>()
209        );
210    }
211
212    #[test]
213    fn test_alive_datanodes_filters_by_activity_and_workload() {
214        let timer = MockSystemTimer::new(100);
215        let mut first = node_info(Role::Datanode, 1, "127.0.0.1:4001", 95);
216        let mut second = node_info(Role::Datanode, 2, "127.0.0.1:4002", 95);
217        let stale = node_info(Role::Datanode, 3, "127.0.0.1:4003", 89);
218
219        if let NodeStatus::Datanode(status) = &mut first.status {
220            status.workloads = DatanodeWorkloads { types: vec![1] };
221        }
222        if let NodeStatus::Datanode(status) = &mut second.status {
223            status.workloads = DatanodeWorkloads { types: vec![2] };
224        }
225
226        let peers = alive_datanodes(
227            &timer,
228            vec![first, second, stale],
229            Duration::from_millis(10),
230            Some(ingest_only),
231        );
232
233        assert_eq!(
234            vec![1],
235            peers
236                .into_iter()
237                .map(|node| node.peer.id)
238                .collect::<Vec<_>>()
239        );
240    }
241
242    #[test]
243    fn test_alive_flownodes_uses_empty_workload_semantics() {
244        let timer = MockSystemTimer::new(100);
245        let peers = alive_flownodes(
246            &timer,
247            vec![
248                node_info(Role::Flownode, 1, "127.0.0.1:5001", 95),
249                node_info(Role::Flownode, 2, "127.0.0.1:5002", 89),
250                node_info(Role::Frontend, 3, "127.0.0.1:3001", 99),
251            ],
252            Duration::from_millis(10),
253            Some(empty_flownode_workloads),
254        );
255
256        assert_eq!(
257            vec![1],
258            peers
259                .into_iter()
260                .map(|node| node.peer.id)
261                .collect::<Vec<_>>()
262        );
263    }
264
265    #[test]
266    fn test_alive_flownodes_filters_by_workloads() {
267        let timer = MockSystemTimer::new(100);
268        let mut first = node_info(Role::Flownode, 1, "127.0.0.1:5001", 95);
269        let mut second = node_info(Role::Flownode, 2, "127.0.0.1:5002", 95);
270
271        if let NodeStatus::Flownode(status) = &mut first.status {
272            status.workloads = FlownodeWorkloads { types: vec![7] };
273        }
274        if let NodeStatus::Flownode(status) = &mut second.status {
275            status.workloads = FlownodeWorkloads { types: vec![8] };
276        }
277
278        fn workload_type_is_7(workloads: &NodeWorkloads) -> bool {
279            matches!(
280                workloads,
281                NodeWorkloads::Flownode(FlownodeWorkloads { types }) if types.as_slice() == [7]
282            )
283        }
284
285        let peers = alive_flownodes(
286            &timer,
287            vec![first, second],
288            Duration::from_millis(10),
289            Some(workload_type_is_7),
290        );
291
292        assert_eq!(
293            vec![1],
294            peers
295                .into_iter()
296                .map(|node| node.peer.id)
297                .collect::<Vec<_>>()
298        );
299    }
300}
meta_client/client/util.rs

meta_client/client/
util.rs
