use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
pub struct HeartbeatManager {
missed_count: Arc<Mutex<u8>>,
max_missed: u8,
interval: Duration,
}
impl Default for HeartbeatManager {
fn default() -> Self {
Self::new()
}
}
impl HeartbeatManager {
pub fn new() -> Self {
Self {
missed_count: Arc::new(Mutex::new(0)),
max_missed: 3,
interval: Duration::from_secs(5),
}
}
pub async fn reset(&self) {
let mut count = self.missed_count.lock().await;
*count = 0;
}
pub async fn increment_missed(&self) -> u8 {
let mut count = self.missed_count.lock().await;
*count += 1;
*count
}
pub async fn is_dead(&self) -> bool {
let count = self.missed_count.lock().await;
*count >= self.max_missed
}
pub fn interval(&self) -> Duration {
self.interval
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_heartbeat_manager() {
let mgr = HeartbeatManager::new();
assert_eq!(mgr.increment_missed().await, 1);
assert_eq!(mgr.increment_missed().await, 2);
assert!(!mgr.is_dead().await);
assert_eq!(mgr.increment_missed().await, 3);
assert!(mgr.is_dead().await);
mgr.reset().await;
assert!(!mgr.is_dead().await);
}
}