Watchdog forces reader close when peer ignores FIN

Author: davidsteiner

crates/hotfix/src/transport/connection.rs

@@ -1,20 +1,157 @@
+use std::time::Duration;
+
+use tokio::sync::oneshot;
+use tracing::warn;
+
 use crate::transport::reader::ReaderRef;
 use crate::transport::writer::WriterRef;
 
+const FORCE_CLOSE_TIMEOUT: Duration = Duration::from_secs(10);
+
 pub struct FixConnection {
     writer: WriterRef,
     reader: ReaderRef,
+    writer_exit: oneshot::Receiver<()>,
 }
 
 impl FixConnection {
-    pub fn new(writer: WriterRef, reader: ReaderRef) -> Self {
-        Self { writer, reader }
+    pub fn new(writer: WriterRef, reader: ReaderRef, writer_exit: oneshot::Receiver<()>) -> Self {
+        Self {
+            writer,
+            reader,
+            writer_exit,
+        }
     }
+
     pub fn get_writer(&self) -> WriterRef {
         self.writer.clone()
     }
 
     pub async fn run_until_disconnect(self) {
-        self.reader.wait_for_disconnect().await
+        let Self {
+            reader,
+            mut writer_exit,
+            ..
+        } = self;
+        let ReaderRef {
+            mut disconnect_signal,
+            kill,
+        } = reader;
+
+        tokio::select! {
+            _ = &mut disconnect_signal => return,
+            _ = &mut writer_exit => {}
+        }
+
+        match tokio::time::timeout(FORCE_CLOSE_TIMEOUT, &mut disconnect_signal).await {
+            Ok(_) => {}
+            Err(_) => {
+                warn!(
+                    "reader did not observe EOF within {:?}, forcing close",
+                    FORCE_CLOSE_TIMEOUT
+                );
+                let _ = kill.send(());
+                let _ = disconnect_signal.await;
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::transport::writer::WriterMessage;
+    use tokio::sync::mpsc;
+
+    /// Build a `FixConnection` and return the ends the test controls:
+    /// dc_sender to fire from the "reader", writer_exit_tx to fire from the "writer",
+    /// and kill_rx so the test can observe or simulate the reader being killed.
+    fn test_connection() -> (
+        FixConnection,
+        oneshot::Sender<()>,
+        oneshot::Sender<()>,
+        oneshot::Receiver<()>,
+    ) {
+        let (dc_tx, dc_rx) = oneshot::channel::<()>();
+        let (kill_tx, kill_rx) = oneshot::channel::<()>();
+        let reader_ref = ReaderRef::new(dc_rx, kill_tx);
+
+        let (writer_mpsc_tx, _writer_mpsc_rx) = mpsc::channel::<WriterMessage>(1);
+        let writer_ref = WriterRef::new(writer_mpsc_tx);
+
+        let (writer_exit_tx, writer_exit_rx) = oneshot::channel::<()>();
+
+        let conn = FixConnection::new(writer_ref, reader_ref, writer_exit_rx);
+        (conn, dc_tx, writer_exit_tx, kill_rx)
+    }
+
+    /// Reader signals disconnect first — return immediately, kill is never sent.
+    #[tokio::test(start_paused = true)]
+    async fn returns_on_reader_disconnect_before_writer_exit() {
+        let (conn, dc_tx, _writer_exit_tx, mut kill_rx) = test_connection();
+
+        dc_tx.send(()).expect("dc receiver dropped");
+
+        conn.run_until_disconnect().await;
+
+        // Kill should not have been sent. The sender has been dropped by now
+        // (scope ended inside run_until_disconnect), so try_recv returns Closed
+        // rather than Empty. Either way, an Ok(()) would mean kill was sent.
+        assert!(
+            !matches!(kill_rx.try_recv(), Ok(())),
+            "kill signal should not have been sent"
+        );
+    }
+
+    /// Writer exits first, reader disconnects within the watchdog window — no kill.
+    #[tokio::test(start_paused = true)]
+    async fn returns_when_reader_disconnects_after_writer_exit_within_timeout() {
+        let (conn, dc_tx, writer_exit_tx, mut kill_rx) = test_connection();
+
+        writer_exit_tx
+            .send(())
+            .expect("writer_exit receiver dropped");
+
+        // Fire the reader disconnect from a task that runs on the same paused clock.
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_secs(1)).await;
+            let _ = dc_tx.send(());
+        });
+
+        conn.run_until_disconnect().await;
+
+        assert!(
+            !matches!(kill_rx.try_recv(), Ok(())),
+            "kill signal should not have been sent when reader disconnected within timeout"
+        );
+    }
+
+    /// Writer exits first, reader stays blocked past the watchdog — kill fires,
+    /// and a simulated reader fires dc once it sees the kill.
+    #[tokio::test(start_paused = true)]
+    async fn watchdog_fires_kill_when_reader_stuck() {
+        let (conn, dc_tx, writer_exit_tx, kill_rx) = test_connection();
+
+        writer_exit_tx
+            .send(())
+            .expect("writer_exit receiver dropped");
+
+        // Stand in for the reader: when the watchdog kills us, we publish dc.
+        tokio::spawn(async move {
+            if kill_rx.await.is_ok() {
+                let _ = dc_tx.send(());
+            }
+        });
+
+        let start = tokio::time::Instant::now();
+        conn.run_until_disconnect().await;
+        let elapsed = start.elapsed();
+
+        assert!(
+            elapsed >= FORCE_CLOSE_TIMEOUT,
+            "expected watchdog to take at least {:?}, took {:?}",
+            FORCE_CLOSE_TIMEOUT,
+            elapsed
+        );
     }
 }

crates/hotfix/src/transport/reader.rs

@@ -7,9 +7,6 @@ pub struct ReaderMessage;
 
 pub struct ReaderRef {
     pub(crate) disconnect_signal: oneshot::Receiver<()>,
-    // Consumed by the watchdog in FixConnection::run_until_disconnect (follow-up commit).
-    // Kept wired now so the reader's select! already accepts an external kill signal.
-    #[allow(dead_code)]
     pub(crate) kill: oneshot::Sender<()>,
 }
 

crates/hotfix/src/transport/socket.rs

@@ -47,8 +47,8 @@ where
 {
     let (reader, writer) = tokio::io::split(stream);
 
-    let (writer_ref, _writer_exit) = spawn_socket_writer(writer);
+    let (writer_ref, writer_exit) = spawn_socket_writer(writer);
     let reader_ref = spawn_socket_reader(reader, session_ref);
 
-    FixConnection::new(writer_ref, reader_ref)
+    FixConnection::new(writer_ref, reader_ref, writer_exit)
 }

crates/hotfix/src/transport/socket/socket_writer.rs

@@ -314,10 +314,7 @@ mod tests {
             Poll::Ready(Ok(()))
         }
 
-        fn poll_shutdown(
-            self: Pin<&mut Self>,
-            _cx: &mut Context<'_>,
-        ) -> Poll<std::io::Result<()>> {
+        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
             Poll::Pending
         }
     }
@@ -335,10 +332,11 @@ mod tests {
         writer_ref.disconnect().await;
 
         // Advance virtual time past the shutdown timeout.
-        tokio::time::advance(WRITER_SHUTDOWN_TIMEOUT + std::time::Duration::from_millis(100))
-            .await;
+        tokio::time::advance(WRITER_SHUTDOWN_TIMEOUT + std::time::Duration::from_millis(100)).await;
 
         // Exit should have fired by now.
-        exit_rx.await.expect("exit sender dropped without signalling");
+        exit_rx
+            .await
+            .expect("exit sender dropped without signalling");
     }
 }

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

@@ -37,7 +37,8 @@ where
     ) -> Result<Self> {
         let (writer_ref, receiver) = Self::create_writer();
         let (reader_ref, dc_sender) = Self::create_reader();
-        let connection = FixConnection::new(writer_ref, reader_ref);
+        let (_writer_exit_tx, writer_exit_rx) = oneshot::channel();
+        let connection = FixConnection::new(writer_ref, reader_ref, writer_exit_rx);
         let message_config = MessageConfig::default();
         let message_builder = MessageBuilder::new(Dictionary::fix44(), message_config)?;
 
@@ -61,7 +62,8 @@ where
     pub async fn reconnect(&mut self, reset_store: bool) -> Result<()> {
         let (writer_ref, receiver) = Self::create_writer();
         let (reader_ref, dc_sender) = Self::create_reader();
-        let connection = FixConnection::new(writer_ref, reader_ref);
+        let (_writer_exit_tx, writer_exit_rx) = oneshot::channel();
+        let connection = FixConnection::new(writer_ref, reader_ref, writer_exit_rx);
 
         self.receiver = receiver;
         self._dc_sender = dc_sender;