meta_srv/discovery/
lease.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 common_meta::kv_backend::KvBackend;
16use common_meta::rpc::KeyValue;
17use common_meta::rpc::store::RangeRequest;
18
19use crate::cluster::MetaPeerClient;
20use crate::error::Result;
21use crate::key::{DatanodeLeaseKey, FlownodeLeaseKey, LeaseValue};
22
23#[derive(Clone, Copy)]
24pub enum LeaseValueType {
25    Flownode,
26    Datanode,
27}
28
29#[async_trait::async_trait]
30pub trait LeaseValueAccessor: Send + Sync {
31    /// Returns the peer id and lease value.
32    async fn lease_values(
33        &self,
34        lease_value_type: LeaseValueType,
35    ) -> Result<Vec<(u64, LeaseValue)>>;
36
37    async fn lease_value(
38        &self,
39        lease_value_type: LeaseValueType,
40        node_id: u64,
41    ) -> Result<Option<(u64, LeaseValue)>>;
42}
43
44fn decoder(lease_value_type: LeaseValueType, kv: KeyValue) -> Result<(u64, LeaseValue)> {
45    match lease_value_type {
46        LeaseValueType::Flownode => {
47            let lease_key: FlownodeLeaseKey = kv.key.try_into()?;
48            let lease_value: LeaseValue = kv.value.try_into()?;
49            Ok((lease_key.node_id, lease_value))
50        }
51        LeaseValueType::Datanode => {
52            let lease_key: DatanodeLeaseKey = kv.key.try_into()?;
53            let lease_value: LeaseValue = kv.value.try_into()?;
54            Ok((lease_key.node_id, lease_value))
55        }
56    }
57}
58
59#[async_trait::async_trait]
60impl LeaseValueAccessor for MetaPeerClient {
61    async fn lease_values(
62        &self,
63        lease_value_type: LeaseValueType,
64    ) -> Result<Vec<(u64, LeaseValue)>> {
65        let prefix = match lease_value_type {
66            LeaseValueType::Flownode => FlownodeLeaseKey::prefix_key_by_cluster(),
67            LeaseValueType::Datanode => DatanodeLeaseKey::prefix_key(),
68        };
69        let range_request = RangeRequest::new().with_prefix(prefix);
70        let response = self.range(range_request).await?;
71        response
72            .kvs
73            .into_iter()
74            .map(|kv| {
75                let (lease_key, lease_value) = decoder(lease_value_type, kv)?;
76                Ok((lease_key, lease_value))
77            })
78            .collect::<Result<Vec<_>>>()
79    }
80
81    async fn lease_value(
82        &self,
83        lease_value_type: LeaseValueType,
84        node_id: u64,
85    ) -> Result<Option<(u64, LeaseValue)>> {
86        let key: Vec<u8> = match lease_value_type {
87            LeaseValueType::Flownode => FlownodeLeaseKey { node_id }.try_into()?,
88            LeaseValueType::Datanode => DatanodeLeaseKey { node_id }.try_into()?,
89        };
90
91        let response = self.get(&key).await?;
92        response.map(|kv| decoder(lease_value_type, kv)).transpose()
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use std::collections::HashMap;
99
100    use api::v1::meta::{DatanodeWorkloads, FlownodeWorkloads, FrontendWorkloads};
101    use common_meta::cluster::{
102        DatanodeStatus, FlownodeStatus, FrontendStatus, NodeInfo, NodeInfoKey, NodeStatus, Role,
103    };
104    use common_meta::distributed_time_constants::default_distributed_time_constants;
105    use common_meta::kv_backend::ResettableKvBackendRef;
106    use common_meta::peer::Peer;
107    use common_meta::rpc::store::PutRequest;
108    use common_time::util::current_time_millis;
109    use common_workload::DatanodeWorkloadType;
110
111    use crate::discovery::utils::accept_ingest_workload;
112    use crate::test_util::create_meta_peer_client;
113
114    async fn put_node_info(kv_backend: &ResettableKvBackendRef, key: NodeInfoKey, value: NodeInfo) {
115        kv_backend
116            .put(PutRequest {
117                key: (&key).into(),
118                value: value.try_into().unwrap(),
119                prev_kv: false,
120            })
121            .await
122            .unwrap();
123    }
124
125    #[tokio::test]
126    async fn test_active_datanodes_returns_node_info_with_env_vars() {
127        let client = create_meta_peer_client();
128        let in_memory = client.memory_backend();
129
130        let mut env_vars = HashMap::new();
131        env_vars.insert("AZ".to_string(), "az-a".to_string());
132
133        put_node_info(
134            &in_memory,
135            NodeInfoKey {
136                role: Role::Datanode,
137                node_id: 1,
138            },
139            NodeInfo {
140                peer: Peer::new(1, "127.0.0.1:4001".to_string()),
141                last_activity_ts: current_time_millis(),
142                status: NodeStatus::Datanode(DatanodeStatus {
143                    rcus: 0,
144                    wcus: 0,
145                    leader_regions: 0,
146                    follower_regions: 0,
147                    workloads: DatanodeWorkloads {
148                        types: vec![DatanodeWorkloadType::Hybrid as i32],
149                    },
150                }),
151                version: String::new(),
152                git_commit: String::new(),
153                start_time_ms: 0,
154                total_cpu_millicores: 0,
155                total_memory_bytes: 0,
156                cpu_usage_millicores: 0,
157                memory_usage_bytes: 0,
158                hostname: String::new(),
159                env_vars,
160                schedulable: None,
161            },
162        )
163        .await;
164        let nodes = client
165            .active_datanodes(Some(accept_ingest_workload))
166            .await
167            .unwrap();
168
169        assert_eq!(nodes.len(), 1);
170        assert_eq!(nodes[0].peer.id, 1);
171        assert_eq!(
172            nodes[0].env_vars.get("AZ").map(String::as_str),
173            Some("az-a")
174        );
175    }
176
177    #[tokio::test]
178    async fn test_active_flownodes_returns_node_info() {
179        let client = create_meta_peer_client();
180        let in_memory = client.memory_backend();
181
182        put_node_info(
183            &in_memory,
184            NodeInfoKey {
185                role: Role::Flownode,
186                node_id: 11,
187            },
188            NodeInfo {
189                peer: Peer::new(11, "127.0.0.1:5001".to_string()),
190                last_activity_ts: current_time_millis(),
191                status: NodeStatus::Flownode(FlownodeStatus {
192                    workloads: FlownodeWorkloads { types: vec![7] },
193                }),
194                version: String::new(),
195                git_commit: String::new(),
196                start_time_ms: 0,
197                total_cpu_millicores: 0,
198                total_memory_bytes: 0,
199                cpu_usage_millicores: 0,
200                memory_usage_bytes: 0,
201                hostname: String::new(),
202                env_vars: Default::default(),
203                schedulable: None,
204            },
205        )
206        .await;
207
208        let nodes = client.active_flownodes(None).await.unwrap();
209
210        assert_eq!(nodes.len(), 1);
211        assert_eq!(nodes[0].peer.id, 11);
212    }
213
214    #[tokio::test]
215    async fn test_no_active_frontends() {
216        let client = create_meta_peer_client();
217        let in_memory = client.memory_backend();
218
219        let frontend_heartbeat_interval =
220            default_distributed_time_constants().frontend_heartbeat_interval;
221        let last_activity_ts =
222            current_time_millis() - frontend_heartbeat_interval.as_millis() as i64 - 1000;
223        let active_frontend_node = NodeInfo {
224            peer: Peer {
225                id: 0,
226                addr: "127.0.0.1:20201".to_string(),
227            },
228            last_activity_ts,
229            status: NodeStatus::Frontend(FrontendStatus {
230                workloads: FrontendWorkloads { types: vec![] },
231            }),
232            version: "1.0.0".to_string(),
233            git_commit: "1234567890".to_string(),
234            start_time_ms: last_activity_ts as u64,
235            total_cpu_millicores: 0,
236            total_memory_bytes: 0,
237            cpu_usage_millicores: 0,
238            memory_usage_bytes: 0,
239            hostname: "test_hostname".to_string(),
240            env_vars: Default::default(),
241            schedulable: None,
242        };
243
244        let key_prefix = NodeInfoKey::key_prefix_with_role(Role::Frontend);
245
246        in_memory
247            .put(PutRequest {
248                key: format!("{}{}", key_prefix, "0").into(),
249                value: active_frontend_node.try_into().unwrap(),
250                prev_kv: false,
251            })
252            .await
253            .unwrap();
254
255        let peers = client.active_frontends().await.unwrap();
256        assert_eq!(peers.len(), 0);
257    }
258}
meta_srv/discovery/lease.rs

meta_srv/discovery/
lease.rs
