Make active session state creation nicer by moving it into a dedicated function

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::{ActiveState, AwaitingResendState};
+use crate::session::state::AwaitingResendState;
 use crate::session_schedule::SessionSchedule;
 use event::SessionEvent;
 use hotfix_message::parsed_message::ParsedMessage;
@@ -253,12 +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(ActiveState {
-                    writer: state.writer.clone(),
-                    heartbeat_timer: Box::pin(sleep(Duration::from_secs(
-                        self.config.heartbeat_interval,
-                    ))),
-                });
+                let new_state =
+                    SessionState::new_active(state.writer.clone(), self.config.heartbeat_interval);
                 Some(std::mem::replace(&mut self.state, new_state))
             } else {
                 None
@@ -356,12 +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(ActiveState {
-                        writer: writer.clone(),
-                        heartbeat_timer: Box::pin(sleep(Duration::from_secs(
-                            self.config.heartbeat_interval,
-                        ))),
-                    })
+                    self.state =
+                        SessionState::new_active(writer.clone(), self.config.heartbeat_interval);
                 }
                 Err(err) => match err {
                     MessageVerificationError::SeqNumberTooLow { actual, expected } => {
@@ -746,7 +738,7 @@ where
                 }
             }
             () = async {
-                if let Some(ref mut timer) = session.state.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

@@ -6,7 +6,7 @@ use std::collections::VecDeque;
 use std::pin::Pin;
 use std::time::Duration;
 use tokio::sync::oneshot;
-use tokio::time::{Instant, Sleep};
+use tokio::time::{Instant, Sleep, sleep};
 use tracing::{debug, error};
 
 pub enum SessionState {
@@ -32,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,
@@ -163,8 +171,8 @@ impl SessionState {
 }
 
 pub struct ActiveState {
-    pub(crate) writer: WriterRef,
-    pub(crate) heartbeat_timer: Pin<Box<Sleep>>,
+    writer: WriterRef,
+    heartbeat_timer: Pin<Box<Sleep>>,
 }
 
 /// Session state we're in while processing messages we requested to be resent.