use std::time::{Duration, Instant};
#[derive(Debug)]
pub struct HeartbeatManager {
interval: Duration,
last_sent: Instant,
last_received: Instant,
test_request_pending: Option<String>,
test_request_sent_at: Option<Instant>,
}
impl HeartbeatManager {
#[must_use]
pub fn new(interval: Duration) -> Self {
let now = Instant::now();
Self {
interval,
last_sent: now,
last_received: now,
test_request_pending: None,
test_request_sent_at: None,
}
}
#[inline]
pub fn on_message_sent(&mut self) {
self.last_sent = Instant::now();
}
pub fn on_message_received(&mut self, is_heartbeat: bool, test_req_id: Option<&str>) {
self.last_received = Instant::now();
if is_heartbeat
&& let (Some(pending), Some(received)) = (&self.test_request_pending, test_req_id)
&& pending == received
{
self.test_request_pending = None;
self.test_request_sent_at = None;
}
}
#[must_use]
pub fn should_send_heartbeat(&self) -> bool {
self.last_sent.elapsed() >= self.interval
}
#[must_use]
pub fn should_send_test_request(&self) -> bool {
if self.test_request_pending.is_some() {
return false;
}
let grace = Duration::from_secs(1);
self.last_received.elapsed() >= self.interval + grace
}
#[must_use]
pub fn is_timed_out(&self) -> bool {
if let Some(sent_at) = self.test_request_sent_at {
sent_at.elapsed() >= self.interval
} else {
false
}
}
pub fn on_test_request_sent(&mut self, test_req_id: String) {
self.test_request_pending = Some(test_req_id);
self.test_request_sent_at = Some(Instant::now());
self.last_sent = Instant::now();
}
#[must_use]
pub fn pending_test_request(&self) -> Option<&str> {
self.test_request_pending.as_deref()
}
#[must_use]
pub fn time_since_last_received(&self) -> Duration {
self.last_received.elapsed()
}
#[must_use]
pub fn time_since_last_sent(&self) -> Duration {
self.last_sent.elapsed()
}
#[must_use]
pub const fn interval(&self) -> Duration {
self.interval
}
pub fn reset(&mut self) {
let now = Instant::now();
self.last_sent = now;
self.last_received = now;
self.test_request_pending = None;
self.test_request_sent_at = None;
}
}
#[must_use]
pub fn generate_test_req_id() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
format!("TEST{}", nanos)
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread::sleep;
#[test]
fn test_heartbeat_manager_new() {
let mgr = HeartbeatManager::new(Duration::from_secs(30));
assert_eq!(mgr.interval(), Duration::from_secs(30));
assert!(mgr.pending_test_request().is_none());
}
#[test]
fn test_should_send_heartbeat() {
let mgr = HeartbeatManager::new(Duration::from_millis(10));
assert!(!mgr.should_send_heartbeat());
sleep(Duration::from_millis(15));
assert!(mgr.should_send_heartbeat());
}
#[test]
fn test_on_message_sent() {
let mut mgr = HeartbeatManager::new(Duration::from_millis(10));
sleep(Duration::from_millis(15));
assert!(mgr.should_send_heartbeat());
mgr.on_message_sent();
assert!(!mgr.should_send_heartbeat());
}
#[test]
fn test_test_request_pending() {
let mut mgr = HeartbeatManager::new(Duration::from_secs(30));
mgr.on_test_request_sent("TEST123".to_string());
assert_eq!(mgr.pending_test_request(), Some("TEST123"));
mgr.on_message_received(true, Some("TEST123"));
assert!(mgr.pending_test_request().is_none());
}
#[test]
fn test_generate_test_req_id() {
let id1 = generate_test_req_id();
std::thread::sleep(std::time::Duration::from_nanos(1));
let id2 = generate_test_req_id();
assert!(id1.starts_with("TEST"));
assert!(id2.starts_with("TEST"));
assert!(id1.len() > 4);
assert!(id2.len() > 4);
}
}