Make the heartbeat timer an option so it can be set to None when it's irrelevant

Author: davidsteiner

crates/hotfix/Cargo.toml

@@ -39,4 +39,5 @@ tracing = { workspace = true }
 uuid = { version = "1.5.0", features = ["v4"] }
 
 [dev-dependencies]
-testcontainers = "^0.24"
+testcontainers = "^0.24"
+tokio = { version = "^1", features = ["test-util"] }

crates/hotfix/src/message.rs

@@ -5,13 +5,13 @@ pub use hotfix_message::fix44;
 pub(crate) use hotfix_message::message::{Config, Message};
 pub use hotfix_message::{Part, RepeatingGroup};
 
-pub(crate) mod heartbeat;
-pub(crate) mod logon;
-pub(crate) mod logout;
-pub(crate) mod parser;
-pub(crate) mod resend_request;
-pub(crate) mod sequence_reset;
-pub(crate) mod test_request;
+pub mod heartbeat;
+pub mod logon;
+pub mod logout;
+pub mod parser;
+pub mod resend_request;
+pub mod sequence_reset;
+pub mod test_request;
 
 pub use parser::RawFixMessage;
 
@@ -23,7 +23,7 @@ pub trait FixMessage: Clone + Send + 'static {
     fn parse(message: &Message) -> Self;
 }
 
-pub(crate) fn generate_message(
+pub fn generate_message(
     sender_comp_id: &str,
     target_comp_id: &str,
     msg_seq_num: usize,

crates/hotfix/src/session.rs

@@ -120,7 +120,7 @@ struct Session<M, S> {
     store: S,
     awaiting_test_response: Option<TestRequestId>,
 
-    heartbeat_timer: Pin<Box<Sleep>>,
+    heartbeat_timer: Option<Pin<Box<Sleep>>>,
     peer_timer: Pin<Box<Sleep>>,
     schedule_check_timer: Pin<Box<Sleep>>,
 }
@@ -132,7 +132,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
         application: ApplicationRef<M>,
         store: S,
     ) -> Session<M, S> {
-        let heartbeat_timer = sleep(Duration::from_secs(config.heartbeat_interval));
         let peer_timer = sleep(Duration::from_secs(
             (config.heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64,
         ));
@@ -151,7 +150,7 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
             application,
             store,
             awaiting_test_response: None,
-            heartbeat_timer: Box::pin(heartbeat_timer),
+            heartbeat_timer: None,
             peer_timer: Box::pin(peer_timer),
             schedule_check_timer: Box::pin(schedule_check_timer),
         }
@@ -545,7 +544,12 @@ 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);
-        self.heartbeat_timer.as_mut().reset(deadline);
+        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));
+        }
     }
 
     fn reset_peer_timer(&mut self, awaiting_req_id: Option<TestRequestId>) {
@@ -742,7 +746,13 @@ where
                     None => break,
                 }
             }
-            () = &mut session.heartbeat_timer.as_mut() => {
+            () = async {
+                if let Some(ref mut timer) = session.heartbeat_timer {
+                    timer.as_mut().await
+                } else {
+                    std::future::pending().await
+                }
+            } => {
                 session.handle_heartbeat_timeout().await;
             }
             () = &mut session.peer_timer.as_mut() => {

crates/hotfix/tests/common/mock_counterparty.rs

@@ -1,4 +1,6 @@
-use hotfix::message::{FixMessage, RawFixMessage};
+use hotfix::config::SessionConfig;
+use hotfix::message::logon::{Logon, ResetSeqNumConfig};
+use hotfix::message::{FixMessage, RawFixMessage, generate_message};
 use hotfix::session::SessionRef;
 use hotfix::transport::FixConnection;
 use hotfix::transport::reader::ReaderRef;
@@ -12,18 +14,23 @@ use tokio::sync::{mpsc, oneshot};
 
 const DEFAULT_TIMEOUT: Duration = Duration::from_millis(10);
 
-pub struct MockCounterparty {
+pub struct MockCounterparty<M> {
     receiver: Receiver<WriterMessage>,
-    // History of received messages from the session
-    messages: Vec<Message>,
+    received_messages: Vec<Message>,
+    sent_messages: Vec<Vec<u8>>,
+    session_ref: SessionRef<M>,
+    session_config: SessionConfig,
     dictionary: Dictionary,
     message_config: MessageConfig,
     _connection: FixConnection,
     _dc_sender: oneshot::Sender<()>,
 }
 
-impl MockCounterparty {
-    pub async fn start(session_ref: SessionRef<impl FixMessage>) -> Self {
+impl<M> MockCounterparty<M>
+where
+    M: FixMessage,
+{
+    pub async fn start(session_ref: SessionRef<M>, session_config: SessionConfig) -> Self {
         let (writer_ref, receiver) = Self::create_writer();
         let (reader_ref, dc_sender) = Self::create_reader();
         let connection = FixConnection::new(writer_ref, reader_ref);
@@ -32,14 +39,39 @@ impl MockCounterparty {
 
         Self {
             receiver,
-            messages: vec![],
+            received_messages: vec![],
+            sent_messages: vec![],
+            session_ref,
+            session_config,
             dictionary: Dictionary::fix44(),
             message_config: MessageConfig::default(),
             _connection: connection,
             _dc_sender: dc_sender,
         }
     }
 
+    pub async fn send_logon(&mut self) {
+        let logon = Logon::new(
+            self.session_config.heartbeat_interval,
+            ResetSeqNumConfig::NoReset(None),
+        );
+        self.send_message(logon).await;
+    }
+
+    pub async fn send_message(&mut self, message: impl FixMessage) {
+        let raw_message = generate_message(
+            &self.session_config.sender_comp_id,
+            &self.session_config.target_comp_id,
+            self.sent_messages.len() + 1,
+            message,
+        )
+        .expect("failed to generate message");
+        self.sent_messages.push(raw_message.clone());
+        self.session_ref
+            .new_fix_message_received(RawFixMessage::new(raw_message))
+            .await;
+    }
+
     /// Waits for and returns the next message received from the session.
     ///
     /// A `None` response indicates we have been disconnected, either through the channel
@@ -51,8 +83,8 @@ impl MockCounterparty {
             .and_then(|writer_message| match writer_message {
                 WriterMessage::SendMessage(raw_message) => {
                     let message = self.parse_message(&raw_message);
-                    self.messages.push(message);
-                    self.messages.last()
+                    self.received_messages.push(message);
+                    self.received_messages.last()
                 }
                 WriterMessage::Disconnect => None,
             })
@@ -73,15 +105,15 @@ impl MockCounterparty {
 
     pub async fn assert_next<F>(&mut self, assertion: F)
     where
-        F: FnOnce(&Message) -> bool,
+        F: FnOnce(&Message),
     {
         self.assert_next_with_timeout(assertion, DEFAULT_TIMEOUT)
             .await;
     }
 
     pub async fn assert_next_with_timeout<F>(&mut self, assertion: F, timeout: Duration)
     where
-        F: FnOnce(&Message) -> bool,
+        F: FnOnce(&Message),
     {
         match tokio::time::timeout(timeout, self.get_next()).await {
             Ok(Some(message)) => {

crates/hotfix/tests/session_tests.rs

@@ -10,14 +10,15 @@ use hotfix_message::fix44::MSG_TYPE;
 
 mod common;
 
+const HEARTBEAT_INTERVAL: u64 = 30;
+
 #[tokio::test]
 async fn test_happy_login_flow() {
-    let session = create_session();
-    let mut mock_counterparty = MockCounterparty::start(session.clone()).await;
+    let (session, mut mock_counterparty) = setup().await;
 
     // assert that a logon message is received (type 'A')
     mock_counterparty
-        .assert_next(|msg| msg.header().get::<&str>(MSG_TYPE).unwrap() == "A")
+        .assert_next(|msg| assert_eq!(msg.header().get::<&str>(MSG_TYPE).unwrap(), "A"))
         .await;
 
     session
@@ -26,12 +27,40 @@ async fn test_happy_login_flow() {
     mock_counterparty.assert_disconnected().await;
 }
 
-fn create_session() -> SessionRef<TestMessage> {
+#[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;
+
+    // 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())
+        .await;
+    mock_counterparty.assert_disconnected().await;
+}
+
+async fn setup() -> (SessionRef<TestMessage>, MockCounterparty<TestMessage>) {
     let config = create_session_config();
+    let counterparty_config = create_counterparty_session_config(config.clone());
+
     let application_ref = ApplicationRef::new(MockApplication {});
     let message_store = InMemoryMessageStore::default();
 
-    SessionRef::new(config, application_ref, message_store)
+    let session = SessionRef::new(config, application_ref, message_store);
+    let mock_counterparty = MockCounterparty::start(session.clone(), counterparty_config).await;
+
+    (session, mock_counterparty)
 }
 
 fn create_session_config() -> SessionConfig {
@@ -43,9 +72,18 @@ fn create_session_config() -> SessionConfig {
         connection_host: "".to_string(),
         connection_port: 0,
         tls_config: None,
-        heartbeat_interval: 30,
+        heartbeat_interval: HEARTBEAT_INTERVAL,
         reconnect_interval: 30,
         reset_on_logon: false,
         schedule: None,
     }
 }
+
+/// Create a session configuration for the counterparty from our configuration.
+fn create_counterparty_session_config(session_config: SessionConfig) -> SessionConfig {
+    SessionConfig {
+        sender_comp_id: session_config.target_comp_id.clone(),
+        target_comp_id: session_config.sender_comp_id.clone(),
+        ..session_config
+    }
+}