use a small timeout?

brayniac · Jul 27, 2023 · a418911 · a418911
1 parent 9698786
commit a418911
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/clients/mod.rs b/src/clients/mod.rs
@@ -2,7 +2,7 @@ use crate::workload::ClientRequest;
 use crate::workload::ClientWorkItem as WorkItem;
 use crate::*;
 use ::momento::MomentoError;
-use crossbeam::channel::{Receiver, TryRecvError};
+use crossbeam::channel::{Receiver, RecvTimeoutError};
 use std::io::{Error, ErrorKind, Result};
 use tokio::io::*;
 use tokio::runtime::Runtime;

diff --git a/src/clients/redis/mod.rs b/src/clients/redis/mod.rs
@@ -76,13 +76,13 @@ async fn task(work_receiver: Receiver<WorkItem>, endpoint: String, config: Confi
         }
 
         let mut con = connection.take().unwrap();
-        let work_item = match work_receiver.try_recv() {
+        let work_item = match work_receiver.recv_timeout(Duration::from_micros(10)) {
             Ok(item) => item,
-            Err(TryRecvError::Empty) => {
+            Err(RecvTimeoutError::Timeout) => {
                 tokio::task::block_in_place(|| work_receiver.recv())
                     .map_err(|_| Error::new(ErrorKind::Other, "channel closed"))?
             }
-            Err(_) => {
+            Err(RecvTimeoutError::Disconnected) => {
                 return Err(Error::new(ErrorKind::Other, "channel closed"));
             }
         };