Move heartbeat timer into session state

Author: davidsteiner

crates/hotfix/src/session.rs

@@ -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::{ActiveState, AwaitingResendState};
 use crate::session_schedule::SessionSchedule;
 use event::SessionEvent;
 use hotfix_message::parsed_message::ParsedMessage;
@@ -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,12 @@ 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 {
+                let new_state = SessionState::Active(ActiveState {
                     writer: state.writer.clone(),
-                };
+                    heartbeat_timer: Box::pin(sleep(Duration::from_secs(
+                        self.config.heartbeat_interval,
+                    ))),
+                });
                 Some(std::mem::replace(&mut self.state, new_state))
             } else {
                 None
@@ -355,9 +356,12 @@ 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 {
+                    self.state = SessionState::Active(ActiveState {
                         writer: writer.clone(),
-                    }
+                        heartbeat_timer: Box::pin(sleep(Duration::from_secs(
+                            self.config.heartbeat_interval,
+                        ))),
+                    })
                 }
                 Err(err) => match err {
                     MessageVerificationError::SeqNumberTooLow { actual, expected } => {
@@ -543,13 +547,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 +746,7 @@ where
                 }
             }
             () = async {
-                if let Some(ref mut timer) = session.heartbeat_timer {
+                if let Some(ref mut 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};
 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.
@@ -38,7 +41,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 +80,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 +90,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 +141,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 {
+    pub(crate) writer: WriterRef,
+    pub(crate) heartbeat_timer: Pin<Box<Sleep>>,
 }
 
 /// Session state we're in while processing messages we requested to be resent.