Avoid hot loop by using timeout on channel receive

Signed-off-by: Sreekanth <[email protected]>
numaproj · Jan 6, 2025 · 1dcad81 · 1dcad81
1 parent d85079a
commit 1dcad81
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/rust/serving/src/source.rs b/rust/serving/src/source.rs
@@ -122,10 +122,10 @@ impl ServingSourceActor {
             if messages.len() >= count || Instant::now() >= timeout_at {
                 break;
             }
-            let message = match self.messages.try_recv() {
-                Ok(msg) => msg,
-                Err(mpsc::error::TryRecvError::Empty) => break,
-                Err(mpsc::error::TryRecvError::Disconnected) => {
+            let next_msg = self.messages.recv();
+            let message = match tokio::time::timeout_at(timeout_at, next_msg).await {
+                Ok(Some(msg)) => msg,
+                Ok(None) => {
                     // If we have collected at-least one message, we return those messages.
                     // The error will happen on all the subsequent read attempts too.
                     if messages.is_empty() {
@@ -136,6 +136,7 @@ impl ServingSourceActor {
                     tracing::error!("Sending half of the Serving channel has disconnected");
                     return Ok(messages);
                 }
+                Err(_) => return Ok(messages),
             };
             let MessageWrapper {
                 confirm_save,