datanode/
heartbeat.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;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicBool, Ordering};
18use std::time::Duration;
19
20use api::v1::meta::heartbeat_request::NodeWorkloads;
21use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest, NodeInfo, Peer, RegionRole, RegionStat};
22use common_base::Plugins;
23use common_meta::cache_invalidator::CacheInvalidatorRef;
24use common_meta::datanode::REGION_STATISTIC_KEY;
25use common_meta::distributed_time_constants::META_KEEP_ALIVE_INTERVAL_SECS;
26use common_meta::heartbeat::handler::invalidate_table_cache::InvalidateCacheHandler;
27use common_meta::heartbeat::handler::parse_mailbox_message::ParseMailboxMessageHandler;
28use common_meta::heartbeat::handler::{
29    HandlerGroupExecutor, HeartbeatResponseHandlerContext, HeartbeatResponseHandlerExecutorRef,
30};
31use common_meta::heartbeat::mailbox::{HeartbeatMailbox, MailboxRef};
32use common_meta::heartbeat::utils::outgoing_message_to_mailbox_message;
33use common_stat::ResourceStatRef;
34use common_telemetry::{debug, error, info, trace, warn};
35use common_workload::DatanodeWorkloadType;
36use meta_client::MetaClientRef;
37use meta_client::client::{HeartbeatSender, MetaClient};
38use servers::addrs;
39use snafu::ResultExt;
40use tokio::sync::{Notify, mpsc};
41use tokio::time::Instant;
42
43use self::handler::RegionHeartbeatResponseHandler;
44use crate::alive_keeper::{CountdownTaskHandlerExtRef, RegionAliveKeeper};
45use crate::config::DatanodeOptions;
46use crate::error::{self, MetaClientInitSnafu, Result};
47use crate::event_listener::RegionServerEventReceiver;
48use crate::metrics::{self, HEARTBEAT_RECV_COUNT, HEARTBEAT_SENT_COUNT};
49use crate::region_server::RegionServer;
50
51pub(crate) mod handler;
52pub(crate) mod task_tracker;
53
54/// The datanode heartbeat task which sending `[HeartbeatRequest]` to Metasrv periodically in background.
55pub struct HeartbeatTask {
56    node_id: u64,
57    workload_types: Vec<DatanodeWorkloadType>,
58    node_epoch: u64,
59    peer_addr: String,
60    running: Arc<AtomicBool>,
61    meta_client: MetaClientRef,
62    region_server: RegionServer,
63    interval: u64,
64    resp_handler_executor: HeartbeatResponseHandlerExecutorRef,
65    region_alive_keeper: Arc<RegionAliveKeeper>,
66    resource_stat: ResourceStatRef,
67}
68
69impl Drop for HeartbeatTask {
70    fn drop(&mut self) {
71        self.running.store(false, Ordering::Release);
72    }
73}
74
75impl HeartbeatTask {
76    /// Create a new heartbeat task instance.
77    pub async fn try_new(
78        opts: &DatanodeOptions,
79        region_server: RegionServer,
80        meta_client: MetaClientRef,
81        cache_invalidator: CacheInvalidatorRef,
82        plugins: Plugins,
83        resource_stat: ResourceStatRef,
84    ) -> Result<Self> {
85        let countdown_task_handler_ext = plugins.get::<CountdownTaskHandlerExtRef>();
86        let region_alive_keeper = Arc::new(RegionAliveKeeper::new(
87            region_server.clone(),
88            countdown_task_handler_ext,
89            opts.heartbeat.interval.as_millis() as u64,
90        ));
91        let resp_handler_executor = Arc::new(HandlerGroupExecutor::new(vec![
92            region_alive_keeper.clone(),
93            Arc::new(ParseMailboxMessageHandler),
94            Arc::new(
95                RegionHeartbeatResponseHandler::new(region_server.clone())
96                    .with_open_region_parallelism(opts.init_regions_parallelism),
97            ),
98            Arc::new(InvalidateCacheHandler::new(cache_invalidator)),
99        ]));
100
101        Ok(Self {
102            node_id: opts.node_id.unwrap_or(0),
103            workload_types: opts.workload_types.clone(),
104            // We use datanode's start time millis as the node's epoch.
105            node_epoch: common_time::util::current_time_millis() as u64,
106            peer_addr: addrs::resolve_addr(&opts.grpc.bind_addr, Some(&opts.grpc.server_addr)),
107            running: Arc::new(AtomicBool::new(false)),
108            meta_client,
109            region_server,
110            interval: opts.heartbeat.interval.as_millis() as u64,
111            resp_handler_executor,
112            region_alive_keeper,
113            resource_stat,
114        })
115    }
116
117    pub async fn create_streams(
118        meta_client: &MetaClient,
119        running: Arc<AtomicBool>,
120        handler_executor: HeartbeatResponseHandlerExecutorRef,
121        mailbox: MailboxRef,
122        mut notify: Option<Arc<Notify>>,
123        quit_signal: Arc<Notify>,
124    ) -> Result<HeartbeatSender> {
125        let client_id = meta_client.id();
126        let (tx, mut rx) = meta_client.heartbeat().await.context(MetaClientInitSnafu)?;
127
128        let mut last_received_lease = Instant::now();
129
130        let _handle = common_runtime::spawn_hb(async move {
131            while let Some(res) = rx.message().await.unwrap_or_else(|e| {
132                error!(e; "Error while reading heartbeat response");
133                None
134            }) {
135                if let Some(msg) = res.mailbox_message.as_ref() {
136                    info!("Received mailbox message: {msg:?}, meta_client id: {client_id:?}");
137                }
138                if let Some(lease) = res.region_lease.as_ref() {
139                    metrics::LAST_RECEIVED_HEARTBEAT_ELAPSED
140                        .set(last_received_lease.elapsed().as_millis() as i64);
141                    // Resets the timer.
142                    last_received_lease = Instant::now();
143
144                    let mut leader_region_lease_count = 0;
145                    let mut follower_region_lease_count = 0;
146                    for lease in &lease.regions {
147                        match lease.role() {
148                            RegionRole::Leader | RegionRole::DowngradingLeader => {
149                                leader_region_lease_count += 1
150                            }
151                            RegionRole::Follower => follower_region_lease_count += 1,
152                        }
153                    }
154
155                    metrics::HEARTBEAT_REGION_LEASES
156                        .with_label_values(&["leader"])
157                        .set(leader_region_lease_count);
158                    metrics::HEARTBEAT_REGION_LEASES
159                        .with_label_values(&["follower"])
160                        .set(follower_region_lease_count);
161                }
162                let ctx = HeartbeatResponseHandlerContext::new(mailbox.clone(), res);
163                if let Err(e) = Self::handle_response(ctx, handler_executor.clone()).await {
164                    error!(e; "Error while handling heartbeat response");
165                }
166                if let Some(notify) = notify.take() {
167                    notify.notify_one();
168                }
169                if !running.load(Ordering::Acquire) {
170                    info!("Heartbeat task shutdown");
171                }
172            }
173            quit_signal.notify_one();
174            info!("Heartbeat handling loop exit.");
175        });
176        Ok(tx)
177    }
178
179    async fn handle_response(
180        ctx: HeartbeatResponseHandlerContext,
181        handler_executor: HeartbeatResponseHandlerExecutorRef,
182    ) -> Result<()> {
183        trace!("Heartbeat response: {:?}", ctx.response);
184        handler_executor
185            .handle(ctx)
186            .await
187            .context(error::HandleHeartbeatResponseSnafu)
188    }
189
190    #[allow(deprecated)]
191    /// Start heartbeat task, spawn background task.
192    pub async fn start(
193        &self,
194        event_receiver: RegionServerEventReceiver,
195        notify: Option<Arc<Notify>>,
196    ) -> Result<()> {
197        let running = self.running.clone();
198        if running
199            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
200            .is_err()
201        {
202            warn!("Heartbeat task started multiple times");
203            return Ok(());
204        }
205        let interval = self.interval;
206        let node_id = self.node_id;
207        let node_epoch = self.node_epoch;
208        let addr = &self.peer_addr;
209        info!(
210            "Starting heartbeat to Metasrv with interval {interval}. My node id is {node_id}, address is {addr}."
211        );
212
213        let meta_client = self.meta_client.clone();
214        let region_server_clone = self.region_server.clone();
215
216        let handler_executor = self.resp_handler_executor.clone();
217
218        let (outgoing_tx, mut outgoing_rx) = mpsc::channel(16);
219        let mailbox = Arc::new(HeartbeatMailbox::new(outgoing_tx));
220
221        let quit_signal = Arc::new(Notify::new());
222
223        let mut tx = Self::create_streams(
224            &meta_client,
225            running.clone(),
226            handler_executor.clone(),
227            mailbox.clone(),
228            notify,
229            quit_signal.clone(),
230        )
231        .await?;
232
233        let self_peer = Some(Peer {
234            id: node_id,
235            addr: addr.clone(),
236        });
237        let epoch = self.region_alive_keeper.epoch();
238        let workload_types = self.workload_types.clone();
239
240        self.region_alive_keeper.start(Some(event_receiver)).await?;
241        let mut last_sent = Instant::now();
242        let total_cpu_millicores = self.resource_stat.get_total_cpu_millicores();
243        let total_memory_bytes = self.resource_stat.get_total_memory_bytes();
244        let resource_stat = self.resource_stat.clone();
245
246        common_runtime::spawn_hb(async move {
247            let sleep = tokio::time::sleep(Duration::from_millis(0));
248            tokio::pin!(sleep);
249
250            let build_info = common_version::build_info();
251            let heartbeat_request = HeartbeatRequest {
252                peer: self_peer,
253                node_epoch,
254                info: Some(NodeInfo {
255                    version: build_info.version.to_string(),
256                    git_commit: build_info.commit_short.to_string(),
257                    start_time_ms: node_epoch,
258                    total_cpu_millicores,
259                    total_memory_bytes,
260                    cpu_usage_millicores: 0,
261                    memory_usage_bytes: 0,
262                    // TODO(zyy17): Remove these deprecated fields when the deprecated fields are removed from the proto.
263                    cpus: total_cpu_millicores as u32,
264                    memory_bytes: total_memory_bytes as u64,
265                    hostname: hostname::get()
266                        .unwrap_or_default()
267                        .to_string_lossy()
268                        .to_string(),
269                }),
270                node_workloads: Some(NodeWorkloads::Datanode(DatanodeWorkloads {
271                    types: workload_types.iter().map(|w| w.to_i32()).collect(),
272                })),
273                ..Default::default()
274            };
275
276            loop {
277                if !running.load(Ordering::Relaxed) {
278                    info!("shutdown heartbeat task");
279                    break;
280                }
281                let req = tokio::select! {
282                    message = outgoing_rx.recv() => {
283                        if let Some(message) = message {
284                            match outgoing_message_to_mailbox_message(message) {
285                                Ok(message) => {
286                                    let req = HeartbeatRequest {
287                                        mailbox_message: Some(message),
288                                        ..heartbeat_request.clone()
289                                    };
290                                    HEARTBEAT_RECV_COUNT.with_label_values(&["success"]).inc();
291                                    Some(req)
292                                }
293                                Err(e) => {
294                                    error!(e; "Failed to encode mailbox messages!");
295                                    HEARTBEAT_RECV_COUNT.with_label_values(&["error"]).inc();
296                                    None
297                                }
298                            }
299                        } else {
300                            None
301                        }
302                    }
303                    _ = &mut sleep => {
304                        let region_stats = Self::load_region_stats(&region_server_clone);
305                        let topic_stats = region_server_clone.topic_stats();
306                        let now = Instant::now();
307                        let duration_since_epoch = (now - epoch).as_millis() as u64;
308                        let mut req = HeartbeatRequest {
309                            region_stats,
310                            topic_stats,
311                            duration_since_epoch,
312                            ..heartbeat_request.clone()
313                        };
314
315                        if let Some(info) = req.info.as_mut() {
316                            info.cpu_usage_millicores = resource_stat.get_cpu_usage_millicores();
317                            info.memory_usage_bytes = resource_stat.get_memory_usage_bytes();
318                        }
319
320                        sleep.as_mut().reset(now + Duration::from_millis(interval));
321                        Some(req)
322                    }
323                    // If the heartbeat stream is broken, send a dummy heartbeat request to re-create the heartbeat stream.
324                    _ = quit_signal.notified() => {
325                        let req = HeartbeatRequest::default();
326                        Some(req)
327                    }
328                };
329                if let Some(req) = req {
330                    metrics::LAST_SENT_HEARTBEAT_ELAPSED
331                        .set(last_sent.elapsed().as_millis() as i64);
332                    // Resets the timer.
333                    last_sent = Instant::now();
334                    debug!("Sending heartbeat request: {:?}", req);
335                    if let Err(e) = tx.send(req).await {
336                        error!(e; "Failed to send heartbeat to metasrv");
337                        match Self::create_streams(
338                            &meta_client,
339                            running.clone(),
340                            handler_executor.clone(),
341                            mailbox.clone(),
342                            None,
343                            quit_signal.clone(),
344                        )
345                        .await
346                        {
347                            Ok(new_tx) => {
348                                info!("Reconnected to metasrv");
349                                tx = new_tx;
350                                // Triggers to send heartbeat immediately.
351                                sleep.as_mut().reset(Instant::now());
352                            }
353                            Err(e) => {
354                                // Before the META_LEASE_SECS expires,
355                                // any retries are meaningless, it always reads the old meta leader address.
356                                // Triggers to retry after META_KEEP_ALIVE_INTERVAL_SECS.
357                                sleep.as_mut().reset(
358                                    Instant::now()
359                                        + Duration::from_secs(META_KEEP_ALIVE_INTERVAL_SECS),
360                                );
361                                error!(e; "Failed to reconnect to metasrv!");
362                            }
363                        }
364                    } else {
365                        HEARTBEAT_SENT_COUNT.inc();
366                    }
367                }
368            }
369        });
370
371        Ok(())
372    }
373
374    fn load_region_stats(region_server: &RegionServer) -> Vec<RegionStat> {
375        region_server
376            .reportable_regions()
377            .into_iter()
378            .map(|stat| {
379                let region_stat = region_server
380                    .region_statistic(stat.region_id)
381                    .unwrap_or_default();
382                let mut extensions = HashMap::new();
383                if let Some(serialized) = region_stat.serialize_to_vec() {
384                    extensions.insert(REGION_STATISTIC_KEY.to_string(), serialized);
385                }
386
387                RegionStat {
388                    region_id: stat.region_id.as_u64(),
389                    engine: stat.engine,
390                    role: RegionRole::from(stat.role).into(),
391                    // TODO(weny): w/rcus
392                    rcus: 0,
393                    wcus: 0,
394                    approximate_bytes: region_stat.estimated_disk_size() as i64,
395                    extensions,
396                }
397            })
398            .collect()
399    }
400
401    pub fn close(&self) -> Result<()> {
402        let running = self.running.clone();
403        if running
404            .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire)
405            .is_err()
406        {
407            warn!("Call close heartbeat task multiple times");
408        }
409
410        Ok(())
411    }
412}