An attempt at refactoring peer timer

Author: davidsteiner

crates/hotfix/src/session.rs

@@ -44,7 +44,6 @@ pub struct SessionRef<M> {
     sender: mpsc::Sender<SessionEvent<M>>,
 }
 
-const TEST_REQUEST_THRESHOLD: f64 = 1.2;
 type TestRequestId = String;
 
 impl<M: FixMessage> SessionRef<M> {
@@ -119,8 +118,6 @@ struct Session<M, S> {
     application: ApplicationRef<M>,
     store: S,
     awaiting_test_response: Option<TestRequestId>,
-
-    peer_timer: Pin<Box<Sleep>>,
     schedule_check_timer: Pin<Box<Sleep>>,
 }
 
@@ -131,9 +128,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
         application: ApplicationRef<M>,
         store: S,
     ) -> Session<M, S> {
-        let peer_timer = sleep(Duration::from_secs(
-            (config.heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64,
-        ));
         let schedule_check_timer = sleep(Duration::from_secs(SCHEDULE_CHECK_INTERVAL));
 
         let dictionary = Self::get_data_dictionary(&config);
@@ -149,7 +143,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
             application,
             store,
             awaiting_test_response: None,
-            peer_timer: Box::pin(peer_timer),
             schedule_check_timer: Box::pin(schedule_check_timer),
         }
     }
@@ -544,11 +537,8 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
     }
 
     fn reset_peer_timer(&mut self, awaiting_req_id: Option<TestRequestId>) {
-        let seconds =
-            (self.config.heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64;
-        let deadline = Instant::now() + Duration::from_secs(seconds);
-        self.peer_timer.as_mut().reset(deadline);
         self.awaiting_test_response = awaiting_req_id;
+        self.state.reset_peer_timer(self.config.heartbeat_interval);
     }
 
     async fn send_message(&mut self, message: impl FixMessage) {
@@ -616,7 +606,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
     async fn logout_and_terminate(&mut self, reason: &str) {
         self.logout(reason).await;
         self.state.disconnect().await;
-        self.disable_peer_timer().await;
     }
 
     async fn initiate_graceful_logout(&mut self, reason: &str) {
@@ -711,13 +700,6 @@ impl<M: FixMessage, S: MessageStore> Session<M, S> {
         let deadline = Instant::now() + Duration::from_secs(SCHEDULE_CHECK_INTERVAL);
         self.schedule_check_timer.as_mut().reset(deadline);
     }
-
-    async fn disable_peer_timer(&mut self) {
-        // push the timer out by about a month - there is no way to disable it entirely afaik
-        self.peer_timer
-            .as_mut()
-            .reset(Instant::now() + Duration::from_secs(2592000));
-    }
 }
 
 async fn run_session<M, S>(mut session: Session<M, S>)
@@ -727,6 +709,8 @@ where
 {
     loop {
         let next_message = session.mailbox.recv();
+        let heartbeat_timer = session.state.heartbeat_timer();
+        let peer_timer = session.state.peer_timer();
 
         select! {
             next = next_message => {
@@ -738,15 +722,21 @@ where
                 }
             }
             () = async {
-                if let Some(timer) = session.state.heartbeat_timer() {
+                if let Some(timer) = heartbeat_timer {
                     timer.as_mut().await
                 } else {
                     std::future::pending().await
                 }
             } => {
                 session.handle_heartbeat_timeout().await;
             }
-            () = &mut session.peer_timer.as_mut() => {
+            () = async {
+                if let Some(timer) = peer_timer {
+                    timer.as_mut().await
+                } else {
+                    std::future::pending().await
+                }
+            } => {
                 session.handle_peer_timeout().await;
             }
             () = &mut session.schedule_check_timer.as_mut() => {

crates/hotfix/src/session/state.rs

@@ -9,6 +9,8 @@ use tokio::sync::oneshot;
 use tokio::time::{Instant, Sleep, sleep};
 use tracing::{debug, error};
 
+const TEST_REQUEST_THRESHOLD: f64 = 1.2;
+
 pub enum SessionState {
     /// We have established a connection, sent a logon message and await a response.
     AwaitingLogon { writer: WriterRef, logon_sent: bool },
@@ -34,9 +36,14 @@ impl SessionState {
 
     pub fn new_active(writer: WriterRef, heartbeat_interval: u64) -> Self {
         let heartbeat_timer = Box::pin(sleep(Duration::from_secs(heartbeat_interval)));
+
+        let peer_interval = calculate_peer_interval(heartbeat_interval);
+        let peer_timer = Box::pin(sleep(Duration::from_secs(peer_interval)));
+
         Self::Active(ActiveState {
             writer,
             heartbeat_timer,
+            peer_timer,
         })
     }
 
@@ -168,11 +175,32 @@ impl SessionState {
             heartbeat_timer.as_mut().reset(deadline);
         }
     }
+
+    pub fn peer_timer(&mut self) -> Option<&mut Pin<Box<Sleep>>> {
+        match self {
+            Self::Active(ActiveState { peer_timer, .. }) => Some(peer_timer),
+            _ => None,
+        }
+    }
+
+    pub fn reset_peer_timer(&mut self, heartbeat_interval: u64) {
+        if let Self::Active(ActiveState { peer_timer, .. }) = self {
+            let interval = calculate_peer_interval(heartbeat_interval);
+            let deadline = Instant::now() + Duration::from_secs(interval);
+            peer_timer.as_mut().reset(deadline);
+        }
+    }
+}
+
+#[inline]
+fn calculate_peer_interval(heartbeat_interval: u64) -> u64 {
+    (heartbeat_interval as f64 * TEST_REQUEST_THRESHOLD).round() as u64
 }
 
 pub struct ActiveState {
     writer: WriterRef,
     heartbeat_timer: Pin<Box<Sleep>>,
+    peer_timer: Pin<Box<Sleep>>,
 }
 
 /// Session state we're in while processing messages we requested to be resent.