Files
imks/engine/heartbeat.rs
T
zhenyi 821537186e refactor(tests): reformat code and update dependency management
- Reorganized import statements in adapter tests for better readability
- Replaced or_insert_with(Vec::new) with or_default() in test closures
- Updated Cargo.lock with new dependency versions and checksums
- Added TLS features to tonic dependency configuration
- Included sqlx, chrono, and uuid dependencies with specific features
- Added jsonwebtoken and arc-swap as project dependencies
- Reformatted assertion statements to comply with line length limits
- Adjusted base64 import order in engine codec module
- Updated protobuf include statement formatting
2026-06-11 12:11:05 +08:00

78 lines
2.2 KiB
Rust

use std::sync::Arc;
use std::time::Duration;
use crate::engine::packet::Packet;
use crate::engine::session::{SessionState, SessionStore, TransportType};
pub struct HeartbeatManager {
store: SessionStore,
ping_interval: u64,
ping_timeout: u64,
}
impl HeartbeatManager {
pub fn new(store: SessionStore, ping_interval: u64, ping_timeout: u64) -> Self {
Self {
store,
ping_interval,
ping_timeout,
}
}
pub fn start(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
let this = self.clone();
tokio::spawn(async move {
this.run().await;
})
}
async fn run(&self) {
let mut interval = tokio::time::interval(Duration::from_millis(self.ping_interval));
loop {
interval.tick().await;
self.check_sessions().await;
}
}
async fn check_sessions(&self) {
let now = std::time::Instant::now();
let timeout_duration = Duration::from_millis(self.ping_interval + self.ping_timeout);
let mut to_remove = Vec::new();
for entry in self.store.sessions.iter() {
let sid = entry.key().clone();
let session = entry.value().clone();
let (state, last_ping, transport) = {
let s = session.read().await;
(s.state, s.last_ping, s.transport)
};
if state == SessionState::Closed {
to_remove.push(sid);
continue;
}
if now.duration_since(last_ping) > timeout_duration {
tracing::warn!("Session {} timed out", sid);
to_remove.push(sid);
continue;
}
// For polling sessions: buffer a ping packet for the next GET request.
// WS/WT sessions rely on their own dedicated ping tasks; the timeout
// check above already serves as the safety net for all transports.
if state == SessionState::Open && transport == TransportType::Polling {
let mut s = session.write().await;
s.buffer_packet(Packet::ping(""));
}
}
for sid in to_remove {
self.store.remove(&sid);
}
}
}