meta_srv/handler/
collect_cluster_info_handler.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::collections::HashMap;
16
17use api::v1::meta::{HeartbeatRequest, NodeInfo as PbNodeInfo, Role};
18use common_meta::cluster::{
19    DatanodeStatus, FlownodeStatus, FrontendStatus, NodeInfo, NodeInfoKey, NodeStatus,
20};
21use common_meta::datanode::EnvVars;
22use common_meta::heartbeat::utils::{get_flownode_workloads, get_frontend_workloads};
23use common_meta::peer::Peer;
24use common_meta::rpc::store::PutRequest;
25use common_telemetry::warn;
26use snafu::ResultExt;
27use store_api::region_engine::RegionRole;
28
29use crate::Result;
30use crate::error::{InvalidClusterInfoFormatSnafu, SaveClusterInfoSnafu};
31use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
32use crate::metasrv::Context;
33
34/// The handler to collect cluster info from the heartbeat request of frontend.
35pub struct CollectFrontendClusterInfoHandler;
36
37#[async_trait::async_trait]
38impl HeartbeatHandler for CollectFrontendClusterInfoHandler {
39    fn is_acceptable(&self, role: Role) -> bool {
40        role == Role::Frontend
41    }
42
43    async fn handle(
44        &self,
45        req: &HeartbeatRequest,
46        ctx: &mut Context,
47        _acc: &mut HeartbeatAccumulator,
48    ) -> Result<HandleControl> {
49        let Some(update) = NodeInfoUpdate::from_request(req) else {
50            return Ok(HandleControl::Continue);
51        };
52
53        let frontend_workloads = get_frontend_workloads(req.node_workloads.as_ref());
54
55        update
56            .save_to_memory(
57                ctx,
58                common_time::util::current_time_millis(),
59                NodeStatus::Frontend(FrontendStatus {
60                    workloads: frontend_workloads,
61                }),
62            )
63            .await?;
64
65        Ok(HandleControl::Continue)
66    }
67}
68
69/// The handler to collect cluster info from the heartbeat request of flownode.
70pub struct CollectFlownodeClusterInfoHandler;
71#[async_trait::async_trait]
72impl HeartbeatHandler for CollectFlownodeClusterInfoHandler {
73    fn is_acceptable(&self, role: Role) -> bool {
74        role == Role::Flownode
75    }
76
77    async fn handle(
78        &self,
79        req: &HeartbeatRequest,
80        ctx: &mut Context,
81        _acc: &mut HeartbeatAccumulator,
82    ) -> Result<HandleControl> {
83        let Some(update) = NodeInfoUpdate::from_request(req) else {
84            return Ok(HandleControl::Continue);
85        };
86        let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref());
87
88        update
89            .save_to_memory(
90                ctx,
91                common_time::util::current_time_millis(),
92                NodeStatus::Flownode(FlownodeStatus {
93                    workloads: flownode_workloads,
94                }),
95            )
96            .await?;
97
98        Ok(HandleControl::Continue)
99    }
100}
101
102/// The handler to collect cluster info from the heartbeat request of datanode.
103pub struct CollectDatanodeClusterInfoHandler;
104
105#[async_trait::async_trait]
106impl HeartbeatHandler for CollectDatanodeClusterInfoHandler {
107    fn is_acceptable(&self, role: Role) -> bool {
108        role == Role::Datanode
109    }
110
111    async fn handle(
112        &self,
113        req: &HeartbeatRequest,
114        ctx: &mut Context,
115        acc: &mut HeartbeatAccumulator,
116    ) -> Result<HandleControl> {
117        let Some(update) = NodeInfoUpdate::from_request(req) else {
118            return Ok(HandleControl::Continue);
119        };
120
121        let Some(stat) = &acc.stat else {
122            return Ok(HandleControl::Continue);
123        };
124
125        let leader_regions = stat
126            .region_stats
127            .iter()
128            .filter(|s| matches!(s.role, RegionRole::Leader | RegionRole::StagingLeader))
129            .count();
130        let follower_regions = stat.region_stats.len() - leader_regions;
131
132        update
133            .save_to_memory(
134                ctx,
135                stat.timestamp_millis,
136                NodeStatus::Datanode(DatanodeStatus {
137                    rcus: stat.rcus,
138                    wcus: stat.wcus,
139                    leader_regions,
140                    follower_regions,
141                    workloads: stat.datanode_workloads.clone(),
142                }),
143            )
144            .await?;
145
146        Ok(HandleControl::Continue)
147    }
148}
149
150/// Common fields collected for one heartbeat's in-memory node information
151/// update. Persistent routing addresses have a separate lifecycle.
152struct NodeInfoUpdate {
153    key: NodeInfoKey,
154    peer: Peer,
155    info: PbNodeInfo,
156    env_vars: HashMap<String, String>,
157}
158
159impl NodeInfoUpdate {
160    fn from_request(request: &HeartbeatRequest) -> Option<Self> {
161        let key = NodeInfoKey::new(request)?;
162        let peer = request.peer.clone()?;
163        let info = request.info.clone()?;
164        let env_vars = EnvVars::from_extensions(&request.extensions)
165            .inspect_err(|e| {
166                warn!(e; "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer);
167            })
168            .unwrap_or_default()
169            .map(|e| e.vars)
170            .unwrap_or_default();
171        Some(Self {
172            key,
173            peer,
174            info,
175            env_vars,
176        })
177    }
178
179    async fn save_to_memory(
180        self,
181        ctx: &Context,
182        last_activity_ts: i64,
183        status: NodeStatus,
184    ) -> Result<()> {
185        let info = self.info;
186        let value = NodeInfo {
187            peer: self.peer,
188            last_activity_ts,
189            status,
190            version: info.version,
191            git_commit: info.git_commit,
192            start_time_ms: info.start_time_ms,
193            total_cpu_millicores: info.total_cpu_millicores,
194            total_memory_bytes: info.total_memory_bytes,
195            cpu_usage_millicores: info.cpu_usage_millicores,
196            memory_usage_bytes: info.memory_usage_bytes,
197            hostname: info.hostname,
198            env_vars: self.env_vars,
199        };
200        let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?;
201        ctx.in_memory
202            .put(PutRequest {
203                key: (&self.key).into(),
204                value,
205                ..Default::default()
206            })
207            .await
208            .context(SaveClusterInfoSnafu)?;
209        Ok(())
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use api::v1::meta::RequestHeader;
216    use common_meta::datanode::Stat;
217
218    use super::*;
219    use crate::handler::test_utils::TestEnv;
220
221    #[tokio::test]
222    async fn test_node_info_refreshes_without_epoch_change() {
223        let mut ctx = TestEnv::new().ctx();
224        let handlers: Vec<(Role, Box<dyn HeartbeatHandler>)> = vec![
225            (Role::Frontend, Box::new(CollectFrontendClusterInfoHandler)),
226            (Role::Datanode, Box::new(CollectDatanodeClusterInfoHandler)),
227            (Role::Flownode, Box::new(CollectFlownodeClusterInfoHandler)),
228        ];
229        for (role, handler) in handlers {
230            let mut req = HeartbeatRequest {
231                header: Some(RequestHeader {
232                    role: role as i32,
233                    ..Default::default()
234                }),
235                peer: Some(Peer {
236                    id: 7,
237                    addr: "node:3001".into(),
238                }),
239                info: Some(PbNodeInfo {
240                    version: "v1".into(),
241                    hostname: "host".into(),
242                    ..Default::default()
243                }),
244                node_epoch: 10,
245                ..Default::default()
246            };
247            let mut acc = HeartbeatAccumulator {
248                stat: Some(Stat {
249                    timestamp_millis: 123,
250                    rcus: 5,
251                    wcus: 6,
252                    ..Default::default()
253                }),
254                ..Default::default()
255            };
256            handler.handle(&req, &mut ctx, &mut acc).await.unwrap();
257            let key: Vec<u8> = (&NodeInfoKey::new(&req).unwrap()).into();
258            let first =
259                NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap();
260            assert_eq!(first.peer, req.peer.clone().unwrap());
261            assert_eq!(first.version, "v1");
262            assert_eq!(first.hostname, "host");
263            req.info.as_mut().unwrap().version = "v2".into();
264            req.info.as_mut().unwrap().memory_usage_bytes = 42;
265            acc.stat.as_mut().unwrap().timestamp_millis = 456;
266            handler.handle(&req, &mut ctx, &mut acc).await.unwrap();
267            let second =
268                NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap();
269            assert_eq!(second.version, "v2");
270            assert_eq!(second.memory_usage_bytes, 42);
271            assert_eq!(
272                serde_json::to_value(&second.status).unwrap(),
273                serde_json::to_value(&first.status).unwrap()
274            );
275            if role == Role::Datanode {
276                assert_eq!(second.last_activity_ts, 456);
277            }
278            // Node information must not be persisted to the configured backend.
279            assert!(ctx.kv_backend.get(&key).await.unwrap().is_none());
280        }
281    }
282}
meta_srv/handler/collect_cluster_info_handler.rs

meta_srv/handler/
collect_cluster_info_handler.rs
