Move sent test request id into the active 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::{AwaitingResendState, TestRequestId};
 use crate::session_schedule::SessionSchedule;
 use event::SessionEvent;
 use hotfix_message::parsed_message::ParsedMessage;
@@ -44,8 +44,6 @@ pub struct SessionRef<M> {
     sender: mpsc::Sender<SessionEvent<M>>,
 }
 
-type TestRequestId = String;
-
 impl<M: FixMessage> SessionRef<M> {
     pub fn new(
         config: SessionConfig,
@@ -117,7 +115,6 @@ struct Session<M, S> {
     state: SessionState,
     application: ApplicationRef<M>,
     store: S,
-    awaiting_test_response: Option<TestRequestId>,
     schedule_check_timer: Pin<Box<Sleep>>,
 }
 
@@ -142,7 +139,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
             state: SessionState::new_disconnected(true, "initialising"),
             application,
             store,
-            awaiting_test_response: None,
             schedule_check_timer: Box::pin(schedule_check_timer),
         }
     }
@@ -159,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);
@@ -398,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 {
@@ -536,9 +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>) {
-        self.awaiting_test_response = awaiting_req_id;
-        self.state.reset_peer_timer(self.config.heartbeat_interval);
+    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) {
@@ -651,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 {

crates/hotfix/src/session/state.rs

@@ -10,6 +10,8 @@ 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 },
@@ -40,6 +42,7 @@ impl SessionState {
             writer,
             heartbeat_deadline: Instant::now() + Duration::from_secs(heartbeat_interval),
             peer_deadline: Instant::now() + Duration::from_secs(peer_interval),
+            sent_test_request_id: None,
         })
     }
 
@@ -178,12 +181,36 @@ impl SessionState {
         }
     }
 
-    pub fn reset_peer_timer(&mut self, heartbeat_interval: u64) {
-        if let Self::Active(ActiveState { peer_deadline, .. }) = self {
+    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 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]
@@ -192,9 +219,14 @@ fn calculate_peer_interval(heartbeat_interval: u64) -> u64 {
 }
 
 pub struct ActiveState {
+    /// The writer's reference to send messages to the counterparty
     writer: WriterRef,
+    /// 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

@@ -67,8 +67,17 @@ async fn test_peer_timeout() {
     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;
+
     // we wait enough time for the peer deadline to pass
-    tokio::time::advance(std::time::Duration::from_secs(peer_interval)).await;
+    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"))