Skip to main content

hyperion_framework/heartbeat/
sender.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::time::Instant;
26
27// Package
28use serde::Serialize;
29use tokio::sync::mpsc;
30use tokio::time::{Duration, interval};
31
32// Local
33use crate::containerisation::traits::HyperionHeartbeatMessage;
34use crate::heartbeat::handler::HeartbeatMissedHandler;
35use crate::messages::heartbeat::HeartbeatResponse;
36use crate::utilities::time::{current_epoch_ms, fmt_ms};
37
38pub struct HeartbeatSender<T> {
39    interval_ms: u64,
40    response_timeout_ms: u64,
41    container_name: String,
42    target_senders: HashMap<String, mpsc::Sender<T>>,
43    response_rx: mpsc::Receiver<HeartbeatResponse>,
44    next_id: u64,
45    pending: HashMap<u64, (String, Instant)>,
46    missed_handler: Option<Box<dyn HeartbeatMissedHandler>>,
47}
48
49impl<T> HeartbeatSender<T>
50where
51    T: HyperionHeartbeatMessage + Send + Clone + Serialize + 'static,
52{
53    /// Returns the sender task and a tx handle the container uses to forward responses into it.
54    pub fn new(
55        interval_ms: u64,
56        response_timeout_ms: u64,
57        container_name: String,
58        all_senders: HashMap<String, mpsc::Sender<T>>,
59        targets: &[String],
60        missed_handler: Option<Box<dyn HeartbeatMissedHandler>>,
61    ) -> (Self, mpsc::Sender<HeartbeatResponse>) {
62        let target_senders: HashMap<String, mpsc::Sender<T>> = targets
63            .iter()
64            .filter_map(|name| {
65                if let Some(s) = all_senders.get(name) {
66                    Some((name.clone(), s.clone()))
67                } else {
68                    log::warn!(
69                        "HeartbeatSender: target '{}' not found in client senders — skipping",
70                        name
71                    );
72                    None
73                }
74            })
75            .collect();
76
77        let (response_tx, response_rx) = mpsc::channel(32);
78        (
79            Self {
80                interval_ms,
81                response_timeout_ms,
82                container_name,
83                target_senders,
84                response_rx,
85                next_id: 0,
86                pending: HashMap::new(),
87                missed_handler,
88            },
89            response_tx,
90        )
91    }
92
93    pub async fn run(mut self) {
94        let mut ticker = interval(Duration::from_millis(self.interval_ms));
95        ticker.tick().await; // consume the immediate first tick
96        log::info!(
97            "HeartbeatSender running — targets: {:?}",
98            self.target_senders.keys().collect::<Vec<_>>()
99        );
100
101        loop {
102            tokio::select! {
103                _ = ticker.tick() => {
104                    let now_ms = current_epoch_ms();
105                    for (name, sender) in &self.target_senders {
106                        self.next_id += 1;
107                        let id = self.next_id;
108                        if let Some(msg) = T::make_heartbeat_request(id, self.container_name.clone(), now_ms) {
109                            if sender.try_send(msg).is_err() {
110                                log::warn!("HeartbeatSender: failed to queue request to '{}'", name);
111                            }
112                            self.pending.insert(id, (name.clone(), Instant::now()));
113                        }
114                    }
115                    // Prune requests that have exceeded the response timeout
116                    let timeout = Duration::from_millis(self.response_timeout_ms);
117                    let missed_handler = &self.missed_handler;
118                    self.pending.retain(|id, (target, sent_at)| {
119                        if sent_at.elapsed() > timeout {
120                            log::warn!("HeartbeatSender: no response from '{}' for request #{}", target, id);
121                            if let Some(handler) = missed_handler {
122                                handler.on_missed(target);
123                            }
124                            false
125                        } else {
126                            true
127                        }
128                    });
129                }
130                Some(response) = self.response_rx.recv() => {
131                    if let Some((target, sent_at)) = self.pending.remove(&response.request_id) {
132                        let latency_ms = sent_at.elapsed().as_millis();
133                        if response.component_alive {
134                            log::debug!(
135                                "Heartbeat report ✓ {} | {} latency | activity {} ago | state {}",
136                                target.to_ascii_uppercase(), fmt_ms(latency_ms), fmt_ms(response.ms_since_last_activity as u128), response.container_state_val
137                            );
138                        } else {
139                            log::warn!(
140                                "Heartbeat report ✓ {} | {} latency | COMPONENT NOT ALIVE | activity {} ago",
141                                target.to_ascii_uppercase(), fmt_ms(latency_ms), fmt_ms(response.ms_since_last_activity as u128)
142                            );
143                        }
144                    }
145                }
146            }
147        }
148    }
149}