Avoid peer timeout checks when there is no established connection

crates/hotfix/src/session.rs

@@ -11,7 +11,7 @@ use std::cmp::Ordering;
 use std::pin::Pin;
 use tokio::select;
 use tokio::sync::{mpsc, oneshot};
-use tokio::time::{Duration, Instant, Sleep, sleep};
+use tokio::time::{Duration, Instant, Sleep, sleep, sleep_until};
 use tracing::{debug, error, info, warn};
 
 use crate::application::{ApplicationMessage, ApplicationRef};
@@ -31,7 +31,7 @@ use crate::message::sequence_reset::SequenceReset;
 use crate::message::test_request::TestRequest;
 use crate::message_utils::is_admin;
 use crate::session::event::AwaitingActiveSessionResponse;
-use crate::session::state::AwaitingResendState;
+use crate::session::state::{AwaitingResendState, TestRequestId};
 use crate::session_schedule::SessionSchedule;
 use event::SessionEvent;
 use hotfix_message::parsed_message::ParsedMessage;
@@ -44,9 +44,6 @@ pub struct SessionRef<M> {
     sender: mpsc::Sender<SessionEvent<M>>,
 }
 
-const TEST_REQUEST_THRESHOLD: f64 = 1.2;
-type TestRequestId = String;
-
 impl<M: FixMessage> SessionRef<M> {
     pub fn new(
         config: SessionConfig,
@@ -118,9 +115,6 @@ struct Session<M, S> {
     state: SessionState,
     application: ApplicationRef<M>,
     store: S,
-    awaiting_test_response: Option<TestRequestId>,
-
-    peer_timer: Pin<Box<Sleep>>,
     schedule_check_timer: Pin<Box<Sleep>>,
 }
 
@@ -131,9 +125,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
         application: ApplicationRef<M>,
         store: S,
     ) -> Session<M, S> {
-        let peer_timer = sleep(Duration::from_secs(
-            (config.heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64,
-        ));
         let schedule_check_timer = sleep(Duration::from_secs(SCHEDULE_CHECK_INTERVAL));
 
         let dictionary = Self::get_data_dictionary(&config);
@@ -148,8 +139,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
             state: SessionState::new_disconnected(true, "initialising"),
             application,
             store,
-            awaiting_test_response: None,
-            peer_timer: Box::pin(peer_timer),
             schedule_check_timer: Box::pin(schedule_check_timer),
         }
     }
@@ -166,7 +155,7 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
 
     async fn on_incoming(&mut self, raw_message: RawFixMessage) -> Result<()> {
         debug!("received message: {}", raw_message);
-        if self.awaiting_test_response.is_none() {
+        if !self.state.is_expecting_test_response() {
             // if we are not awaiting a specific test response, any message can reset the timer
             // otherwise, only a heartbeat with the corresponding TestReqID can
             self.reset_peer_timer(None);
@@ -405,7 +394,7 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
 
     async fn on_heartbeat(&mut self, message: &Message) {
         if let (Some(expected_req_id), Ok(message_req_id)) = (
-            &self.awaiting_test_response,
+            &self.state.expected_test_response_id(),
             message.get::<&str>(fix44::TEST_REQ_ID),
         ) {
             if expected_req_id.as_str() == message_req_id {
@@ -543,12 +532,9 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
             .reset_heartbeat_timer(self.config.heartbeat_interval);
     }
 
-    fn reset_peer_timer(&mut self, awaiting_req_id: Option<TestRequestId>) {
-        let seconds =
-            (self.config.heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64;
-        let deadline = Instant::now() + Duration::from_secs(seconds);
-        self.peer_timer.as_mut().reset(deadline);
-        self.awaiting_test_response = awaiting_req_id;
+    fn reset_peer_timer(&mut self, test_request_id: Option<TestRequestId>) {
+        self.state
+            .reset_peer_timer(self.config.heartbeat_interval, test_request_id);
     }
 
     async fn send_message(&mut self, message: impl FixMessage) {
@@ -616,7 +602,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
     async fn logout_and_terminate(&mut self, reason: &str) {
         self.logout(reason).await;
         self.state.disconnect().await;
-        self.disable_peer_timer().await;
     }
 
     async fn initiate_graceful_logout(&mut self, reason: &str) {
@@ -662,7 +647,7 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
     }
 
     async fn handle_peer_timeout(&mut self) {
-        if self.awaiting_test_response.is_some() {
+        if self.state.is_expecting_test_response() {
             warn!("peer didn't respond, terminating..");
             self.logout_and_terminate("peer timeout").await;
         } else {
@@ -711,13 +696,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
         let deadline = Instant::now() + Duration::from_secs(SCHEDULE_CHECK_INTERVAL);
         self.schedule_check_timer.as_mut().reset(deadline);
     }
-
-    async fn disable_peer_timer(&mut self) {
-        // push the timer out by about a month - there is no way to disable it entirely afaik
-        self.peer_timer
-            .as_mut()
-            .reset(Instant::now() + Duration::from_secs(2592000));
-    }
 }
 
 async fn run_session<M, S>(mut session: Session<M, S>)
@@ -738,15 +716,21 @@ where
                 }
             }
             () = async {
-                if let Some(timer) = session.state.heartbeat_timer() {
-                    timer.as_mut().await
+                if let Some(deadline) = session.state.heartbeat_deadline() {
+                    sleep_until(*deadline).await
                 } else {
                     std::future::pending().await
                 }
             } => {
                 session.handle_heartbeat_timeout().await;
             }
-            () = &mut session.peer_timer.as_mut() => {
+            () = async {
+                if let Some(deadline) = session.state.peer_deadline() {
+                    sleep_until(*deadline).await
+                } else {
+                    std::future::pending().await
+                }
+            } => {
                 session.handle_peer_timeout().await;
             }
             () = &mut session.schedule_check_timer.as_mut() => {

crates/hotfix/src/session/state.rs

@@ -3,12 +3,15 @@ use crate::session::event::AwaitingActiveSessionResponse;
 use crate::transport::writer::WriterRef;
 use hotfix_message::message::Message;
 use std::collections::VecDeque;
-use std::pin::Pin;
 use std::time::Duration;
 use tokio::sync::oneshot;
-use tokio::time::{Instant, Sleep, sleep};
+use tokio::time::Instant;
 use tracing::{debug, error};
 
+const TEST_REQUEST_THRESHOLD: f64 = 1.2;
+
+pub(crate) type TestRequestId = String;
+
 pub enum SessionState {
     /// We have established a connection, sent a logon message and await a response.
     AwaitingLogon { writer: WriterRef, logon_sent: bool },
@@ -33,10 +36,13 @@ impl SessionState {
     }
 
     pub fn new_active(writer: WriterRef, heartbeat_interval: u64) -> Self {
-        let heartbeat_timer = Box::pin(sleep(Duration::from_secs(heartbeat_interval)));
+        let peer_interval = calculate_peer_interval(heartbeat_interval);
+
         Self::Active(ActiveState {
             writer,
-            heartbeat_timer,
+            heartbeat_deadline: Instant::now() + Duration::from_secs(heartbeat_interval),
+            peer_deadline: Instant::now() + Duration::from_secs(peer_interval),
+            sent_test_request_id: None,
         })
     }
 
@@ -150,29 +156,77 @@ impl SessionState {
         }
     }
 
-    pub fn heartbeat_timer(&mut self) -> Option<&mut Pin<Box<Sleep>>> {
+    pub fn heartbeat_deadline(&self) -> Option<&Instant> {
         match self {
             Self::Active(ActiveState {
-                heartbeat_timer, ..
-            }) => Some(heartbeat_timer),
+                heartbeat_deadline, ..
+            }) => Some(heartbeat_deadline),
             _ => None,
         }
     }
 
     pub fn reset_heartbeat_timer(&mut self, heartbeat_interval: u64) {
         if let Self::Active(ActiveState {
-            heartbeat_timer, ..
+            heartbeat_deadline, ..
+        }) = self
+        {
+            *heartbeat_deadline = Instant::now() + Duration::from_secs(heartbeat_interval);
+        }
+    }
+
+    pub fn peer_deadline(&self) -> Option<&Instant> {
+        match self {
+            Self::Active(ActiveState { peer_deadline, .. }) => Some(peer_deadline),
+            _ => None,
+        }
+    }
+
+    pub fn reset_peer_timer(
+        &mut self,
+        heartbeat_interval: u64,
+        test_request_id: Option<TestRequestId>,
+    ) {
+        if let Self::Active(ActiveState {
+            peer_deadline,
+            sent_test_request_id,
+            ..
         }) = self
         {
-            let deadline = Instant::now() + Duration::from_secs(heartbeat_interval);
-            heartbeat_timer.as_mut().reset(deadline);
+            let interval = calculate_peer_interval(heartbeat_interval);
+            *peer_deadline = Instant::now() + Duration::from_secs(interval);
+            *sent_test_request_id = test_request_id;
         }
     }
+
+    pub fn expected_test_response_id(&self) -> Option<&TestRequestId> {
+        match self {
+            Self::Active(ActiveState {
+                sent_test_request_id: expected_test_response_id,
+                ..
+            }) => expected_test_response_id.as_ref(),
+            _ => None,
+        }
+    }
+
+    pub fn is_expecting_test_response(&self) -> bool {
+        self.expected_test_response_id().is_some()
+    }
+}
+
+#[inline]
+fn calculate_peer_interval(heartbeat_interval: u64) -> u64 {
+    (heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64
 }
 
 pub struct ActiveState {
+    /// The writer's reference to send messages to the counterparty
     writer: WriterRef,
-    heartbeat_timer: Pin<Box<Sleep>>,
+    /// When we should send the next heartbeat message to the counterparty
+    heartbeat_deadline: Instant,
+    /// When the next message from the counterparty is expected at the latest
+    peer_deadline: Instant,
+    /// The ID of the test request we sent on peer timer expiry
+    sent_test_request_id: Option<TestRequestId>,
 }
 
 /// Session state we're in while processing messages we requested to be resent.

crates/hotfix/tests/session_tests.rs

@@ -12,24 +12,52 @@ mod common;
 
 const HEARTBEAT_INTERVAL: u64 = 30;
 
-#[tokio::test]
-async fn test_happy_login_flow() {
+/// Tests the automatic heartbeat mechanism in an active FIX session:
+/// 1. Establishes a session by exchanging logon messages with the counterparty
+/// 2. Advances time beyond the configured heartbeat interval
+/// 3. Verifies that a heartbeat message (type '0') is automatically sent
+/// 4. Cleanly disconnects the session
+///
+/// This test ensures that the session maintains connectivity by sending
+/// periodic heartbeat messages when no other messages are being exchanged,
+/// as required by the FIX protocol to prevent timeout disconnections.
+#[tokio::test(start_paused = true)]
+async fn test_heartbeats() {
     let (session, mut mock_counterparty) = setup().await;
 
     // assert that a logon message is received (type 'A')
     mock_counterparty
         .assert_next(|msg| assert_eq!(msg.header().get::<&str>(MSG_TYPE).unwrap(), "A"))
         .await;
+    // counterparty responds with a logon to establish a happy session
+    mock_counterparty.send_logon().await;
+    tokio::task::yield_now().await;
+
+    // let's wait enough time for a heartbeat and assert that the heartbeat was sent
+    tokio::time::advance(std::time::Duration::from_secs(HEARTBEAT_INTERVAL + 1)).await;
+    mock_counterparty
+        .assert_next(|msg| assert_eq!(msg.header().get::<&str>(MSG_TYPE).unwrap(), "0"))
+        .await;
 
     session
         .disconnect("Test Session Finished".to_string())
         .await;
     mock_counterparty.assert_disconnected().await;
 }
 
+/// Tests the peer timeout and disconnection mechanism:
+/// 1. Establishes a session by exchanging logon messages
+/// 2. Simulates peer inactivity by advancing time past the peer timeout threshold
+/// 3. Verifies that a TestRequest message is sent to check peer responsiveness
+/// 4. Continues to simulate peer silence and verifies automatic disconnection
+///
+/// This test ensures the session properly handles unresponsive peers by first
+/// attempting to verify connectivity with a TestRequest, then disconnecting
+/// if no response is received within the timeout period.
 #[tokio::test(start_paused = true)]
-async fn test_heartbeats() {
-    let (session, mut mock_counterparty) = setup().await;
+async fn test_peer_timeout() {
+    let (_session, mut mock_counterparty) = setup().await;
+    let peer_interval = (1.2 * HEARTBEAT_INTERVAL as f64) as u64 + 1;
 
     // assert that a logon message is received (type 'A')
     mock_counterparty
@@ -45,9 +73,18 @@ async fn test_heartbeats() {
         .assert_next(|msg| assert_eq!(msg.header().get::<&str>(MSG_TYPE).unwrap(), "0"))
         .await;
 
-    session
-        .disconnect("Test Session Finished".to_string())
+    // we wait enough time for the peer deadline to pass
+    tokio::time::advance(std::time::Duration::from_secs(
+        peer_interval - HEARTBEAT_INTERVAL,
+    ))
+    .await;
+    // a TestRequest (type '1') is sent to the counterparty
+    mock_counterparty
+        .assert_next(|msg| assert_eq!(msg.header().get::<&str>(MSG_TYPE).unwrap(), "1"))
         .await;
+
+    // we wait even longer and the counterparty never responds, so we disconnect from the counterparty
+    tokio::time::advance(std::time::Duration::from_secs(peer_interval)).await;
     mock_counterparty.assert_disconnected().await;
 }