hyperion_framework/containerisation/
hyperion_container.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::fmt::Debug;
25use std::sync::Arc as StdArc;
26use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
27
28// Package
29use serde::{Serialize, de::DeserializeOwned};
30use tokio::sync::{Notify, mpsc};
31use tokio::task;
32use tokio::time::{Duration, sleep};
33
34// Local
35use crate::containerisation::client_broker::ClientBroker;
36use crate::containerisation::container_state::ContainerState;
37use crate::containerisation::traits::{
38    HyperionContainerDirectiveMessage, HyperionHeartbeatMessage, Run,
39};
40use crate::heartbeat::config::{HeartbeatConfig, HeartbeatMode};
41use crate::heartbeat::handler::{HeartbeatMissedHandler, HeartbeatTimeoutHandler};
42use crate::heartbeat::receiver::HeartbeatReceiver;
43use crate::heartbeat::sender::HeartbeatSender;
44use crate::messages::client_broker_message::ClientBrokerMessage;
45use crate::messages::container_directive::ContainerDirective;
46use crate::messages::heartbeat::{HeartbeatRequest, HeartbeatResponse};
47use crate::utilities::time::current_epoch_ms;
48use crate::utilities::tx_sender::add_to_tx_with_retry;
49
50#[allow(dead_code)]
51pub struct HyperionContainer<T> {
52    component_handle: task::JoinHandle<()>,
53    container_state: StdArc<AtomicUsize>,
54    container_state_notify: StdArc<Notify>,
55    client_broker: ClientBroker<T>,
56    component_in_tx: mpsc::Sender<T>,
57    component_out_rx: mpsc::Receiver<ClientBrokerMessage<T>>,
58    main_rx: mpsc::Receiver<T>,
59    server_rx: mpsc::Receiver<T>,
60    last_activity_ms: StdArc<AtomicU64>,
61    heartbeat_request_tx: Option<mpsc::Sender<HeartbeatRequest>>,
62    heartbeat_response_tx: Option<mpsc::Sender<HeartbeatResponse>>,
63}
64
65impl<T> HyperionContainer<T>
66where
67    T: HyperionContainerDirectiveMessage
68        + HyperionHeartbeatMessage
69        + Debug
70        + Send
71        + 'static
72        + DeserializeOwned
73        + Sync
74        + Clone
75        + Serialize,
76{
77    // TODO: Use component_archetype for component restart? Can we store a clean one inside the container without enforcing clone?
78    #[allow(clippy::too_many_arguments)]
79    pub fn create<A>(
80        component_archetype: A,
81        container_state: StdArc<AtomicUsize>,
82        container_state_notify: StdArc<Notify>,
83        client_broker: ClientBroker<T>,
84        main_rx: mpsc::Receiver<T>,
85        server_rx: mpsc::Receiver<T>,
86        heartbeat_config: Option<HeartbeatConfig>,
87        container_name: String,
88        timeout_handler: Option<Box<dyn HeartbeatTimeoutHandler>>,
89        missed_handler: Option<Box<dyn HeartbeatMissedHandler>>,
90    ) -> Self
91    where
92        A: Run<Message = T> + Send + 'static + Sync + Debug,
93    {
94        log::info!("Starting Hyperion Container...");
95        let (component_in_tx, component_in_rx) = mpsc::channel::<T>(32);
96        let (component_out_tx, component_out_rx) = mpsc::channel::<ClientBrokerMessage<T>>(32);
97        let component_handle = HyperionContainer::start_component(
98            component_archetype,
99            component_in_rx,
100            component_out_tx,
101        );
102
103        let last_activity_ms = StdArc::new(AtomicU64::new(current_epoch_ms()));
104        let all_senders = client_broker.clone_all_senders();
105
106        let (heartbeat_request_tx, heartbeat_response_tx) = Self::spawn_heartbeat_tasks(
107            heartbeat_config,
108            container_name,
109            all_senders,
110            timeout_handler,
111            missed_handler,
112        );
113
114        Self {
115            component_handle,
116            container_state,
117            container_state_notify,
118            client_broker,
119            component_in_tx,
120            component_out_rx,
121            main_rx,
122            server_rx,
123            last_activity_ms,
124            heartbeat_request_tx,
125            heartbeat_response_tx,
126        }
127    }
128
129    fn start_component<A>(
130        component_archetype: A,
131        component_in_rx: mpsc::Receiver<T>,
132        component_out_tx: mpsc::Sender<ClientBrokerMessage<T>>,
133    ) -> task::JoinHandle<()>
134    where
135        A: Run<Message = T> + Send + 'static + Sync + Debug,
136    {
137        task::spawn(async move {
138            component_archetype
139                .run(component_in_rx, component_out_tx)
140                .await;
141        })
142    }
143
144    fn spawn_heartbeat_tasks(
145        heartbeat_config: Option<HeartbeatConfig>,
146        container_name: String,
147        all_senders: std::collections::HashMap<String, mpsc::Sender<T>>,
148        timeout_handler: Option<Box<dyn HeartbeatTimeoutHandler>>,
149        missed_handler: Option<Box<dyn HeartbeatMissedHandler>>,
150    ) -> (
151        Option<mpsc::Sender<HeartbeatRequest>>,
152        Option<mpsc::Sender<HeartbeatResponse>>,
153    ) {
154        let config = match heartbeat_config {
155            Some(c) if c.mode != HeartbeatMode::Disabled => c,
156            _ => return (None, None),
157        };
158
159        match config.mode {
160            HeartbeatMode::Sender => {
161                let sender_cfg = match config.sender {
162                    Some(c) => c,
163                    None => {
164                        log::error!("HeartbeatMode::Sender set but no sender config provided");
165                        return (None, None);
166                    }
167                };
168                let (sender_task, response_tx) = HeartbeatSender::new(
169                    sender_cfg.interval_ms,
170                    sender_cfg.response_timeout_ms,
171                    container_name,
172                    all_senders,
173                    &sender_cfg.targets,
174                    missed_handler,
175                );
176                task::spawn(async move { sender_task.run().await });
177                (None, Some(response_tx))
178            }
179            HeartbeatMode::Receiver => {
180                let receiver_cfg = match config.receiver {
181                    Some(c) => c,
182                    None => {
183                        log::error!("HeartbeatMode::Receiver set but no receiver config provided");
184                        return (None, None);
185                    }
186                };
187                let handler = match timeout_handler {
188                    Some(h) => h,
189                    None => {
190                        log::error!("HeartbeatMode::Receiver set but no timeout handler provided");
191                        return (None, None);
192                    }
193                };
194                let (receiver_task, request_tx) =
195                    HeartbeatReceiver::new(receiver_cfg.timeout_ms, all_senders, handler);
196                task::spawn(async move { receiver_task.run().await });
197                (Some(request_tx), None)
198            }
199            HeartbeatMode::Disabled => (None, None),
200        }
201    }
202
203    /// HyperionContainer main loop
204    pub async fn run(&mut self) {
205        log::info!("Hyperion Container is running!");
206        loop {
207            // Check if Container is dying
208            let state = self.container_state.load(Ordering::SeqCst);
209            if state == ContainerState::ShuttingDown as usize
210                || state == ContainerState::DeadComponent as usize
211            {
212                log::info!("Container is shutting down...");
213                self.container_state
214                    .store(ContainerState::ShuttingDown as usize, Ordering::SeqCst);
215                self.container_state_notify.notify_waiters();
216
217                // Allow time for comms to stop etc. before stopping main.rs
218                sleep(Duration::from_secs(3)).await;
219                self.container_state
220                    .store(ContainerState::Closed as usize, Ordering::SeqCst);
221                self.container_state_notify.notify_waiters();
222                break;
223            }
224
225            // Check Component task handle
226            if self.component_handle.is_finished() {
227                log::warn!("Component task has finished unexpectedly.");
228                self.container_state
229                    .store(ContainerState::DeadComponent as usize, Ordering::SeqCst);
230                self.container_state_notify.notify_waiters();
231            }
232
233            tokio::select! {
234                Some(message) = self.main_rx.recv() => {                // Messages from console
235                    log::trace!("Container received message from console: {message:?}");
236                    self.process_incoming_message(message).await;
237                }
238                Some(message) = self.server_rx.recv() => {              // Messages from Server
239                    log::trace!("Container received message from server: {message:?}");
240                    self.process_incoming_message(message).await;
241                }
242                Some(message) = self.component_out_rx.recv() => {       // Messages from Component
243                    log::trace!("Container received message from Component: {message:?}");
244                    self.last_activity_ms.store(current_epoch_ms(), Ordering::SeqCst);
245                    self.client_broker.handle_message(message).await;
246                }
247            }
248        }
249    }
250
251    async fn process_incoming_message(&mut self, message: T) {
252        // Heartbeat messages are handled entirely at container level — never forwarded to component.
253        if let Some(req) = message.as_heartbeat_request() {
254            if let Some(tx) = &self.heartbeat_request_tx {
255                let enriched = HeartbeatRequest {
256                    request_id: req.request_id,
257                    sender_name: req.sender_name,
258                    timestamp_ms: req.timestamp_ms,
259                    component_alive: !self.component_handle.is_finished(),
260                    ms_since_last_activity: current_epoch_ms()
261                        .saturating_sub(self.last_activity_ms.load(Ordering::SeqCst)),
262                    container_state_val: self.container_state.load(Ordering::SeqCst),
263                };
264                let _ = tx.try_send(enriched);
265            }
266            return;
267        }
268        if let Some(resp) = message.as_heartbeat_response() {
269            if let Some(tx) = &self.heartbeat_response_tx {
270                let _ = tx.try_send(resp);
271            }
272            return;
273        }
274
275        // Standard container directives
276        if let Some(container_directive) = message.get_container_directive_message() {
277            match container_directive {
278                ContainerDirective::Shutdown => {
279                    log::info!("Container received shutdown directive");
280                    // Set shutdown state
281                    self.container_state
282                        .store(ContainerState::ShuttingDown as usize, Ordering::SeqCst);
283                    self.container_state_notify.notify_waiters();
284                }
285                ContainerDirective::SystemShutdown => {
286                    log::info!("Container received system shutdown directive");
287                    // Forward shutdown message
288                    self.client_broker.forward_shutdown(message.clone()).await;
289                    // Set shutdown state
290                    self.container_state
291                        .store(ContainerState::ShuttingDown as usize, Ordering::SeqCst);
292                    self.container_state_notify.notify_waiters();
293                    // Wait for clients to finish
294                    self.client_broker.shutdown().await;
295                }
296                _ => {
297                    log::warn!("Container received unmapped directive: {container_directive:?}");
298                }
299            }
300        } else {
301            // Send to component without inspection (ComponentDirective or non-generic message)
302            log::trace!("Forwarding non-framework message to component: {message:?}");
303            add_to_tx_with_retry(
304                &self.component_in_tx,
305                &message,
306                "Container main loop",
307                "Component main loop",
308            )
309            .await;
310        }
311    }
312}
hyperion_framework/containerisation/hyperion_container.rs

hyperion_framework/containerisation/
hyperion_container.rs
