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, HashSet};
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::mailbox_message::Payload;
22use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest, NodeInfo, Peer, RegionRole, RegionStat};
23use common_base::Plugins;
24use common_meta::cache_invalidator::CacheInvalidatorRef;
25use common_meta::datanode::{EnvVars, REGION_STATISTIC_KEY};
26use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL;
27use common_meta::heartbeat::handler::invalidate_table_cache::InvalidateCacheHandler;
28use common_meta::heartbeat::handler::parse_mailbox_message::ParseMailboxMessageHandler;
29use common_meta::heartbeat::handler::suspend::SuspendHandler;
30use common_meta::heartbeat::handler::{
31    HandlerGroupExecutor, HeartbeatResponseHandlerContext, HeartbeatResponseHandlerExecutorRef,
32};
33use common_meta::heartbeat::mailbox::{HeartbeatMailbox, MailboxRef};
34use common_meta::heartbeat::utils::outgoing_message_to_mailbox_message;
35use common_meta::key::MetadataValue;
36use common_meta::key::datanode_epoch::{DatanodeEpochKey, DatanodeEpochValue};
37use common_meta::kv_backend::KvBackendRef;
38use common_stat::ResourceStatRef;
39use common_telemetry::{debug, error, info, trace, warn};
40use common_workload::DatanodeWorkloadType;
41use meta_client::MetaClientRef;
42use meta_client::client::heartbeat::HeartbeatConfig;
43use meta_client::client::{HeartbeatSender, MetaClient};
44use servers::addrs;
45use snafu::{OptionExt as _, ResultExt};
46use tokio::sync::{Notify, mpsc};
47use tokio::time::Instant;
48
49use self::handler::RegionHeartbeatResponseHandler;
50use crate::alive_keeper::{CountdownTaskHandlerExtRef, RegionAliveKeeper};
51use crate::config::DatanodeOptions;
52use crate::error::{
53    self, GetMetadataSnafu, MetaClientInitSnafu, RegionEngineNotFoundSnafu, Result,
54};
55use crate::event_listener::RegionServerEventReceiver;
56use crate::metrics::{self, HEARTBEAT_RECV_COUNT, HEARTBEAT_SENT_COUNT};
57use crate::region_server::RegionServer;
58
59pub(crate) mod handler;
60pub(crate) mod task_tracker;
61
62/// Returns the epoch of this process: its start time, or one more than the epoch the metasrv has
63/// recorded for the node when the record is not lower than the start time, as after a restart on
64/// a slower clock or after the clock stepped back. The metasrv fences an epoch lower than the
65/// recorded one.
66async fn node_epoch_above_record(
67    kv_backend: &KvBackendRef,
68    node_id: u64,
69    start_time_ms: u64,
70) -> Result<u64> {
71    let key = DatanodeEpochKey::new(node_id).to_bytes();
72    let Some(kv) = kv_backend.get(&key).await.context(GetMetadataSnafu)? else {
73        return Ok(start_time_ms);
74    };
75    let recorded = DatanodeEpochValue::try_from_raw_value(&kv.value).context(GetMetadataSnafu)?;
76    Ok(start_time_ms.max(recorded.node_epoch + 1))
77}
78
79/// The datanode heartbeat task which sending `[HeartbeatRequest]` to Metasrv periodically in background.
80pub struct HeartbeatTask {
81    node_id: u64,
82    workload_types: Vec<DatanodeWorkloadType>,
83    node_epoch: u64,
84    start_time_ms: u64,
85    peer_addr: String,
86    running: Arc<AtomicBool>,
87    meta_client: MetaClientRef,
88    region_server: RegionServer,
89    resp_handler_executor: HeartbeatResponseHandlerExecutorRef,
90    region_alive_keeper: Arc<RegionAliveKeeper>,
91    resource_stat: ResourceStatRef,
92    env_vars: EnvVars,
93}
94
95impl Drop for HeartbeatTask {
96    fn drop(&mut self) {
97        self.running.store(false, Ordering::Release);
98    }
99}
100
101impl HeartbeatTask {
102    /// Create a new heartbeat task instance.
103    pub async fn try_new(
104        opts: &DatanodeOptions,
105        region_server: RegionServer,
106        meta_client: MetaClientRef,
107        kv_backend: KvBackendRef,
108        cache_invalidator: CacheInvalidatorRef,
109        plugins: Plugins,
110        resource_stat: ResourceStatRef,
111    ) -> Result<Self> {
112        let node_id = opts.node_id.unwrap_or(0);
113        let start_time_ms = common_time::util::current_time_millis() as u64;
114        let node_epoch = node_epoch_above_record(&kv_backend, node_id, start_time_ms).await?;
115        let countdown_task_handler_ext = plugins.get::<CountdownTaskHandlerExtRef>();
116        let region_alive_keeper = Arc::new(RegionAliveKeeper::new(
117            region_server.clone(),
118            countdown_task_handler_ext,
119            BASE_HEARTBEAT_INTERVAL,
120        ));
121        let resp_handler_executor = Arc::new(HandlerGroupExecutor::new(vec![
122            region_alive_keeper.clone(),
123            Arc::new(ParseMailboxMessageHandler),
124            Arc::new(SuspendHandler::new(region_server.suspend_state())),
125            Arc::new(
126                RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend)
127                    .with_open_region_parallelism(opts.init_regions_parallelism),
128            ),
129            Arc::new(InvalidateCacheHandler::new(cache_invalidator)),
130        ]));
131
132        Ok(Self {
133            node_id,
134            workload_types: opts.workload_types.clone(),
135            node_epoch,
136            start_time_ms,
137            peer_addr: addrs::resolve_addr(&opts.grpc.bind_addr, Some(&opts.grpc.server_addr)),
138            running: Arc::new(AtomicBool::new(false)),
139            meta_client,
140            region_server,
141            resp_handler_executor,
142            region_alive_keeper,
143            resource_stat,
144            env_vars: EnvVars::from_config(&opts.heartbeat_env_vars),
145        })
146    }
147
148    pub async fn create_streams(
149        meta_client: &MetaClient,
150        local_gc_enabled: bool,
151        running: Arc<AtomicBool>,
152        handler_executor: HeartbeatResponseHandlerExecutorRef,
153        mailbox: MailboxRef,
154        mut notify: Option<Arc<Notify>>,
155        quit_signal: Arc<Notify>,
156    ) -> Result<(HeartbeatSender, HeartbeatConfig)> {
157        let client_id = meta_client.id();
158        let (tx, mut rx, config) = meta_client.heartbeat().await.context(MetaClientInitSnafu)?;
159        if config.gc_enabled != local_gc_enabled {
160            return error::GcConfigMismatchSnafu {
161                metasrv_gc_enabled: config.gc_enabled,
162                datanode_gc_enabled: local_gc_enabled,
163            }
164            .fail();
165        }
166
167        let mut last_received_lease = Instant::now();
168
169        let _handle = common_runtime::spawn_hb(async move {
170            while let Some(res) = rx.message().await.unwrap_or_else(|e| {
171                error!(e; "Error while reading heartbeat response");
172                None
173            }) {
174                if let Some(msg) = res.mailbox_message.as_ref() {
175                    let payload_len = msg
176                        .payload
177                        .as_ref()
178                        .map(|Payload::Json(payload)| payload.len())
179                        .unwrap_or_default();
180                    info!(
181                        message_id = msg.id,
182                        payload_len, client_id, "Received mailbox message"
183                    );
184                }
185                if let Some(lease) = res.region_lease.as_ref() {
186                    metrics::LAST_RECEIVED_HEARTBEAT_ELAPSED
187                        .set(last_received_lease.elapsed().as_millis() as i64);
188                    // Resets the timer.
189                    last_received_lease = Instant::now();
190
191                    let mut leader_region_lease_count = 0;
192                    let mut follower_region_lease_count = 0;
193                    for lease in &lease.regions {
194                        match lease.role() {
195                            RegionRole::Leader
196                            | RegionRole::StagingLeader
197                            | RegionRole::DowngradingLeader => leader_region_lease_count += 1,
198                            RegionRole::Follower => follower_region_lease_count += 1,
199                        }
200                    }
201
202                    metrics::HEARTBEAT_REGION_LEASES
203                        .with_label_values(&["leader"])
204                        .set(leader_region_lease_count);
205                    metrics::HEARTBEAT_REGION_LEASES
206                        .with_label_values(&["follower"])
207                        .set(follower_region_lease_count);
208                }
209                let ctx = HeartbeatResponseHandlerContext::new(mailbox.clone(), res);
210                if let Err(e) = Self::handle_response(ctx, handler_executor.clone()).await {
211                    error!(e; "Error while handling heartbeat response");
212                }
213                if let Some(notify) = notify.take() {
214                    notify.notify_one();
215                }
216                if !running.load(Ordering::Acquire) {
217                    info!("Heartbeat task shutdown");
218                }
219            }
220            quit_signal.notify_one();
221            info!("Heartbeat handling loop exit.");
222        });
223        Ok((tx, config))
224    }
225
226    async fn handle_response(
227        ctx: HeartbeatResponseHandlerContext,
228        handler_executor: HeartbeatResponseHandlerExecutorRef,
229    ) -> Result<()> {
230        let mailbox_message_id = ctx.response.mailbox_message.as_ref().map(|msg| msg.id);
231        trace!(?mailbox_message_id, "Handling heartbeat response");
232        handler_executor
233            .handle(ctx)
234            .await
235            .context(error::HandleHeartbeatResponseSnafu)
236    }
237
238    #[allow(deprecated)]
239    /// Start heartbeat task, spawn background task.
240    pub async fn start(
241        &self,
242        event_receiver: RegionServerEventReceiver,
243        notify: Option<Arc<Notify>>,
244    ) -> Result<()> {
245        let running = self.running.clone();
246        if running
247            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
248            .is_err()
249        {
250            warn!("Heartbeat task started multiple times");
251            return Ok(());
252        }
253        let node_id = self.node_id;
254        let node_epoch = self.node_epoch;
255        let start_time_ms = self.start_time_ms;
256        let addr = &self.peer_addr;
257
258        let meta_client = self.meta_client.clone();
259        let region_server_clone = self.region_server.clone();
260
261        let handler_executor = self.resp_handler_executor.clone();
262
263        let (outgoing_tx, mut outgoing_rx) = mpsc::channel(16);
264        let mailbox = Arc::new(HeartbeatMailbox::new(outgoing_tx));
265
266        let quit_signal = Arc::new(Notify::new());
267        let local_gc_enabled = self
268            .region_server
269            .mito_engine()
270            .context(RegionEngineNotFoundSnafu { name: "mito" })?
271            .mito_config()
272            .gc
273            .enable;
274
275        let (mut tx, config) = Self::create_streams(
276            &meta_client,
277            local_gc_enabled,
278            running.clone(),
279            handler_executor.clone(),
280            mailbox.clone(),
281            notify,
282            quit_signal.clone(),
283        )
284        .await?;
285
286        let interval = config.interval.as_millis() as u64;
287        let mut retry_interval = config.retry_interval;
288
289        // Update RegionAliveKeeper with the interval from Metasrv
290        self.region_alive_keeper.update_heartbeat_interval(interval);
291
292        info!(
293            "Starting heartbeat to Metasrv with config: {}. My node id is {}, address is {}.",
294            config, node_id, addr
295        );
296
297        let self_peer = Some(Peer {
298            id: node_id,
299            addr: addr.clone(),
300        });
301        let epoch = self.region_alive_keeper.epoch();
302        let workload_types = self.workload_types.clone();
303
304        self.region_alive_keeper.start(Some(event_receiver)).await?;
305        let mut last_sent = Instant::now();
306        let total_cpu_millicores = self.resource_stat.get_total_cpu_millicores();
307        let total_memory_bytes = self.resource_stat.get_total_memory_bytes();
308        let resource_stat = self.resource_stat.clone();
309        let region_alive_keeper = self.region_alive_keeper.clone();
310        let gc_limiter = self
311            .region_server
312            .mito_engine()
313            .context(RegionEngineNotFoundSnafu { name: "mito" })?
314            .gc_limiter();
315        let mut env_var_extensions = HashMap::new();
316        self.env_vars.into_extensions(&mut env_var_extensions);
317
318        common_runtime::spawn_hb(async move {
319            let sleep = tokio::time::sleep(Duration::from_millis(0));
320            tokio::pin!(sleep);
321
322            let build_info = common_version::build_info();
323
324            let heartbeat_request = HeartbeatRequest {
325                peer: self_peer,
326                node_epoch,
327                info: Some(NodeInfo {
328                    version: build_info.version.to_string(),
329                    git_commit: build_info.commit_short.to_string(),
330                    start_time_ms,
331                    total_cpu_millicores,
332                    total_memory_bytes,
333                    cpu_usage_millicores: 0,
334                    memory_usage_bytes: 0,
335                    // TODO(zyy17): Remove these deprecated fields when the deprecated fields are removed from the proto.
336                    cpus: total_cpu_millicores as u32,
337                    memory_bytes: total_memory_bytes as u64,
338                    hostname: hostname::get()
339                        .unwrap_or_default()
340                        .to_string_lossy()
341                        .to_string(),
342                }),
343                node_workloads: Some(NodeWorkloads::Datanode(DatanodeWorkloads {
344                    types: workload_types.iter().map(|w| w.to_i32()).collect(),
345                })),
346                ..Default::default()
347            };
348
349            loop {
350                if !running.load(Ordering::Relaxed) {
351                    info!("shutdown heartbeat task");
352                    break;
353                }
354                let req = tokio::select! {
355                    message = outgoing_rx.recv() => {
356                        if let Some(message) = message {
357                            match outgoing_message_to_mailbox_message(message) {
358                                Ok(message) => {
359                                    let mut extensions = env_var_extensions.clone();
360                                    let gc_stat = gc_limiter.gc_stat();
361                                    gc_stat.into_extensions(&mut extensions);
362
363                                    let req = HeartbeatRequest {
364                                        mailbox_message: Some(message),
365                                        extensions,
366                                        ..heartbeat_request.clone()
367                                    };
368                                    HEARTBEAT_RECV_COUNT.with_label_values(&["success"]).inc();
369                                    Some(req)
370                                }
371                                Err(e) => {
372                                    error!(e; "Failed to encode mailbox messages!");
373                                    HEARTBEAT_RECV_COUNT.with_label_values(&["error"]).inc();
374                                    None
375                                }
376                            }
377                        } else {
378                            None
379                        }
380                    }
381                    _ = &mut sleep => {
382                        // Read opening first so a completed open is covered by either snapshot.
383                        let opening_regions = region_server_clone.opening_regions();
384                        let region_stats = Self::load_region_stats(&region_server_clone);
385                        let opening_regions = if opening_regions.is_empty() {
386                            Vec::new()
387                        } else {
388                            let reported_regions: HashSet<_> = region_stats.iter()
389                                .map(|stat| stat.region_id)
390                                .collect();
391                            opening_regions.into_iter()
392                                .filter(|region_id| !reported_regions.contains(&region_id.as_u64()))
393                                .map(|region_id| api::v1::meta::OpeningRegion {
394                                    region_id: region_id.as_u64(),
395                                })
396                                .collect()
397                        };
398                        let topic_stats = region_server_clone.topic_stats();
399                        let now = Instant::now();
400                        let duration_since_epoch = (now - epoch).as_millis() as u64;
401
402                        let mut extensions = env_var_extensions.clone();
403                        let gc_stat = gc_limiter.gc_stat();
404                        gc_stat.into_extensions(&mut extensions);
405
406                        let mut req = HeartbeatRequest {
407                            opening_regions,
408                            region_stats,
409                            topic_stats,
410                            duration_since_epoch,
411                            extensions,
412                            ..heartbeat_request.clone()
413                        };
414
415                        if let Some(info) = req.info.as_mut() {
416                            info.cpu_usage_millicores = resource_stat.get_cpu_usage_millicores();
417                            info.memory_usage_bytes = resource_stat.get_memory_usage_bytes();
418                        }
419
420                        sleep.as_mut().reset(now + Duration::from_millis(interval));
421                        Some(req)
422                    }
423                    // If the heartbeat stream is broken, send a dummy heartbeat request to re-create the heartbeat stream.
424                    _ = quit_signal.notified() => {
425                        let req = HeartbeatRequest::default();
426                        Some(req)
427                    }
428                };
429                if let Some(req) = req {
430                    metrics::LAST_SENT_HEARTBEAT_ELAPSED
431                        .set(last_sent.elapsed().as_millis() as i64);
432                    // Resets the timer.
433                    last_sent = Instant::now();
434                    let mailbox_message_id = req.mailbox_message.as_ref().map(|msg| msg.id);
435                    debug!(?mailbox_message_id, "Sending heartbeat request");
436                    if let Err(e) = tx.send(req).await {
437                        error!(e; "Failed to send heartbeat to metasrv");
438                        match Self::create_streams(
439                            &meta_client,
440                            local_gc_enabled,
441                            running.clone(),
442                            handler_executor.clone(),
443                            mailbox.clone(),
444                            None,
445                            quit_signal.clone(),
446                        )
447                        .await
448                        {
449                            Ok((new_tx, new_config)) => {
450                                info!("Reconnected to metasrv, heartbeat config: {}", new_config);
451                                tx = new_tx;
452                                // Update retry_interval from new config
453                                retry_interval = new_config.retry_interval;
454                                // Update region_alive_keeper's heartbeat interval
455                                region_alive_keeper.update_heartbeat_interval(
456                                    new_config.interval.as_millis() as u64,
457                                );
458                                // Triggers to send heartbeat immediately.
459                                sleep.as_mut().reset(Instant::now());
460                            }
461                            Err(e) => {
462                                // Before the META_LEASE_SECS expires,
463                                // any retries are meaningless, it always reads the old meta leader address.
464                                // Triggers to retry after retry_interval from Metasrv config.
465                                sleep.as_mut().reset(Instant::now() + retry_interval);
466                                error!(e; "Failed to reconnect to metasrv!");
467                            }
468                        }
469                    } else {
470                        HEARTBEAT_SENT_COUNT.inc();
471                    }
472                }
473            }
474        });
475
476        Ok(())
477    }
478
479    fn load_region_stats(region_server: &RegionServer) -> Vec<RegionStat> {
480        region_server
481            .reportable_regions()
482            .into_iter()
483            .map(|stat| {
484                let region_stat = region_server
485                    .region_statistic(stat.region_id)
486                    .unwrap_or_default();
487                let mut extensions = HashMap::new();
488                if let Some(serialized) = region_stat.serialize_to_vec() {
489                    extensions.insert(REGION_STATISTIC_KEY.to_string(), serialized);
490                }
491                RegionStat {
492                    region_id: stat.region_id.as_u64(),
493                    engine: stat.engine,
494                    role: RegionRole::from(stat.role).into(),
495                    rcus: 0,
496                    wcus: 0,
497                    approximate_bytes: region_stat.estimated_disk_size() as i64,
498                    extensions,
499                }
500            })
501            .collect()
502    }
503
504    pub fn close(&self) -> Result<()> {
505        let running = self.running.clone();
506        if running
507            .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire)
508            .is_err()
509        {
510            warn!("Call close heartbeat task multiple times");
511        }
512
513        Ok(())
514    }
515}
516
517#[cfg(test)]
518mod tests {
519    use common_meta::cache_invalidator::DummyCacheInvalidator;
520    use common_meta::kv_backend::memory::MemoryKvBackend;
521    use common_meta::rpc::store::PutRequest;
522    use common_stat::ResourceStatImpl;
523    use meta_client::client::MetaClientBuilder;
524
525    use super::*;
526    use crate::tests::mock_region_server;
527
528    const NODE_ID: u64 = 7;
529
530    async fn kv_backend_with_record(node_epoch: Option<u64>) -> KvBackendRef {
531        let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::new());
532        if let Some(node_epoch) = node_epoch {
533            let key = DatanodeEpochKey::new(NODE_ID).to_bytes();
534            let value = DatanodeEpochValue { node_epoch }
535                .try_as_raw_value()
536                .unwrap();
537            let req = PutRequest::new().with_key(key).with_value(value);
538            kv_backend.put(req).await.unwrap();
539        }
540        kv_backend
541    }
542
543    #[tokio::test]
544    async fn test_node_epoch_above_record() {
545        for (record, expected) in [
546            (None, 1000),
547            (Some(500), 1000),
548            (Some(1000), 1001),
549            (Some(5000), 5001),
550        ] {
551            let kv_backend = kv_backend_with_record(record).await;
552            let node_epoch = node_epoch_above_record(&kv_backend, NODE_ID, 1000).await;
553            assert_eq!(node_epoch.unwrap(), expected);
554        }
555    }
556
557    #[tokio::test]
558    async fn test_start_time_stays_local_above_a_record() {
559        let now = common_time::util::current_time_millis() as u64;
560        let record = now + 60_000;
561        let task = HeartbeatTask::try_new(
562            &DatanodeOptions {
563                node_id: Some(NODE_ID),
564                ..Default::default()
565            },
566            mock_region_server(),
567            Arc::new(MetaClientBuilder::datanode_default_options(NODE_ID).build()),
568            kv_backend_with_record(Some(record)).await,
569            Arc::new(DummyCacheInvalidator),
570            Plugins::default(),
571            Arc::new(ResourceStatImpl::default()),
572        )
573        .await
574        .unwrap();
575        assert_eq!(task.node_epoch, record + 1);
576        assert!(task.start_time_ms >= now && task.start_time_ms < record);
577    }
578}
datanode/heartbeat.rs

datanode/
heartbeat.rs
