From c868ff5086e774aea6f5e070acd8f66479fa8214 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Wed, 27 Nov 2024 10:48:46 -0600
Subject: [PATCH] receive_and_buffer exit condition fix

---
 .../transaction_scheduler/receive_and_buffer.rs    | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/core/src/banking_stage/transaction_scheduler/receive_and_buffer.rs b/core/src/banking_stage/transaction_scheduler/receive_and_buffer.rs
index f394b4d41f4ab3..ec0e1b493b830e 100644
--- a/core/src/banking_stage/transaction_scheduler/receive_and_buffer.rs
+++ b/core/src/banking_stage/transaction_scheduler/receive_and_buffer.rs
@@ -52,7 +52,8 @@ pub(crate) trait ReceiveAndBuffer {
     type Transaction: TransactionWithMeta + Send + Sync;
     type Container: StateContainer<Self::Transaction> + Send + Sync;
 
-    /// Returns whether the packet receiver is still connected.
+    /// Returns false only if no packets were received
+    /// AND the receiver is disconnected.
     fn receive_and_buffer_packets(
         &mut self,
         container: &mut Self::Container,
@@ -312,6 +313,7 @@ impl ReceiveAndBuffer for TransactionViewReceiveAndBuffer {
         // Receive packet batches.
         const TIMEOUT: Duration = Duration::from_millis(10);
         let start = Instant::now();
+        let mut received_message = false;
 
         // If not leader, do a blocking-receive initially. This lets the thread
         // sleep when there is not work to do.
@@ -327,6 +329,7 @@ impl ReceiveAndBuffer for TransactionViewReceiveAndBuffer {
         ) {
             match self.receiver.recv_timeout(TIMEOUT) {
                 Ok(packet_batch_message) => {
+                    received_message = true;
                     self.handle_packet_batch_message(
                         container,
                         timing_metrics,
@@ -336,13 +339,16 @@ impl ReceiveAndBuffer for TransactionViewReceiveAndBuffer {
                     );
                 }
                 Err(RecvTimeoutError::Timeout) => return true,
-                Err(RecvTimeoutError::Disconnected) => return false,
+                Err(RecvTimeoutError::Disconnected) => {
+                    return received_message;
+                }
             }
         }
 
         while start.elapsed() < TIMEOUT {
             match self.receiver.try_recv() {
                 Ok(packet_batch_message) => {
+                    received_message = true;
                     self.handle_packet_batch_message(
                         container,
                         timing_metrics,
@@ -352,7 +358,9 @@ impl ReceiveAndBuffer for TransactionViewReceiveAndBuffer {
                     );
                 }
                 Err(TryRecvError::Empty) => return true,
-                Err(TryRecvError::Disconnected) => return false,
+                Err(TryRecvError::Disconnected) => {
+                    return received_message;
+                }
             }
         }