update behaviour to stop abrupt writer close

Author: tharinduamila-insta

shotover/src/connection.rs

@@ -285,8 +285,10 @@ fn spawn_read_write_tasks<
     // 1.   The reader task detects that the destination has closed the connection via reader returning None and terminates,
     // 1.1. connection_closed_tx is sent `ConnectionError::OtherSideClosed`
     // 1.2. in_tx, out_tx (reader's copy), and reader_shutdown_tx are dropped
-    // 2.   The writer task detects that reader_shutdown_tx was dropped via reader_shutdown_rx and terminates promptly,
-    //      closing the TCP write half and preventing CLOSE-WAIT connection leaks.
+    // 2.   The writer task detects that reader_shutdown_tx was dropped via reader_shutdown_rx.
+    // 2.1. The writer task drains any already-buffered messages using try_recv() (non-blocking).
+    // 2.2. Once the buffer is empty, the writer task terminates, closing the TCP write half
+    //      and preventing CLOSE-WAIT connection leaks.
     //      (Note: without reader_shutdown_rx, the writer would wait forever on out_rx.recv() since SinkConnection.out_tx still exists)
     // 3.   Eventually when SinkConnection::recv/try_recv is called, it detects that in_tx is dropped by in_rx returning None
     //      and returns the ConnectionError::OtherSideClosed received from connection_closed_rx.
@@ -305,7 +307,8 @@ fn spawn_read_write_tasks<
 
     // This channel is used to signal the writer task when the reader task terminates.
     // When the reader task terminates (for any reason), it drops reader_shutdown_tx,
-    // which causes reader_shutdown_rx to return an error, signaling the writer to shut down.
+    // which causes reader_shutdown_rx to complete, signaling the writer to begin graceful shutdown.
+    // The writer then drains any buffered messages before terminating.
     // This prevents CLOSE-WAIT connection leaks that would occur if the writer task
     // kept waiting on out_rx.recv() while holding the TCP write half open.
     let (reader_shutdown_tx, reader_shutdown_rx) = oneshot::channel::<()>();
@@ -441,38 +444,58 @@ async fn writer_task<C: CodecBuilder + 'static, W: AsyncWrite + Unpin + Send + '
     request_pending: Arc<RequestPending>,
     mut reader_shutdown_rx: oneshot::Receiver<()>,
 ) -> Result<(), ConnectionError> {
+    let mut reader_has_shutdown = false;
+
     loop {
-        tokio::select! {
-            biased;
-            // Check if reader task has terminated - this prevents CLOSE-WAIT leaks
-            // when the broker closes the connection and shotover doesn't actively poll the connection.
-            _ = &mut reader_shutdown_rx => {
-                // Reader task has terminated, we should shut down too.
-                // This ensures the TCP write half is closed promptly.
-                return Ok(());
-            }
-            result = out_rx.recv() => {
-                if let Some(messages) = result {
+        // When reader has shutdown, only process already-buffered messages (non-blocking).
+        // This prevents CLOSE-WAIT leaks while still allowing pending sends to complete.
+        if reader_has_shutdown {
+            match out_rx.try_recv() {
+                Ok(messages) => {
                     request_pending.add(messages.len() as u64);
-                    match writer.send(messages).await {
-                        Err(CodecWriteError::Encoder(err)) => {
-                            return Err(ConnectionError::MessageEncode(Arc::new(err)));
-                        }
-                        Err(CodecWriteError::Io(err)) => {
-                            if matches!(
-                                err.kind(),
-                                ErrorKind::BrokenPipe | ErrorKind::ConnectionReset
-                            ) {
-                                return Err(ConnectionError::OtherSideClosed);
-                            } else {
-                                return Err(ConnectionError::Io(Arc::new(err)));
+                    // Ignore send errors when draining - connection is closing anyway
+                    let _ = writer.send(messages).await;
+                }
+                Err(_) => {
+                    // No more buffered messages, safe to terminate
+                    //TODO not sure if skipping error handling is correct, at least might need to log em
+                    return Ok(());
+                }
+            }
+        } else {
+            tokio::select! {
+                // Don't use biased - prefer sending messages over checking shutdown
+                // to avoid dropping messages that are ready to send.
+                result = out_rx.recv() => {
+                    if let Some(messages) = result {
+                        request_pending.add(messages.len() as u64);
+                        match writer.send(messages).await {
+                            Err(CodecWriteError::Encoder(err)) => {
+                                return Err(ConnectionError::MessageEncode(Arc::new(err)));
                             }
+                            Err(CodecWriteError::Io(err)) => {
+                                if matches!(
+                                    err.kind(),
+                                    ErrorKind::BrokenPipe | ErrorKind::ConnectionReset
+                                ) {
+                                    return Err(ConnectionError::OtherSideClosed);
+                                } else {
+                                    return Err(ConnectionError::Io(Arc::new(err)));
+                                }
+                            }
+                            Ok(()) => {}
                         }
-                        Ok(()) => {}
+                    } else {
+                        // shotover is no longer sending responses, this task is no longer needed
+                        return Ok(());
                     }
-                } else {
-                    // shotover is no longer sending responses, this task is no longer needed
-                    return Ok(());
+                }
+                // Check if reader task has terminated - this prevents CLOSE-WAIT leaks
+                // when the broker closes the connection and shotover doesn't actively poll the connection.
+                _ = &mut reader_shutdown_rx => {
+                    // Reader task has terminated. Don't terminate immediately -
+                    // first drain any buffered messages to avoid dropping them.
+                    reader_has_shutdown = true;
                 }
             }
         }