From a418911247bdbf8fd146167c3fce72b4ae6e4577 Mon Sep 17 00:00:00 2001
From: Brian Martin <brayniac@gmail.com>
Date: Thu, 27 Jul 2023 13:33:41 -0700
Subject: [PATCH] use a small timeout?

---
 src/clients/mod.rs       | 2 +-
 src/clients/redis/mod.rs | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/clients/mod.rs b/src/clients/mod.rs
index de2ab39c..4b8425d4 100644
--- 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
index 9fbb13a2..6e5aed82 100644
--- 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"));
             }
         };