hyperion_framework/heartbeat/
receiver.rs

1// -------------------------------------------------------------------------------------------------
2// Hyperion Framework
3// https://github.com/robert-hannah/hyperion-framework
4//
5// A lightweight component-based TCP framework for building service-oriented Rust applications with
6// CLI control, async messaging, and lifecycle management.
7//
8// Copyright 2025 Robert Hannah
9//
10// Licensed under the Apache License, Version 2.0 (the "License");
11// you may not use this file except in compliance with the License.
12// You may obtain a copy of the License at
13//
14//     http://www.apache.org/licenses/LICENSE-2.0
15//
16// Unless required by applicable law or agreed to in writing, software
17// distributed under the License is distributed on an "AS IS" BASIS,
18// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19// See the License for the specific language governing permissions and
20// limitations under the License.
21// -------------------------------------------------------------------------------------------------
22
23// Standard
24use std::collections::HashMap;
25use std::sync::Arc;
26use std::sync::atomic::{AtomicU64, Ordering};
27
28// Package
29use serde::Serialize;
30use tokio::sync::mpsc;
31use tokio::time::{Duration, interval};
32
33// Local
34use crate::containerisation::traits::HyperionHeartbeatMessage;
35use crate::heartbeat::handler::HeartbeatTimeoutHandler;
36use crate::messages::heartbeat::HeartbeatRequest;
37use crate::utilities::time::current_epoch_ms;
38
39pub struct HeartbeatReceiver<T> {
40    timeout_ms: u64,
41    last_received_ms: Arc<AtomicU64>,
42    request_rx: mpsc::Receiver<HeartbeatRequest>,
43    all_senders: HashMap<String, mpsc::Sender<T>>, // Copy of client map from ClientBroker
44    handler: Box<dyn HeartbeatTimeoutHandler>,
45}
46
47impl<T> HeartbeatReceiver<T>
48where
49    T: HyperionHeartbeatMessage + Send + Clone + Serialize + 'static,
50{
51    /// Returns the receiver task and a tx handle the container uses to forward enriched requests
52    /// into it.
53    pub fn new(
54        timeout_ms: u64,
55        all_senders: HashMap<String, mpsc::Sender<T>>,
56        handler: Box<dyn HeartbeatTimeoutHandler>,
57    ) -> (Self, mpsc::Sender<HeartbeatRequest>) {
58        let (request_tx, request_rx) = mpsc::channel(32);
59        (
60            Self {
61                timeout_ms,
62                last_received_ms: Arc::new(AtomicU64::new(current_epoch_ms())),
63                request_rx,
64                all_senders,
65                handler,
66            },
67            request_tx,
68        )
69    }
70
71    pub async fn run(mut self) {
72        // Check three times per timeout window so we catch a miss promptly.
73        let check_interval_ms = (self.timeout_ms / 3).max(1000);
74        let mut watchdog = interval(Duration::from_millis(check_interval_ms));
75        log::info!("HeartbeatReceiver running — timeout: {}ms", self.timeout_ms);
76
77        loop {
78            tokio::select! {
79                Some(request) = self.request_rx.recv() => {
80                    self.last_received_ms.store(current_epoch_ms(), Ordering::SeqCst);
81
82                    if let Some(sender) = self.all_senders.get(&request.sender_name) {
83                        if let Some(msg) = T::make_heartbeat_response(
84                            request.request_id,
85                            current_epoch_ms(),
86                            request.component_alive,
87                            request.ms_since_last_activity,
88                            request.container_state_val,
89                        )
90                        && sender.try_send(msg).is_err() {
91                            log::warn!("HeartbeatReceiver: failed to send response to '{}'", request.sender_name);
92                        }
93                    } else {
94                        log::warn!("HeartbeatReceiver: no sender found for '{}' — cannot respond", request.sender_name);
95                    }
96                }
97                _ = watchdog.tick() => {
98                    let elapsed_ms = current_epoch_ms()
99                        .saturating_sub(self.last_received_ms.load(Ordering::SeqCst));
100                    if elapsed_ms > self.timeout_ms {
101                        log::warn!(
102                            "HeartbeatReceiver: {}ms since last request (timeout: {}ms) — calling handler",
103                            elapsed_ms, self.timeout_ms
104                        );
105                        self.handler.on_timeout();
106                    }
107                }
108            }
109        }
110    }
111}
hyperion_framework/heartbeat/receiver.rs

hyperion_framework/heartbeat/
receiver.rs
