use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Notify;
use crate::containerisation::container_state::ContainerState;
pub trait HeartbeatTimeoutHandler: Send + Sync + 'static {
fn on_timeout(&self);
}
pub trait HeartbeatMissedHandler: Send + Sync + 'static {
fn on_missed(&self, target: &str);
}
pub struct FnMissedHandler(Box<dyn Fn(&str) + Send + Sync>);
impl FnMissedHandler {
pub fn new(f: impl Fn(&str) + Send + Sync + 'static) -> Self {
Self(Box::new(f))
}
}
impl HeartbeatMissedHandler for FnMissedHandler {
fn on_missed(&self, target: &str) {
(self.0)(target);
}
}
pub struct FnHandler(Box<dyn Fn() + Send + Sync>);
impl FnHandler {
pub fn new(f: impl Fn() + Send + Sync + 'static) -> Self {
Self(Box::new(f))
}
}
impl HeartbeatTimeoutHandler for FnHandler {
fn on_timeout(&self) {
(self.0)();
}
}
pub struct ShutdownOnTimeout {
container_state: Arc<AtomicUsize>,
container_state_notify: Arc<Notify>,
}
impl ShutdownOnTimeout {
pub fn new(container_state: Arc<AtomicUsize>, container_state_notify: Arc<Notify>) -> Self {
Self {
container_state,
container_state_notify,
}
}
}
impl HeartbeatTimeoutHandler for ShutdownOnTimeout {
fn on_timeout(&self) {
log::warn!(
"HeartbeatReceiver: timeout — no request received within window. Initiating shutdown."
);
self.container_state
.store(ContainerState::ShuttingDown as usize, Ordering::SeqCst);
self.container_state_notify.notify_waiters();
}
}