feat: notify application of state changes (#326)

* Make most state transitions go through apply_transition

* Make check_end_of_resend go through apply_transition

* Make graceful logout go through apply_transition

* Remove logic to send logout message in AwaitingLogout state as it's no longer allowed

* Notify application on state changes

* Fix failing tests

Author: davidsteiner

crates/hotfix/src/application.rs

@@ -1,4 +1,5 @@
 use crate::message::OutboundMessage;
+use crate::session::Status;
 use hotfix_message::message::Message;
 
 #[async_trait::async_trait]
@@ -18,6 +19,11 @@ pub trait Application: Send + Sync + 'static {
     async fn on_logout(&mut self, reason: &str);
     /// Called when the session is logged on.
     async fn on_logon(&mut self);
+    /// Called when the session state changes.
+    ///
+    /// This is invoked after every state transition, providing the previous
+    /// and new status. The default implementation does nothing.
+    async fn on_state_change(&self, from: &Status, to: &Status);
 }
 
 /// Standard FIX Business Reject Reason values (tag 380).

crates/hotfix/src/initiator.rs

@@ -161,6 +161,7 @@ mod tests {
     use crate::message::logon::{Logon, ResetSeqNumConfig};
     use crate::message::logout::Logout;
     use crate::message::parser::Parser;
+    use crate::session::Status;
     use crate::store::in_memory::InMemoryMessageStore;
     use hotfix_message::Part;
     use hotfix_message::message::Message;
@@ -195,6 +196,8 @@ mod tests {
         }
         async fn on_logout(&mut self, _reason: &str) {}
         async fn on_logon(&mut self) {}
+
+        async fn on_state_change(&self, _from: &Status, _to: &Status) {}
     }
 
     /// A minimal FIX counterparty for testing the Initiator over TCP.

crates/hotfix/src/session.rs

@@ -223,7 +223,7 @@ where
             .handle_verification_issue(&mut self.ctx, &message, flags)
             .await?
         {
-            self.apply_transition(result);
+            self.apply_transition(result).await;
             return Ok(());
         }
 
@@ -286,25 +286,28 @@ where
     }
 
     async fn check_end_of_resend(&mut self) -> Result<(), SessionOperationError> {
-        let ended_state = if let SessionState::AwaitingResend(state) = &mut self.state {
+        let backlog = if let SessionState::AwaitingResend(state) = &mut self.state {
             if self.ctx.store.next_target_seq_number() > state.end_seq_number {
+                let inbound_queue = std::mem::take(&mut state.inbound_queue);
                 let new_state = SessionState::new_active(
                     state.writer.clone(),
                     self.ctx.config.heartbeat_interval,
                 );
-                Some(std::mem::replace(&mut self.state, new_state))
+                self.apply_transition(TransitionResult::TransitionTo(new_state))
+                    .await;
+                Some(inbound_queue)
             } else {
                 None
             }
         } else {
             None
         };
 
-        if let Some(SessionState::AwaitingResend(mut state)) = ended_state {
+        if let Some(mut inbound_queue) = backlog {
             // we have reached the end of the resend,
             // process queued messages and resume normal operation
             debug!("resend is done, processing backlog");
-            while let Some(msg) = state.inbound_queue.pop_front() {
+            while let Some(msg) = inbound_queue.pop_front() {
                 let seq_number: u64 = msg.get(MSG_SEQ_NUM).unwrap_or_else(|e| {
                     error!("failed to get seq number: {:?}", e);
                     0
@@ -328,39 +331,48 @@ where
     }
 
     async fn on_connect(&mut self, writer: WriterRef) -> Result<(), SessionOperationError> {
-        self.state = SessionState::AwaitingLogon(AwaitingLogonState {
-            writer,
-            logon_sent: false,
-            logon_timeout: Instant::now() + Duration::from_secs(self.ctx.config.logon_timeout),
-        });
+        self.apply_transition(TransitionResult::TransitionTo(SessionState::AwaitingLogon(
+            AwaitingLogonState {
+                writer,
+                logon_sent: false,
+                logon_timeout: Instant::now() + Duration::from_secs(self.ctx.config.logon_timeout),
+            },
+        )))
+        .await;
         self.reset_peer_timer(None);
         self.send_logon().await?;
 
         Ok(())
     }
 
     async fn on_disconnect(&mut self, reason: String) {
-        match self.state {
+        let transition = match self.state {
             SessionState::Active(_)
             | SessionState::AwaitingLogon(_)
             | SessionState::AwaitingResend(_) => {
                 self.state.disconnect_writer().await;
-                self.state = SessionState::new_disconnected(true, &reason);
+                TransitionResult::TransitionTo(SessionState::new_disconnected(true, &reason))
             }
             SessionState::Disconnected(_) => {
-                warn!("disconnect message was received, but the session is already disconnected")
+                warn!("disconnect message was received, but the session is already disconnected");
+                TransitionResult::Stay
             }
             SessionState::AwaitingLogout(AwaitingLogoutState { reconnect, .. }) => {
-                self.state = SessionState::new_disconnected(reconnect, &reason);
+                TransitionResult::TransitionTo(SessionState::new_disconnected(reconnect, &reason))
             }
-        }
+        };
+        self.apply_transition(transition).await;
     }
 
     async fn on_logon(&mut self) -> Result<(), SessionOperationError> {
         if let SessionState::AwaitingLogon(AwaitingLogonState { writer, .. }) = &self.state {
             let writer = writer.clone();
             // happy logon flow, the session is now active
-            self.state = SessionState::new_active(writer, self.ctx.config.heartbeat_interval);
+            self.apply_transition(TransitionResult::TransitionTo(SessionState::new_active(
+                writer,
+                self.ctx.config.heartbeat_interval,
+            )))
+            .await;
             self.ctx.application.on_logon().await;
             self.ctx.store.increment_target_seq_number().await?;
         } else {
@@ -388,12 +400,18 @@ where
             // if we initiated the logout, preserve the reconnect flag
             SessionState::AwaitingLogout(AwaitingLogoutState { reconnect, .. }) => {
                 self.state.disconnect_writer().await;
-                self.state = SessionState::new_disconnected(reconnect, "logout completed");
+                self.apply_transition(TransitionResult::TransitionTo(
+                    SessionState::new_disconnected(reconnect, "logout completed"),
+                ))
+                .await;
             }
             // otherwise assume it makes sense to try to reconnect
             _ => {
                 self.state.disconnect_writer().await;
-                self.state = SessionState::new_disconnected(true, "peer has logged us out")
+                self.apply_transition(TransitionResult::TransitionTo(
+                    SessionState::new_disconnected(true, "peer has logged us out"),
+                ))
+                .await;
             }
         }
 
@@ -462,9 +480,17 @@ where
         Ok(())
     }
 
-    fn apply_transition(&mut self, result: TransitionResult) {
+    async fn apply_transition(&mut self, result: TransitionResult) {
         if let TransitionResult::TransitionTo(new_state) = result {
+            let old_status = self.state.as_status();
             self.state = new_state;
+            let new_status = self.state.as_status();
+            if old_status != new_status {
+                self.ctx
+                    .application
+                    .on_state_change(&old_status, &new_status)
+                    .await;
+            }
         }
     }
 
@@ -532,7 +558,10 @@ where
                     self.state
                         .logout_and_terminate(&mut self.ctx, "internal error")
                         .await;
-                    self.state = SessionState::new_disconnected(true, &reason);
+                    self.apply_transition(TransitionResult::TransitionTo(
+                        SessionState::new_disconnected(true, &reason),
+                    ))
+                    .await;
                 }
             }
             SessionEvent::Disconnected(reason) => {
@@ -575,12 +604,13 @@ where
         match request {
             AdminRequest::InitiateGracefulShutdown { reconnect } => {
                 warn!("initiating shutdown on request from admin..");
-                if let Err(err) = self
+                match self
                     .state
                     .initiate_graceful_logout(&mut self.ctx, "explicitly requested", reconnect)
                     .await
                 {
-                    error!(err = ?err, "initiating graceful shutdown");
+                    Ok(result) => self.apply_transition(result).await,
+                    Err(err) => error!(err = ?err, "initiating graceful shutdown"),
                 }
             }
             AdminRequest::RequestSessionInfo(responder) => {
@@ -646,8 +676,10 @@ where
                         .await;
                     if let Err(err) = self.ctx.store.reset().await {
                         error!("error resetting session store: {err:}");
-                        self.state =
-                            SessionState::new_disconnected(false, "unexpected error in reset");
+                        self.apply_transition(TransitionResult::TransitionTo(
+                            SessionState::new_disconnected(false, "unexpected error in reset"),
+                        ))
+                        .await;
                     }
                 }
                 Ok(SessionPeriodComparison::OutsideSessionTime { .. }) => {
@@ -659,8 +691,10 @@ where
                         .await;
                     if let Err(err) = self.ctx.store.reset().await {
                         error!("error resetting session store: {err:}");
-                        self.state =
-                            SessionState::new_disconnected(false, "unexpected error in reset");
+                        self.apply_transition(TransitionResult::TransitionTo(
+                            SessionState::new_disconnected(false, "unexpected error in reset"),
+                        ))
+                        .await;
                     }
                 }
                 Err(err) => {
@@ -673,12 +707,13 @@ where
             }
         } else if self.state.is_connected() {
             // we are currently outside scheduled session time
-            if let Err(err) = self
+            match self
                 .state
                 .initiate_graceful_logout(&mut self.ctx, "End of session time", true)
                 .await
             {
-                error!(err = ?err, "failed to initiate graceful logout");
+                Ok(result) => self.apply_transition(result).await,
+                Err(err) => error!(err = ?err, "failed to initiate graceful logout"),
             }
         }
 
@@ -887,6 +922,8 @@ mod tests {
         }
         async fn on_logout(&mut self, _: &str) {}
         async fn on_logon(&mut self) {}
+
+        async fn on_state_change(&self, _from: &Status, _to: &Status) {}
     }
 
     fn create_writer_ref() -> WriterRef {

crates/hotfix/src/session/state.rs

@@ -15,7 +15,7 @@ use crate::message::OutboundMessage;
 use crate::message::logon::Logon;
 use crate::message::logout::Logout;
 use crate::message::verification::VerificationFlags;
-use crate::session::ctx::{SessionCtx, VerificationResult};
+use crate::session::ctx::{SessionCtx, TransitionResult, VerificationResult};
 use crate::session::error::{InternalSendError, InternalSendResultExt, SessionOperationError};
 use crate::session::event::ScheduleResponse;
 use crate::session::info::Status as SessionInfoStatus;
@@ -108,12 +108,8 @@ impl SessionState {
                 }
                 _ => 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(raw).await
-                }
+            Self::AwaitingLogout(_) => {
+                error!("trying to send message while awaiting logout");
             }
             _ => error!("trying to write without an established connection"),
         }
@@ -144,25 +140,24 @@ impl SessionState {
     }
 
     pub fn try_transition_to_awaiting_logout(
-        &mut self,
+        &self,
         logout_timeout: Duration,
         reconnect: bool,
-    ) -> bool {
+    ) -> TransitionResult {
         if matches!(self, SessionState::AwaitingLogout(_)) {
             debug!("already in awaiting logout state");
-            return false;
+            return TransitionResult::Stay;
         }
 
         if let Some(writer) = self.get_writer() {
-            *self = SessionState::AwaitingLogout(AwaitingLogoutState {
+            TransitionResult::TransitionTo(SessionState::AwaitingLogout(AwaitingLogoutState {
                 writer: writer.clone(),
                 logout_timeout: Instant::now() + logout_timeout,
                 reconnect,
-            });
-            true
+            }))
         } else {
             error!("trying to transition to awaiting logout without an established connection");
-            false
+            TransitionResult::Stay
         }
     }
 
@@ -202,19 +197,20 @@ impl SessionState {
         ctx: &mut SessionCtx<A, S>,
         reason: &str,
         reconnect: bool,
-    ) -> Result<(), SessionOperationError>
+    ) -> Result<TransitionResult, SessionOperationError>
     where
         A: Application,
         S: MessageStore,
     {
-        if self.try_transition_to_awaiting_logout(
+        let result = self.try_transition_to_awaiting_logout(
             Duration::from_secs(ctx.config.logout_timeout),
             reconnect,
-        ) {
+        );
+        if matches!(result, TransitionResult::TransitionTo(_)) {
             self.send_logout(ctx, reason).await?;
         }
 
-        Ok(())
+        Ok(result)
     }
 
     /// Sends a logout message and immediately disconnects the counterparty.

crates/hotfix/tests/connection_test_cases/helpers.rs

@@ -10,6 +10,7 @@ use std::sync::{Arc, Once};
 use hotfix::Application;
 use hotfix::application::{InboundDecision, OutboundDecision};
 use hotfix::message::OutboundMessage;
+use hotfix::session::Status;
 use hotfix_message::message::Message;
 use rcgen::{CertificateParams, DnType, IsCa, KeyPair, KeyUsagePurpose, SanType};
 use rustls::ServerConfig;
@@ -355,4 +356,6 @@ impl Application for MinimalApplication {
     async fn on_logout(&mut self, _reason: &str) {}
 
     async fn on_logon(&mut self) {}
+
+    async fn on_state_change(&self, _from: &Status, _to: &Status) {}
 }

crates/hotfix/tests/session_test_cases/common/fakes/fake_application.rs

@@ -1,6 +1,7 @@
 use crate::common::test_messages::TestMessage;
 use hotfix::Application;
 use hotfix::application::{InboundDecision, OutboundDecision};
+use hotfix::session::Status;
 use hotfix_message::message::Message;
 use std::sync::Mutex;
 
@@ -75,4 +76,6 @@ impl Application for FakeApplication {
     async fn on_logout(&mut self, _reason: &str) {}
 
     async fn on_logon(&mut self) {}
+
+    async fn on_state_change(&self, _from: &Status, _to: &Status) {}
 }

examples/load-testing/src/application.rs

@@ -2,6 +2,7 @@ use crate::messages::{ExecutionReport, InboundMsg, OutboundMsg};
 use hotfix::Application;
 use hotfix::Message;
 use hotfix::application::{InboundDecision, OutboundDecision};
+use hotfix::session::Status;
 use tokio::sync::mpsc::UnboundedSender;
 use tracing::info;
 
@@ -51,4 +52,8 @@ impl Application for LoadTestingApplication {
     async fn on_logon(&mut self) {
         info!("we've been logged in");
     }
+
+    async fn on_state_change(&self, from: &Status, to: &Status) {
+        info!("we've changed from {:?} to {:?}", from, to);
+    }
 }