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;
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((key, peer, info, env_vars)) = extract_base_info(req) else {
50            return Ok(HandleControl::Continue);
51        };
52
53        let value = NodeInfo {
54            peer,
55            last_activity_ts: common_time::util::current_time_millis(),
56            status: NodeStatus::Frontend(FrontendStatus {}),
57            version: info.version,
58            git_commit: info.git_commit,
59            start_time_ms: info.start_time_ms,
60            total_cpu_millicores: info.total_cpu_millicores,
61            total_memory_bytes: info.total_memory_bytes,
62            cpu_usage_millicores: info.cpu_usage_millicores,
63            memory_usage_bytes: info.memory_usage_bytes,
64            hostname: info.hostname,
65            env_vars,
66        };
67
68        put_into_memory_store(ctx, key, value).await?;
69
70        Ok(HandleControl::Continue)
71    }
72}
73
74/// The handler to collect cluster info from the heartbeat request of flownode.
75pub struct CollectFlownodeClusterInfoHandler;
76#[async_trait::async_trait]
77impl HeartbeatHandler for CollectFlownodeClusterInfoHandler {
78    fn is_acceptable(&self, role: Role) -> bool {
79        role == Role::Flownode
80    }
81
82    async fn handle(
83        &self,
84        req: &HeartbeatRequest,
85        ctx: &mut Context,
86        _acc: &mut HeartbeatAccumulator,
87    ) -> Result<HandleControl> {
88        let Some((key, peer, info, env_vars)) = extract_base_info(req) else {
89            return Ok(HandleControl::Continue);
90        };
91        let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref());
92
93        let value = NodeInfo {
94            peer,
95            last_activity_ts: common_time::util::current_time_millis(),
96            status: NodeStatus::Flownode(FlownodeStatus {
97                workloads: flownode_workloads,
98            }),
99            version: info.version,
100            git_commit: info.git_commit,
101            start_time_ms: info.start_time_ms,
102            total_cpu_millicores: info.total_cpu_millicores,
103            total_memory_bytes: info.total_memory_bytes,
104            cpu_usage_millicores: info.cpu_usage_millicores,
105            memory_usage_bytes: info.memory_usage_bytes,
106            hostname: info.hostname,
107            env_vars,
108        };
109
110        put_into_memory_store(ctx, key, value).await?;
111
112        Ok(HandleControl::Continue)
113    }
114}
115
116/// The handler to collect cluster info from the heartbeat request of datanode.
117pub struct CollectDatanodeClusterInfoHandler;
118
119#[async_trait::async_trait]
120impl HeartbeatHandler for CollectDatanodeClusterInfoHandler {
121    fn is_acceptable(&self, role: Role) -> bool {
122        role == Role::Datanode
123    }
124
125    async fn handle(
126        &self,
127        req: &HeartbeatRequest,
128        ctx: &mut Context,
129        acc: &mut HeartbeatAccumulator,
130    ) -> Result<HandleControl> {
131        let Some((key, peer, info, env_vars)) = extract_base_info(req) else {
132            return Ok(HandleControl::Continue);
133        };
134
135        let Some(stat) = &acc.stat else {
136            return Ok(HandleControl::Continue);
137        };
138
139        let leader_regions = stat
140            .region_stats
141            .iter()
142            .filter(|s| matches!(s.role, RegionRole::Leader | RegionRole::StagingLeader))
143            .count();
144        let follower_regions = stat.region_stats.len() - leader_regions;
145
146        let value = NodeInfo {
147            peer,
148            last_activity_ts: stat.timestamp_millis,
149            status: NodeStatus::Datanode(DatanodeStatus {
150                rcus: stat.rcus,
151                wcus: stat.wcus,
152                leader_regions,
153                follower_regions,
154                workloads: stat.datanode_workloads.clone(),
155            }),
156            version: info.version,
157            git_commit: info.git_commit,
158            start_time_ms: info.start_time_ms,
159            total_cpu_millicores: info.total_cpu_millicores,
160            total_memory_bytes: info.total_memory_bytes,
161            cpu_usage_millicores: info.cpu_usage_millicores,
162            memory_usage_bytes: info.memory_usage_bytes,
163            hostname: info.hostname,
164            env_vars,
165        };
166
167        put_into_memory_store(ctx, key, value).await?;
168
169        Ok(HandleControl::Continue)
170    }
171}
172
173fn extract_base_info(
174    request: &HeartbeatRequest,
175) -> Option<(NodeInfoKey, Peer, PbNodeInfo, HashMap<String, String>)> {
176    let HeartbeatRequest { peer, info, .. } = request;
177    let key = NodeInfoKey::new(request)?;
178    let Some(peer) = &peer else {
179        return None;
180    };
181    let Some(info) = &info else {
182        return None;
183    };
184
185    let env_vars = EnvVars::from_extensions(&request.extensions)
186        .inspect_err(|e| {
187            warn!(e;
188                "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer
189            );
190        })
191        .unwrap_or_default()
192        .map(|e| e.vars)
193        .unwrap_or_default();
194
195    Some((key, peer.clone(), info.clone(), env_vars))
196}
197
198async fn put_into_memory_store(ctx: &mut Context, key: NodeInfoKey, value: NodeInfo) -> Result<()> {
199    let key = (&key).into();
200    let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?;
201    let put_req = PutRequest {
202        key,
203        value,
204        ..Default::default()
205    };
206
207    ctx.in_memory
208        .put(put_req)
209        .await
210        .context(SaveClusterInfoSnafu)?;
211
212    Ok(())
213}
meta_srv/handler/collect_cluster_info_handler.rs

meta_srv/handler/
collect_cluster_info_handler.rs
