Move logout functions into session state

Author: davidsteiner

crates/hotfix/src/session.rs

@@ -392,7 +392,9 @@ where
         }
 
         if self.state.is_logged_on() {
-            self.send_logout("Logout acknowledged").await?;
+            self.state
+                .send_logout(&mut self.ctx, "Logout acknowledged")
+                .await?;
         }
 
         self.ctx
@@ -649,10 +651,12 @@ where
     }
 
     async fn handle_incorrect_begin_string(&mut self, received_begin_string: String) {
-        self.logout_and_terminate(&format!(
-            "beginString={received_begin_string} is not supported"
-        ))
-        .await;
+        self.state
+            .logout_and_terminate(
+                &mut self.ctx,
+                &format!("beginString={received_begin_string} is not supported"),
+            )
+            .await;
     }
 
     async fn handle_incorrect_comp_id(
@@ -671,7 +675,8 @@ where
             error!("failed to send reject message with invalid comp ID: {err}");
         };
 
-        self.logout_and_terminate("incorrect comp ID received")
+        self.state
+            .logout_and_terminate(&mut self.ctx, "incorrect comp ID received")
             .await;
     }
 
@@ -691,7 +696,9 @@ where
             "we expected {expected} sequence number, but target sent lower ({actual}), terminating..."
         );
         let reason = format!("sequence number too low (actual {actual}, expected {expected})");
-        self.logout_and_terminate(&reason).await;
+        self.state
+            .logout_and_terminate(&mut self.ctx, &reason)
+            .await;
         self.state = SessionState::new_disconnected(false, &reason);
     }
 
@@ -817,11 +824,7 @@ where
         &mut self,
         message: impl OutboundMessage,
     ) -> Result<u64, InternalSendError> {
-        let msg_type = message.message_type().to_string();
-        let prepared = self.ctx.prepare_message(message).await?;
-        self.state.send_message(&msg_type, prepared.raw).await;
-        self.reset_heartbeat_timer();
-        Ok(prepared.seq_num)
+        self.state.send_message(&mut self.ctx, message).await
     }
 
     async fn send_resend_request(
@@ -850,48 +853,6 @@ where
         Ok(())
     }
 
-    async fn send_logout(&mut self, reason: &str) -> Result<(), SessionOperationError> {
-        let logout = Logout::with_reason(reason.to_string());
-        self.send_message(logout)
-            .await
-            .with_send_context("logout")?;
-        Ok(())
-    }
-
-    /// Sends a logout message and immediately disconnects the counterparty.
-    ///
-    /// This should be used sparingly in scenarios where there is a major issue
-    /// requiring operational intervention, such as the sequence number being lower
-    /// than expected, or some other key header field containing an invalid value.
-    ///
-    /// In other scenarios, [`initiate_graceful_logout`] should be preferred.
-    async fn logout_and_terminate(&mut self, reason: &str) {
-        if let Err(err) = self.send_logout(reason).await {
-            warn!("failed to send logout during session termination: {}", err);
-        }
-        self.state.disconnect_writer().await;
-    }
-
-    /// Sends a logout message and puts the session state into an [`AwaitingLogout`] state.
-    ///
-    /// The session waits for a configurable timeout period for the counterparty to
-    /// respond with a `Logout` message. If no response is received within the timeout
-    /// period, it disconnects the counterparty.
-    async fn initiate_graceful_logout(
-        &mut self,
-        reason: &str,
-        reconnect: bool,
-    ) -> Result<(), SessionOperationError> {
-        if self.state.try_transition_to_awaiting_logout(
-            Duration::from_secs(self.ctx.config.logout_timeout),
-            reconnect,
-        ) {
-            self.send_logout(reason).await?;
-        }
-
-        Ok(())
-    }
-
     async fn handle_session_event(&mut self, event: SessionEvent) {
         self.handle_schedule_check().await;
 
@@ -900,7 +861,9 @@ where
                 if let Err(err) = self.on_incoming(fix_message).await {
                     let reason = err.to_string();
                     error!(reason, "fatal error in message processing");
-                    self.logout_and_terminate("internal error").await;
+                    self.state
+                        .logout_and_terminate(&mut self.ctx, "internal error")
+                        .await;
                     self.state = SessionState::new_disconnected(true, &reason);
                 }
             }
@@ -945,7 +908,8 @@ where
             AdminRequest::InitiateGracefulShutdown { reconnect } => {
                 warn!("initiating shutdown on request from admin..");
                 if let Err(err) = self
-                    .initiate_graceful_logout("explicitly requested", reconnect)
+                    .state
+                    .initiate_graceful_logout(&mut self.ctx, "explicitly requested", reconnect)
                     .await
                 {
                     error!(err = ?err, "initiating graceful shutdown");
@@ -973,7 +937,9 @@ where
     async fn handle_peer_timeout(&mut self) {
         if self.state.is_expecting_test_response() {
             warn!("peer didn't respond, terminating..");
-            self.logout_and_terminate("peer timeout").await;
+            self.state
+                .logout_and_terminate(&mut self.ctx, "peer timeout")
+                .await;
         } else if self.state.is_awaiting_logon() {
             warn!("peer didn't respond to our Logon, disconnecting..");
             self.state.disconnect_writer().await;
@@ -1007,7 +973,9 @@ where
                 Ok(SessionPeriodComparison::DifferentPeriod) => {
                     // the message store is for a previous session,
                     // we need to terminate this session, reset the store, and reestablish the session
-                    self.logout_and_terminate("session period changed").await;
+                    self.state
+                        .logout_and_terminate(&mut self.ctx, "session period changed")
+                        .await;
                     if let Err(err) = self.ctx.store.reset().await {
                         error!("error resetting session store: {err:}");
                         self.state =
@@ -1018,7 +986,8 @@ where
                     // the creation_time was recorded outside the session schedule,
                     // treat this similarly to a different period - reset the store
                     warn!("store creation time is outside session schedule, resetting store");
-                    self.logout_and_terminate("creation time outside schedule")
+                    self.state
+                        .logout_and_terminate(&mut self.ctx, "creation time outside schedule")
                         .await;
                     if let Err(err) = self.ctx.store.reset().await {
                         error!("error resetting session store: {err:}");
@@ -1029,13 +998,16 @@ where
                 Err(err) => {
                     // actual schedule calculation error (e.g., DST transition, date overflow)
                     error!("error checking session period: {err:?}");
-                    self.logout_and_terminate("internal error").await;
+                    self.state
+                        .logout_and_terminate(&mut self.ctx, "internal error")
+                        .await;
                 }
             }
         } else if self.state.is_connected() {
             // we are currently outside scheduled session time
             if let Err(err) = self
-                .initiate_graceful_logout("End of session time", true)
+                .state
+                .initiate_graceful_logout(&mut self.ctx, "End of session time", true)
                 .await
             {
                 error!(err = ?err, "failed to initiate graceful logout");

crates/hotfix/src/session/state.rs

@@ -10,16 +10,20 @@ pub(crate) use awaiting_logout::AwaitingLogoutState;
 pub(crate) use awaiting_resend::{AwaitingResendState, AwaitingResendTransitionOutcome};
 pub(crate) use disconnected::DisconnectedState;
 
+use crate::Application;
+use crate::message::OutboundMessage;
 use crate::message::logon::Logon;
 use crate::message::logout::Logout;
-use crate::message::parser::RawFixMessage;
+use crate::session::ctx::SessionCtx;
+use crate::session::error::{InternalSendError, InternalSendResultExt, SessionOperationError};
 use crate::session::event::ScheduleResponse;
 use crate::session::info::Status as SessionInfoStatus;
 use crate::transport::writer::WriterRef;
+use hotfix_store::MessageStore;
 use std::time::Duration;
 use tokio::sync::oneshot;
 use tokio::time::Instant;
-use tracing::{debug, error};
+use tracing::{debug, error, warn};
 
 const TEST_REQUEST_THRESHOLD: f64 = 1.2;
 
@@ -64,41 +68,57 @@ impl SessionState {
         }
     }
 
-    pub async fn send_message(&mut self, message_type: &str, message: RawFixMessage) {
+    pub async fn send_message<A, S>(
+        &mut self,
+        ctx: &mut SessionCtx<A, S>,
+        message: impl OutboundMessage,
+    ) -> Result<u64, InternalSendError>
+    where
+        A: Application,
+        S: MessageStore,
+    {
+        let message_type = message.message_type().to_string();
+        let prepared = ctx.prepare_message(message).await?;
+        let raw = prepared.raw;
+
         match self {
             Self::Active(ActiveState { writer, .. })
             | Self::AwaitingResend(AwaitingResendState { writer, .. }) => {
                 if message_type == Logon::MSG_TYPE {
                     error!("logon message is invalid for active sessions")
                 } else {
-                    writer.send_raw_message(message).await
+                    writer.send_raw_message(raw).await
                 }
             }
             Self::AwaitingLogon(AwaitingLogonState {
                 writer, logon_sent, ..
-            }) => match message_type {
+            }) => match message_type.as_str() {
                 Logon::MSG_TYPE => {
                     if *logon_sent {
                         error!("trying to send logon twice");
                     } else {
-                        writer.send_raw_message(message).await;
+                        writer.send_raw_message(raw).await;
                         *logon_sent = true;
                     }
                 }
                 Logout::MSG_TYPE => {
-                    writer.send_raw_message(message).await;
+                    writer.send_raw_message(raw).await;
                 }
                 _ => error!("invalid outgoing message for AwaitingLogon state"),
             },
             Self::AwaitingLogout(AwaitingLogoutState { writer, .. }) => {
                 // Logout messages are allowed because we first transition into AwaitingLogout
                 // and only then send the logout message
                 if message_type == Logout::MSG_TYPE {
-                    writer.send_raw_message(message).await
+                    writer.send_raw_message(raw).await
                 }
             }
             _ => error!("trying to write without an established connection"),
         }
+
+        self.reset_heartbeat_timer(ctx.config.heartbeat_interval);
+
+        Ok(prepared.seq_num)
     }
 
     pub async fn disconnect_writer(&self) {
@@ -167,6 +187,65 @@ impl SessionState {
         }
     }
 
+    /// Sends a logout message and puts the session state into an [`AwaitingLogout`] state.
+    ///
+    /// The session waits for a configurable timeout period for the counterparty to
+    /// respond with a `Logout` message. If no response is received within the timeout
+    /// period, it disconnects the counterparty.
+    pub async fn initiate_graceful_logout<A, S>(
+        &mut self,
+        ctx: &mut SessionCtx<A, S>,
+        reason: &str,
+        reconnect: bool,
+    ) -> Result<(), SessionOperationError>
+    where
+        A: Application,
+        S: MessageStore,
+    {
+        if self.try_transition_to_awaiting_logout(
+            Duration::from_secs(ctx.config.logout_timeout),
+            reconnect,
+        ) {
+            self.send_logout(ctx, reason).await?;
+        }
+
+        Ok(())
+    }
+
+    /// Sends a logout message and immediately disconnects the counterparty.
+    ///
+    /// This should be used sparingly in scenarios where there is a major issue
+    /// requiring operational intervention, such as the sequence number being lower
+    /// than expected, or some other key header field containing an invalid value.
+    ///
+    /// In other scenarios, [`initiate_graceful_logout`] should be preferred.
+    pub async fn logout_and_terminate<A, S>(&mut self, ctx: &mut SessionCtx<A, S>, reason: &str)
+    where
+        A: Application,
+        S: MessageStore,
+    {
+        if let Err(err) = self.send_logout(ctx, reason).await {
+            warn!("failed to send logout during session termination: {}", err);
+        }
+        self.disconnect_writer().await;
+    }
+
+    pub async fn send_logout<A, S>(
+        &mut self,
+        ctx: &mut SessionCtx<A, S>,
+        reason: &str,
+    ) -> Result<(), SessionOperationError>
+    where
+        A: Application,
+        S: MessageStore,
+    {
+        let logout = Logout::with_reason(reason.to_string());
+        self.send_message(ctx, logout)
+            .await
+            .with_send_context("logout")?;
+        Ok(())
+    }
+
     pub fn register_schedule_awaiter(&mut self, responder: oneshot::Sender<ScheduleResponse>) {
         match self {
             SessionState::Disconnected(state) => {