diff --git a/crates/hotfix/src/session.rs b/crates/hotfix/src/session.rs index 23164cad..7575febe 100644 --- a/crates/hotfix/src/session.rs +++ b/crates/hotfix/src/session.rs @@ -120,7 +120,6 @@ struct Session { store: S, awaiting_test_response: Option, - heartbeat_timer: Option>>, peer_timer: Pin>, schedule_check_timer: Pin>, } @@ -150,7 +149,6 @@ impl Session { application, store, awaiting_test_response: None, - heartbeat_timer: None, peer_timer: Box::pin(peer_timer), schedule_check_timer: Box::pin(schedule_check_timer), } @@ -255,9 +253,8 @@ impl Session { async fn check_end_of_resend(&mut self) -> Result<()> { let ended_state = if let SessionState::AwaitingResend(state) = &mut self.state { if self.store.next_target_seq_number() > state.end_seq_number { - let new_state = SessionState::Active { - writer: state.writer.clone(), - }; + let new_state = + SessionState::new_active(state.writer.clone(), self.config.heartbeat_interval); Some(std::mem::replace(&mut self.state, new_state)) } else { None @@ -355,9 +352,8 @@ impl Session { match self.verify_message(message).await { Ok(_) => { // happy logon flow, the session is now active - self.state = SessionState::Active { - writer: writer.clone(), - } + self.state = + SessionState::new_active(writer.clone(), self.config.heartbeat_interval); } Err(err) => match err { MessageVerificationError::SeqNumberTooLow { actual, expected } => { @@ -543,13 +539,8 @@ impl Session { } fn reset_heartbeat_timer(&mut self) { - let deadline = Instant::now() + Duration::from_secs(self.config.heartbeat_interval); - if let Some(heartbeat_timer) = &mut self.heartbeat_timer { - heartbeat_timer.as_mut().reset(deadline); - } else { - let timer = sleep(Duration::from_secs(self.config.heartbeat_interval)); - self.heartbeat_timer = Some(Box::pin(timer)); - } + self.state + .reset_heartbeat_timer(self.config.heartbeat_interval); } fn reset_peer_timer(&mut self, awaiting_req_id: Option) { @@ -747,7 +738,7 @@ where } } () = async { - if let Some(ref mut timer) = session.heartbeat_timer { + if let Some(timer) = session.state.heartbeat_timer() { timer.as_mut().await } else { std::future::pending().await diff --git a/crates/hotfix/src/session/state.rs b/crates/hotfix/src/session/state.rs index 80bd6946..276914b6 100644 --- a/crates/hotfix/src/session/state.rs +++ b/crates/hotfix/src/session/state.rs @@ -3,7 +3,10 @@ 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 tracing::{debug, error}; pub enum SessionState { @@ -14,7 +17,7 @@ pub enum SessionState { /// We are in the process of gracefully logging out AwaitingLogout { writer: WriterRef }, // we need the writer so we can disconnect it on successful logout /// The session is active, we have connected and mutually logged on. - Active { writer: WriterRef }, + Active(ActiveState), /// The peer has logged us out. LoggedOut { reconnect: bool }, /// The TCP connection has been dropped. @@ -29,6 +32,14 @@ impl SessionState { Self::Disconnected(DisconnectedState::new(reconnect, reason)) } + pub fn new_active(writer: WriterRef, heartbeat_interval: u64) -> Self { + let heartbeat_timer = Box::pin(sleep(Duration::from_secs(heartbeat_interval))); + Self::Active(ActiveState { + writer, + heartbeat_timer, + }) + } + pub fn should_reconnect(&self) -> bool { match self { SessionState::Disconnected(DisconnectedState { reconnect, .. }) => *reconnect, @@ -38,7 +49,8 @@ impl SessionState { pub async fn send_message(&mut self, message_type: &[u8], message: RawFixMessage) { match self { - Self::Active { writer } | Self::AwaitingResend(AwaitingResendState { writer, .. }) => { + Self::Active(ActiveState { writer, .. }) + | Self::AwaitingResend(AwaitingResendState { writer, .. }) => { if message_type == b"A" { error!("logon message is invalid for active sessions") } else { @@ -76,7 +88,7 @@ impl SessionState { pub async fn disconnect(&self) { match self { - Self::Active { writer } + Self::Active(ActiveState { writer, .. }) | Self::AwaitingLogon { writer, .. } | Self::AwaitingLogout { writer } | Self::AwaitingResend(AwaitingResendState { writer, .. }) => writer.disconnect().await, @@ -86,7 +98,7 @@ impl SessionState { pub fn try_transition_to_awaiting_logout(&mut self) -> bool { match self { - Self::Active { writer } + Self::Active(ActiveState { writer, .. }) | Self::AwaitingLogon { writer, .. } | Self::AwaitingResend(AwaitingResendState { writer, .. }) => { *self = SessionState::AwaitingLogout { @@ -137,6 +149,30 @@ impl SessionState { } } } + + pub fn heartbeat_timer(&mut self) -> Option<&mut Pin>> { + match self { + Self::Active(ActiveState { + heartbeat_timer, .. + }) => Some(heartbeat_timer), + _ => None, + } + } + + pub fn reset_heartbeat_timer(&mut self, heartbeat_interval: u64) { + if let Self::Active(ActiveState { + heartbeat_timer, .. + }) = self + { + let deadline = Instant::now() + Duration::from_secs(heartbeat_interval); + heartbeat_timer.as_mut().reset(deadline); + } + } +} + +pub struct ActiveState { + writer: WriterRef, + heartbeat_timer: Pin>, } /// 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 7244d670..91a05448 100644 --- a/crates/hotfix/tests/session_tests.rs +++ b/crates/hotfix/tests/session_tests.rs @@ -37,6 +37,7 @@ async fn test_heartbeats() { .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;