frontend/
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
15#[cfg(test)]
16mod tests;
17
18use std::collections::{HashMap, HashSet};
19use std::fmt::{Display, Formatter};
20use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
21use std::sync::{Arc, Mutex as StdMutex, MutexGuard, PoisonError};
22
23use api::v1::meta::heartbeat_request::NodeWorkloads;
24use api::v1::meta::{FrontendWorkloads, HeartbeatRequest, HeartbeatResponse, NodeInfo, Peer};
25use async_trait::async_trait;
26use common_error::ext::BoxedError;
27use common_meta::cache_invalidator::CacheInvalidatorRef;
28use common_meta::datanode::EnvVars;
29use common_meta::heartbeat::handler::invalidate_table_cache::InvalidateCacheHandler;
30use common_meta::heartbeat::handler::parse_mailbox_message::ParseMailboxMessageHandler;
31use common_meta::heartbeat::handler::suspend::SuspendHandler;
32use common_meta::heartbeat::handler::{
33    HandlerGroupExecutor, HeartbeatResponseHandlerContext, HeartbeatResponseHandlerExecutorRef,
34    HeartbeatResponseHandlerRef,
35};
36use common_meta::heartbeat::mailbox::{HeartbeatMailbox, MailboxRef, OutgoingMessage};
37use common_meta::heartbeat::utils::outgoing_message_to_mailbox_message;
38use common_stat::ResourceStatRef;
39use common_telemetry::{debug, error, info, warn};
40use meta_client::client::heartbeat::HeartbeatConfig;
41use meta_client::client::{HeartbeatSender, HeartbeatStream, MetaClient};
42use snafu::ResultExt;
43use tokio::sync::mpsc::Receiver;
44use tokio::sync::{Mutex, mpsc};
45use tokio::time::{Duration, Instant};
46use tokio_util::sync::CancellationToken;
47
48use crate::error;
49use crate::error::Result;
50use crate::frontend::FrontendOptions;
51use crate::metrics::{HEARTBEAT_RECV_COUNT, HEARTBEAT_SENT_COUNT};
52
53/// The result type returned by a [`FrontendHeartbeatExtension`].
54pub type FrontendHeartbeatExtensionResult<T> = std::result::Result<T, BoxedError>;
55
56/// An extension to frontend heartbeat requests, responses, and lifecycle events.
57///
58/// [`FrontendHeartbeatExtension::request_extensions`] is called for every heartbeat. A failed call
59/// is isolated from the base heartbeat and from other extensions.
60/// [`FrontendHeartbeatExtension::connected`] is called once for every successfully established
61/// connection generation, including reconnects; implementations must therefore be idempotent.
62/// [`FrontendHeartbeatExtension::shutdown`] is called once during task shutdown after heartbeat
63/// I/O has stopped.
64#[async_trait]
65pub trait FrontendHeartbeatExtension: Send + Sync {
66    /// Returns the stable name used to make registration idempotent.
67    fn name(&self) -> &str;
68
69    /// Generates request extensions for one heartbeat.
70    async fn request_extensions(
71        &self,
72    ) -> FrontendHeartbeatExtensionResult<HashMap<String, Vec<u8>>> {
73        Ok(HashMap::new())
74    }
75
76    /// Returns a handler to insert into the heartbeat response handler chain.
77    ///
78    /// Errors and [`common_meta::heartbeat::handler::HandleControl::Done`] are isolated to this
79    /// extension so they cannot skip later extensions or mandatory OSS handlers.
80    fn response_handler(&self) -> Option<HeartbeatResponseHandlerRef> {
81        None
82    }
83
84    /// Notifies the extension that a heartbeat connection generation is ready.
85    async fn connected(&self, _generation: u64) -> FrontendHeartbeatExtensionResult<()> {
86        Ok(())
87    }
88
89    /// Stops and joins background work owned by the extension.
90    async fn shutdown(&self) -> FrontendHeartbeatExtensionResult<()> {
91        Ok(())
92    }
93}
94
95/// A shareable, ordered registry of frontend heartbeat extensions.
96///
97/// Registration is available until [`FrontendHeartbeatExtensions::freeze`] is called. After that,
98/// membership is immutable so every heartbeat consumer observes the same extensions.
99#[derive(Clone, Default)]
100pub struct FrontendHeartbeatExtensions {
101    inner: Arc<StdMutex<FrontendHeartbeatExtensionsInner>>,
102}
103
104/// The reason a frontend heartbeat extension was not registered.
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
106pub enum RegistrationError {
107    /// An extension with the same name is already registered.
108    Duplicate,
109    /// The registry is frozen and no longer accepts registrations.
110    Frozen,
111}
112
113impl Display for RegistrationError {
114    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
115        match self {
116            Self::Duplicate => formatter.write_str("heartbeat extension is already registered"),
117            Self::Frozen => formatter.write_str("heartbeat extension registry is frozen"),
118        }
119    }
120}
121
122impl std::error::Error for RegistrationError {}
123
124#[derive(Default)]
125struct FrontendHeartbeatExtensionsInner {
126    names: HashSet<String>,
127    extensions: Vec<Arc<dyn FrontendHeartbeatExtension>>,
128    frozen: bool,
129}
130
131impl FrontendHeartbeatExtensions {
132    fn lock(&self) -> MutexGuard<'_, FrontendHeartbeatExtensionsInner> {
133        self.inner.lock().unwrap_or_else(PoisonError::into_inner)
134    }
135
136    /// Registers an extension without replacing an existing extension of the same name.
137    ///
138    /// Returns `true` for a new registration and `false` for an idempotent duplicate or when the
139    /// registry is frozen. Use [`FrontendHeartbeatExtensions::try_register`] when the rejection
140    /// reason must not be discarded.
141    pub fn register(&self, extension: Arc<dyn FrontendHeartbeatExtension>) -> bool {
142        self.try_register(extension).is_ok()
143    }
144
145    /// Registers an extension and returns the reason when registration is rejected.
146    ///
147    /// If the registry is already frozen, this method returns [`RegistrationError::Frozen`]
148    /// without calling [`FrontendHeartbeatExtension::name`]. If freezing races with name
149    /// evaluation, the registry state is checked again before insertion.
150    pub fn try_register(
151        &self,
152        extension: Arc<dyn FrontendHeartbeatExtension>,
153    ) -> std::result::Result<(), RegistrationError> {
154        if self.lock().frozen {
155            return Err(RegistrationError::Frozen);
156        }
157
158        let name = extension.name().to_string();
159        let mut inner = self.lock();
160        if inner.frozen {
161            return Err(RegistrationError::Frozen);
162        }
163        if !inner.names.insert(name) {
164            return Err(RegistrationError::Duplicate);
165        }
166        inner.extensions.push(extension);
167        Ok(())
168    }
169
170    /// Prevents further registrations while preserving the current registration order.
171    pub fn freeze(&self) {
172        self.lock().frozen = true;
173    }
174
175    /// Returns the registered extensions in registration order.
176    pub fn extensions(&self) -> Vec<Arc<dyn FrontendHeartbeatExtension>> {
177        self.lock().extensions.clone()
178    }
179
180    /// Returns response handlers in extension registration order.
181    pub fn response_handlers(&self) -> Vec<HeartbeatResponseHandlerRef> {
182        self.extensions()
183            .into_iter()
184            .filter_map(|extension| extension.response_handler())
185            .collect()
186    }
187
188    /// Returns the number of registered extensions.
189    pub fn len(&self) -> usize {
190        self.lock().extensions.len()
191    }
192
193    /// Returns whether no extension is registered.
194    pub fn is_empty(&self) -> bool {
195        self.len() == 0
196    }
197}
198
199/// Builds the frontend heartbeat response handler chain.
200///
201/// Mailbox parsing always runs first, followed by extension handlers, suspension state handling,
202/// and cache invalidation.
203pub fn heartbeat_response_handler_executor(
204    extensions: &FrontendHeartbeatExtensions,
205    suspend_state: Arc<AtomicBool>,
206    cache_invalidator: CacheInvalidatorRef,
207) -> HeartbeatResponseHandlerExecutorRef {
208    let mut handlers: Vec<HeartbeatResponseHandlerRef> = vec![Arc::new(ParseMailboxMessageHandler)];
209    handlers.extend(extensions.response_handlers().into_iter().map(|handler| {
210        Arc::new(IsolatedHeartbeatResponseHandler(handler)) as HeartbeatResponseHandlerRef
211    }));
212    handlers.extend([
213        Arc::new(SuspendHandler::new(suspend_state)) as HeartbeatResponseHandlerRef,
214        Arc::new(InvalidateCacheHandler::new(cache_invalidator)),
215    ]);
216    Arc::new(HandlerGroupExecutor::new(handlers))
217}
218
219struct IsolatedHeartbeatResponseHandler(HeartbeatResponseHandlerRef);
220
221#[async_trait]
222impl common_meta::heartbeat::handler::HeartbeatResponseHandler
223    for IsolatedHeartbeatResponseHandler
224{
225    fn is_acceptable(&self, _ctx: &HeartbeatResponseHandlerContext) -> bool {
226        true
227    }
228
229    async fn handle(
230        &self,
231        ctx: &mut HeartbeatResponseHandlerContext,
232    ) -> common_meta::error::Result<common_meta::heartbeat::handler::HandleControl> {
233        use common_meta::heartbeat::handler::HandleControl;
234
235        if self.0.is_acceptable(ctx)
236            && let Err(error) = self.0.handle(ctx).await
237        {
238            error!(error; "Heartbeat extension response handler failed");
239        }
240        Ok(HandleControl::Continue)
241    }
242}
243
244#[async_trait]
245trait HeartbeatConnector: Send + Sync {
246    async fn connect(&self) -> Result<HeartbeatConnection>;
247}
248
249#[async_trait]
250trait HeartbeatRequestSender: Send + Sync {
251    async fn send(&self, request: HeartbeatRequest) -> FrontendHeartbeatExtensionResult<()>;
252}
253
254#[async_trait]
255trait HeartbeatResponseStream: Send {
256    async fn message(&mut self) -> FrontendHeartbeatExtensionResult<Option<HeartbeatResponse>>;
257}
258
259struct HeartbeatConnection {
260    sender: Arc<dyn HeartbeatRequestSender>,
261    stream: Box<dyn HeartbeatResponseStream>,
262    config: HeartbeatConfig,
263}
264
265struct MetaHeartbeatConnector {
266    client: Arc<MetaClient>,
267}
268
269#[async_trait]
270impl HeartbeatConnector for MetaHeartbeatConnector {
271    async fn connect(&self) -> Result<HeartbeatConnection> {
272        let (sender, stream, config) = self
273            .client
274            .heartbeat()
275            .await
276            .context(error::CreateMetaHeartbeatStreamSnafu)?;
277        Ok(HeartbeatConnection {
278            sender: Arc::new(MetaHeartbeatSender(sender)),
279            stream: Box::new(MetaHeartbeatStream(stream)),
280            config,
281        })
282    }
283}
284
285struct MetaHeartbeatSender(HeartbeatSender);
286
287#[async_trait]
288impl HeartbeatRequestSender for MetaHeartbeatSender {
289    async fn send(&self, request: HeartbeatRequest) -> FrontendHeartbeatExtensionResult<()> {
290        self.0.send(request).await.map_err(BoxedError::new)
291    }
292}
293
294struct MetaHeartbeatStream(HeartbeatStream);
295
296#[async_trait]
297impl HeartbeatResponseStream for MetaHeartbeatStream {
298    async fn message(&mut self) -> FrontendHeartbeatExtensionResult<Option<HeartbeatResponse>> {
299        self.0.message().await.map_err(BoxedError::new)
300    }
301}
302
303#[derive(Clone)]
304struct HeartbeatRunner {
305    peer_addr: String,
306    connector: Arc<dyn HeartbeatConnector>,
307    resp_handler_executor: HeartbeatResponseHandlerExecutorRef,
308    start_time_ms: u64,
309    resource_stat: ResourceStatRef,
310    env_vars: EnvVars,
311    extensions: FrontendHeartbeatExtensions,
312    cancellation: CancellationToken,
313    generation: Arc<AtomicU64>,
314}
315
316impl HeartbeatRunner {
317    async fn connect(&self) -> Result<Option<HeartbeatConnection>> {
318        tokio::select! {
319            _ = self.cancellation.cancelled() => Ok(None),
320            connection = self.connector.connect() => connection.map(Some),
321        }
322    }
323
324    fn next_generation(&self) -> u64 {
325        self.generation
326            .fetch_add(1, Ordering::AcqRel)
327            .wrapping_add(1)
328    }
329
330    async fn notify_connected(&self, generation: u64) -> bool {
331        for extension in self.extensions.extensions() {
332            let result = tokio::select! {
333                _ = self.cancellation.cancelled() => return false,
334                result = extension.connected(generation) => result,
335            };
336            if let Err(error) = result {
337                error!(error; "Heartbeat extension '{}' failed its connected callback", extension.name());
338            }
339        }
340        true
341    }
342
343    async fn shutdown_extensions(&self) {
344        for extension in self.extensions.extensions() {
345            if let Err(error) = extension.shutdown().await {
346                error!(error; "Failed to shut down heartbeat extension '{}'", extension.name());
347            }
348        }
349    }
350
351    async fn run(self, mut connection: HeartbeatConnection) {
352        loop {
353            let retry_interval = connection.config.retry_interval;
354            if self.run_connection(connection).await == ConnectionEnd::Shutdown {
355                return;
356            }
357
358            loop {
359                if !self.wait_retry(retry_interval).await {
360                    return;
361                }
362                info!("Try to re-establish the heartbeat connection to metasrv.");
363
364                match self.connect().await {
365                    Ok(Some(next)) => {
366                        let generation = self.next_generation();
367                        if !self.notify_connected(generation).await {
368                            return;
369                        }
370                        connection = next;
371                        break;
372                    }
373                    Ok(None) => return,
374                    Err(error) => {
375                        error!(error; "Failed to re-establish heartbeat connection to metasrv");
376                    }
377                }
378            }
379        }
380    }
381
382    async fn wait_retry(&self, retry_interval: Duration) -> bool {
383        tokio::select! {
384            _ = self.cancellation.cancelled() => false,
385            _ = tokio::time::sleep(retry_interval) => true,
386        }
387    }
388
389    async fn run_connection(&self, connection: HeartbeatConnection) -> ConnectionEnd {
390        let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
391        let mailbox = Arc::new(HeartbeatMailbox::new(outgoing_tx));
392        let report =
393            self.report_heartbeats(connection.sender, outgoing_rx, connection.config.interval);
394        let responses = self.handle_responses(connection.stream, mailbox);
395        tokio::pin!(report);
396        tokio::pin!(responses);
397
398        tokio::select! {
399            _ = self.cancellation.cancelled() => ConnectionEnd::Shutdown,
400            end = &mut report => end,
401            end = &mut responses => end,
402        }
403    }
404
405    async fn handle_responses(
406        &self,
407        mut stream: Box<dyn HeartbeatResponseStream>,
408        mailbox: MailboxRef,
409    ) -> ConnectionEnd {
410        loop {
411            let response = tokio::select! {
412                _ = self.cancellation.cancelled() => return ConnectionEnd::Shutdown,
413                response = stream.message() => response,
414            };
415            match response {
416                Ok(Some(response)) => {
417                    debug!("Receiving heartbeat response: {:?}", response);
418                    if let Some(message) = &response.mailbox_message {
419                        info!("Received mailbox message: {message:?}");
420                    }
421                    let context = HeartbeatResponseHandlerContext::new(mailbox.clone(), response);
422                    let result = tokio::select! {
423                        _ = self.cancellation.cancelled() => return ConnectionEnd::Shutdown,
424                        result = self.handle_response(context) => result,
425                    };
426                    if let Err(error) = result {
427                        error!(error; "Error while handling heartbeat response");
428                        HEARTBEAT_RECV_COUNT
429                            .with_label_values(&["processing_error"])
430                            .inc();
431                    } else {
432                        HEARTBEAT_RECV_COUNT.with_label_values(&["success"]).inc();
433                    }
434                }
435                Ok(None) => {
436                    warn!("Heartbeat response stream closed");
437                    return ConnectionEnd::Reconnect;
438                }
439                Err(error) => {
440                    HEARTBEAT_RECV_COUNT.with_label_values(&["error"]).inc();
441                    error!(error; "Occur error while reading heartbeat response");
442                    return ConnectionEnd::Reconnect;
443                }
444            }
445        }
446    }
447
448    async fn report_heartbeats(
449        &self,
450        sender: Arc<dyn HeartbeatRequestSender>,
451        mut outgoing_rx: Receiver<OutgoingMessage>,
452        report_interval: Duration,
453    ) -> ConnectionEnd {
454        let total_cpu_millicores = self.resource_stat.get_total_cpu_millicores();
455        let total_memory_bytes = self.resource_stat.get_total_memory_bytes();
456        let mut extensions = HashMap::new();
457        self.env_vars.into_extensions(&mut extensions);
458        let heartbeat_request = HeartbeatRequest {
459            peer: Some(Peer {
460                // Metasrv calculates the frontend id by hashing this reachable address.
461                id: 0,
462                addr: self.peer_addr.clone(),
463            }),
464            info: Self::build_node_info(
465                self.start_time_ms,
466                total_cpu_millicores,
467                total_memory_bytes,
468            ),
469            node_workloads: Some(NodeWorkloads::Frontend(FrontendWorkloads { types: vec![] })),
470            extensions,
471            ..Default::default()
472        };
473        let sleep = tokio::time::sleep(Duration::ZERO);
474        tokio::pin!(sleep);
475
476        loop {
477            let request = tokio::select! {
478                _ = self.cancellation.cancelled() => return ConnectionEnd::Shutdown,
479                message = outgoing_rx.recv() => {
480                    if let Some(message) = message {
481                        Self::new_heartbeat_request(&heartbeat_request, Some(message), 0, 0)
482                    } else {
483                        warn!("Sender has been dropped, exiting the heartbeat loop");
484                        return ConnectionEnd::Reconnect;
485                    }
486                }
487                _ = &mut sleep => {
488                    sleep.as_mut().reset(Instant::now() + report_interval);
489                    Self::new_heartbeat_request(
490                        &heartbeat_request,
491                        None,
492                        self.resource_stat.get_cpu_usage_millicores(),
493                        self.resource_stat.get_memory_usage_bytes(),
494                    )
495                }
496            };
497
498            if let Some(mut request) = request {
499                if !self.add_request_extensions(&mut request).await {
500                    return ConnectionEnd::Shutdown;
501                }
502                debug!(
503                    "Sending a heartbeat request to metasrv, content: {:?}",
504                    request
505                );
506                let result = tokio::select! {
507                    _ = self.cancellation.cancelled() => return ConnectionEnd::Shutdown,
508                    result = sender.send(request) => result,
509                };
510                if let Err(error) = result {
511                    error!(error; "Failed to send heartbeat to metasrv");
512                    return ConnectionEnd::Reconnect;
513                }
514                HEARTBEAT_SENT_COUNT.inc();
515            }
516        }
517    }
518
519    async fn add_request_extensions(&self, request: &mut HeartbeatRequest) -> bool {
520        for extension in self.extensions.extensions() {
521            let generated = tokio::select! {
522                _ = self.cancellation.cancelled() => return false,
523                generated = extension.request_extensions() => generated,
524            };
525            let generated = match generated {
526                Ok(generated) => generated,
527                Err(error) => {
528                    error!(error; "Heartbeat extension '{}' failed to generate request extensions", extension.name());
529                    continue;
530                }
531            };
532
533            if let Some(key) = generated
534                .keys()
535                .find(|key| request.extensions.contains_key(*key))
536            {
537                warn!(
538                    "Heartbeat extension '{}' produced conflicting key '{}'; discarding its output",
539                    extension.name(),
540                    key
541                );
542                continue;
543            }
544            request.extensions.extend(generated);
545        }
546        true
547    }
548
549    fn new_heartbeat_request(
550        heartbeat_request: &HeartbeatRequest,
551        message: Option<OutgoingMessage>,
552        cpu_usage: i64,
553        memory_usage: i64,
554    ) -> Option<HeartbeatRequest> {
555        let mailbox_message = match message.map(outgoing_message_to_mailbox_message) {
556            Some(Ok(message)) => Some(message),
557            Some(Err(error)) => {
558                error!(error; "Failed to encode mailbox messages");
559                return None;
560            }
561            None => None,
562        };
563
564        let mut heartbeat_request = HeartbeatRequest {
565            mailbox_message,
566            ..heartbeat_request.clone()
567        };
568        if let Some(info) = heartbeat_request.info.as_mut() {
569            info.memory_usage_bytes = memory_usage;
570            info.cpu_usage_millicores = cpu_usage;
571        }
572        Some(heartbeat_request)
573    }
574
575    #[allow(deprecated)]
576    fn build_node_info(
577        start_time_ms: u64,
578        total_cpu_millicores: i64,
579        total_memory_bytes: i64,
580    ) -> Option<NodeInfo> {
581        let build_info = common_version::build_info();
582        Some(NodeInfo {
583            version: build_info.version.to_string(),
584            git_commit: build_info.commit_short.to_string(),
585            start_time_ms,
586            total_cpu_millicores,
587            total_memory_bytes,
588            cpu_usage_millicores: 0,
589            memory_usage_bytes: 0,
590            // TODO(zyy17): Remove these deprecated fields when the deprecated fields are removed from the proto.
591            cpus: total_cpu_millicores as u32,
592            memory_bytes: total_memory_bytes as u64,
593            hostname: hostname::get()
594                .unwrap_or_default()
595                .to_string_lossy()
596                .to_string(),
597        })
598    }
599
600    async fn handle_response(&self, context: HeartbeatResponseHandlerContext) -> Result<()> {
601        self.resp_handler_executor
602            .handle(context)
603            .await
604            .context(error::HandleHeartbeatResponseSnafu)
605    }
606}
607
608#[derive(Debug, PartialEq, Eq)]
609enum ConnectionEnd {
610    Reconnect,
611    Shutdown,
612}
613
614/// The frontend task that sends [`HeartbeatRequest`] values to metasrv in the background.
615#[derive(Clone)]
616pub struct HeartbeatTask {
617    runner: HeartbeatRunner,
618    start_lock: Arc<Mutex<()>>,
619    shutdown_lock: Arc<Mutex<()>>,
620    supervisor: Arc<Mutex<Option<common_runtime::JoinHandle<()>>>>,
621}
622
623impl HeartbeatTask {
624    pub fn new(
625        peer_addr: String,
626        opts: &FrontendOptions,
627        meta_client: Arc<MetaClient>,
628        resp_handler_executor: HeartbeatResponseHandlerExecutorRef,
629        resource_stat: ResourceStatRef,
630    ) -> Self {
631        Self::new_with_connector(
632            peer_addr,
633            opts,
634            Arc::new(MetaHeartbeatConnector {
635                client: meta_client,
636            }),
637            resp_handler_executor,
638            resource_stat,
639        )
640    }
641
642    fn new_with_connector(
643        peer_addr: String,
644        opts: &FrontendOptions,
645        connector: Arc<dyn HeartbeatConnector>,
646        resp_handler_executor: HeartbeatResponseHandlerExecutorRef,
647        resource_stat: ResourceStatRef,
648    ) -> Self {
649        Self {
650            runner: HeartbeatRunner {
651                peer_addr,
652                connector,
653                resp_handler_executor,
654                start_time_ms: common_time::util::current_time_millis() as u64,
655                resource_stat,
656                env_vars: EnvVars::from_config(&opts.heartbeat_env_vars),
657                extensions: FrontendHeartbeatExtensions::default(),
658                cancellation: CancellationToken::new(),
659                generation: Arc::new(AtomicU64::new(0)),
660            },
661            start_lock: Arc::new(Mutex::new(())),
662            shutdown_lock: Arc::new(Mutex::new(())),
663            supervisor: Arc::new(Mutex::new(None)),
664        }
665    }
666
667    /// Installs the extensions registered before heartbeat startup.
668    pub fn with_extensions(mut self, extensions: FrontendHeartbeatExtensions) -> Self {
669        self.runner.extensions = extensions;
670        self
671    }
672
673    /// Establishes the initial heartbeat connection and starts its background supervisor.
674    pub async fn start(&self) -> Result<()> {
675        let _start_guard = self.start_lock.lock().await;
676        if self.runner.cancellation.is_cancelled() {
677            return Ok(());
678        }
679
680        let finished = {
681            let mut supervisor = self.supervisor.lock().await;
682            match supervisor.as_ref() {
683                Some(handle) if !handle.is_finished() => return Ok(()),
684                Some(_) => supervisor.take(),
685                None => None,
686            }
687        };
688        if let Some(handle) = finished
689            && let Err(error) = handle.await
690            && !error.is_cancelled()
691        {
692            error!(error; "Heartbeat supervisor join failed");
693        }
694
695        let Some(connection) = self.runner.connect().await? else {
696            return Ok(());
697        };
698        info!(
699            "Heartbeat started with Metasrv config: {}",
700            connection.config
701        );
702
703        let generation = self.runner.next_generation();
704        if !self.runner.notify_connected(generation).await {
705            return Ok(());
706        }
707
708        let runner = self.runner.clone();
709        let handle = common_runtime::spawn_hb(async move {
710            runner.run(connection).await;
711        });
712        *self.supervisor.lock().await = Some(handle);
713        Ok(())
714    }
715
716    /// Cancels and joins heartbeat I/O and all registered extension lifecycles.
717    pub async fn shutdown(&self) {
718        let _shutdown_guard = self.shutdown_lock.lock().await;
719        if self.runner.cancellation.is_cancelled() {
720            return;
721        }
722        self.runner.cancellation.cancel();
723
724        // Wait for a concurrently running handshake or connected callback to observe cancellation.
725        let _start_guard = self.start_lock.lock().await;
726        let handle = self.supervisor.lock().await.take();
727        if let Some(handle) = handle
728            && let Err(error) = handle.await
729            && !error.is_cancelled()
730        {
731            error!(error; "Heartbeat supervisor join failed");
732        }
733        self.runner.shutdown_extensions().await;
734    }
735
736    #[cfg(test)]
737    fn generation(&self) -> u64 {
738        self.runner.generation.load(Ordering::Acquire)
739    }
740
741    #[cfg(test)]
742    async fn has_supervisor(&self) -> bool {
743        self.supervisor.lock().await.is_some()
744    }
745
746    #[cfg(test)]
747    pub(crate) fn is_shutdown(&self) -> bool {
748        self.runner.cancellation.is_cancelled()
749    }
750}
frontend/heartbeat.rs

frontend/
heartbeat.rs
