Move heartbeat timer into session state

crates/hotfix/src/session.rs

@@ -120,7 +120,6 @@ struct Session<M, S> {
     store: S,
     awaiting_test_response: Option<TestRequestId>,
 
-    heartbeat_timer: Option<Pin<Box<Sleep>>>,
     peer_timer: Pin<Box<Sleep>>,
     schedule_check_timer: Pin<Box<Sleep>>,
 }
@@ -150,7 +149,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
             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<M: FixMessage, S: MessageStore> Session<M, S> {
     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<M: FixMessage, S: MessageStore> Session<M, S> {
             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<M: FixMessage, S: MessageStore> Session<M, S> {
     }
 
     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<TestRequestId>) {
@@ -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

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<Box<Sleep>>> {
+        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<Box<Sleep>>,
 }
 
 /// Session state we're in while processing messages we requested to be resent.

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::time::advance(std::time::Duration::from_millis(5)).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;