diff --git a/crates/hotfix/src/session.rs b/crates/hotfix/src/session.rs index 7575febe..bac7e33a 100644 --- a/crates/hotfix/src/session.rs +++ b/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 { sender: mpsc::Sender>, } -const TEST_REQUEST_THRESHOLD: f64 = 1.2; -type TestRequestId = String; - impl SessionRef { pub fn new( config: SessionConfig, @@ -118,9 +115,6 @@ struct Session { state: SessionState, application: ApplicationRef, store: S, - awaiting_test_response: Option, - - peer_timer: Pin>, schedule_check_timer: Pin>, } @@ -131,9 +125,6 @@ impl Session { application: ApplicationRef, store: S, ) -> Session { - 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 Session { 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 Session { 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 Session { 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 Session { .reset_heartbeat_timer(self.config.heartbeat_interval); } - fn reset_peer_timer(&mut self, awaiting_req_id: Option) { - 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) { + 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 Session { 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 Session { } 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 Session { 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(mut session: Session) @@ -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() => { diff --git a/crates/hotfix/src/session/state.rs b/crates/hotfix/src/session/state.rs index 276914b6..56561a4c 100644 --- a/crates/hotfix/src/session/state.rs +++ b/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>> { + 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, + ) { + 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>, + /// 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, } /// Session state we're in while processing messages we requested to be resent. diff --git a/crates/hotfix/tests/session_tests.rs b/crates/hotfix/tests/session_tests.rs index 91a05448..36fab8b9 100644 --- a/crates/hotfix/tests/session_tests.rs +++ b/crates/hotfix/tests/session_tests.rs @@ -12,14 +12,32 @@ 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()) @@ -27,9 +45,19 @@ async fn test_happy_login_flow() { 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; }