network_protocol/protocol/
keepalive.rs

1use std::time::{Instant, Duration};
2use crate::protocol::heartbeat::{build_ping, is_pong};
3use crate::protocol::message::Message;
4use crate::utils::timeout::KEEPALIVE_INTERVAL;
5use tracing::{debug, warn, instrument};
6
7/// Keep-alive manager to maintain active connections and detect dead peers
8#[derive(Debug)]
9pub struct KeepAliveManager {
10    /// Last time a message was sent
11    last_send: Instant,
12    /// Last time a message was received
13    last_recv: Instant,
14    /// Interval for sending ping messages
15    pub ping_interval: Duration,
16    /// Max time without receiving any message before considering connection dead
17    dead_timeout: Duration,
18}
19
20impl KeepAliveManager {
21    /// Create a new keep-alive manager with default settings
22    pub fn new() -> Self {
23        Self {
24            last_send: Instant::now(),
25            last_recv: Instant::now(),
26            ping_interval: KEEPALIVE_INTERVAL,
27            dead_timeout: KEEPALIVE_INTERVAL.mul_f32(4.0), // 4x the ping interval
28        }
29    }
30
31    /// Create a new keep-alive manager with custom settings
32    pub fn with_settings(ping_interval: Duration, dead_timeout: Duration) -> Self {
33        Self {
34            last_send: Instant::now(),
35            last_recv: Instant::now(),
36            ping_interval,
37            dead_timeout,
38        }
39    }
40
41    /// Update last send time
42    pub fn update_send(&mut self) {
43        self.last_send = Instant::now();
44    }
45
46    /// Update last receive time
47    pub fn update_recv(&mut self) {
48        self.last_recv = Instant::now();
49    }
50
51    /// Check if we need to send a ping to keep the connection alive
52    pub fn should_ping(&self) -> bool {
53        self.last_send.elapsed() >= self.ping_interval
54    }
55    
56    /// Get the ping interval duration
57    pub fn ping_interval(&self) -> Duration {
58        self.ping_interval
59    }
60
61    /// Check if the connection is considered dead (no messages received)
62    pub fn is_connection_dead(&self) -> bool {
63        self.last_recv.elapsed() >= self.dead_timeout
64    }
65
66    /// Get time since last received message
67    pub fn time_since_last_recv(&self) -> Duration {
68        self.last_recv.elapsed()
69    }
70
71    /// Build a ping message for keep-alive
72    #[instrument]
73    pub fn build_ping_message() -> Message {
74        debug!("Building ping message for keep-alive");
75        build_ping()
76    }
77
78    /// Process an incoming message, update received timestamp if it's not a pong
79    /// Returns true if the message was a pong response to our ping
80    #[instrument(skip(msg))]
81    pub fn process_message(&mut self, msg: &Message) -> bool {
82        let is_pong_msg = is_pong(msg);
83        
84        // Update last received time for any message (including pong)
85        self.update_recv();
86        
87        if is_pong_msg {
88            debug!("Received pong message");
89        }
90        
91        is_pong_msg
92    }
93}
94
95impl Default for KeepAliveManager {
96    fn default() -> Self {
97        Self::new()
98    }
99}
network_protocol/protocol/keepalive.rs

network_protocol/protocol/
keepalive.rs
