Merge pull request #279 from doyoubi/BackendTimeout

Support backend timeout

Author: doyoubi

conf/server-proxy.toml

@@ -24,6 +24,8 @@ backend_flush_size = 1024
 # In nanoseconds
 backend_low_flush_interval = 200000
 backend_high_flush_interval = 600000
+# In millisecond
+backend_timeout = 3000
 
 # Active Redirection Mode
 # When active_redirection is enabled,

src/bin/server_proxy.rs

@@ -82,6 +82,8 @@ fn gen_conf() -> Result<(ServerProxyConfig, ClusterConfig), &'static str> {
             .unwrap_or(600_000),
     )
     .ok_or("backend_high_flush_interval")?;
+    let backend_timeout = NonZeroU64::new(s.get::<u64>("backend_timeout").unwrap_or(3_000))
+        .ok_or("backend_timeout")?;
 
     let mut max_redirections = s.get::<usize>("max_redirections").unwrap_or(0);
     if max_redirections != 0 {
@@ -117,6 +119,7 @@ fn gen_conf() -> Result<(ServerProxyConfig, ClusterConfig), &'static str> {
         backend_flush_size,
         backend_low_flush_interval: Duration::from_nanos(backend_low_flush_interval.get()),
         backend_high_flush_interval: Duration::from_nanos(backend_high_flush_interval.get()),
+        backend_timeout: Duration::from_millis(backend_timeout.get()),
     };
 
     let mut cluster_config = ClusterConfig::default();

src/proxy/backend.rs

@@ -433,6 +433,7 @@ where
             config.backend_flush_size,
             config.backend_low_flush_interval,
             config.backend_high_flush_interval,
+            config.backend_timeout,
         )
         .await;
         match res {
@@ -463,6 +464,7 @@ async fn handle_conn<H, S>(
     backend_flush_size: NonZeroUsize,
     backend_low_flush_interval: Duration,
     backend_high_flush_interval: Duration,
+    backend_timeout: Duration,
 ) -> Result<(), HandleConnErr<H::Task>>
 where
     H: CmdTaskResultHandler,
@@ -482,6 +484,10 @@ where
         batch_stats.clone(),
     );
 
+    let mut timer = futures_timer::Delay::new(backend_timeout);
+    let mut response_received = false;
+    let mut task_empty = true;
+
     future::poll_fn(
         |cx: &mut Context<'_>| -> Poll<Result<(), HandleConnErr<H::Task>>> {
             let retry_times_opt = match retry_state_opt.take() {
@@ -580,6 +586,8 @@ where
                     Poll::Pending => break Ok(()),
                 };
 
+                response_received = true;
+
                 let mut task = match tasks.pop_front() {
                     Some(task) => task,
                     None => {
@@ -598,6 +606,21 @@ where
                 return Poll::Ready(Err((err, retry_state)));
             }
 
+            if let Poll::Ready(()) = Pin::new(&mut timer).poll(cx) {
+                if !task_empty && !response_received {
+                    let err = BackendError::Timeout;
+                    error!("backend read timeout");
+                    let failed_tasks = tasks.drain(..).collect();
+                    // For timeout we just don't retry as it will take a long time.
+                    let retry_state = handle_conn_err(Some(MAX_BACKEND_RETRY), failed_tasks, &err);
+                    return Poll::Ready(Err((err, retry_state)));
+                }
+
+                timer.reset(backend_timeout);
+                task_empty = tasks.is_empty();
+                response_received = false;
+            }
+
             Poll::Pending
         },
     )
@@ -639,6 +662,7 @@ pub enum BackendError {
     InvalidAddress,
     Canceled,
     InvalidState,
+    Timeout,
 }
 
 impl BackendError {
@@ -650,6 +674,7 @@ impl BackendError {
             SenderBackendError::InvalidAddress => Either::Left(BackendError::InvalidAddress),
             SenderBackendError::Canceled => Either::Left(BackendError::Canceled),
             SenderBackendError::InvalidState => Either::Left(BackendError::InvalidState),
+            SenderBackendError::Timeout => Either::Left(BackendError::Timeout),
             SenderBackendError::Retry(task) => Either::Right(RetryError::new(task)),
         }
     }
@@ -682,6 +707,7 @@ pub enum SenderBackendError<T> {
     Canceled,
     InvalidState,
     Retry(T),
+    Timeout,
 }
 
 impl<T> SenderBackendError<T> {
@@ -693,6 +719,7 @@ impl<T> SenderBackendError<T> {
             BackendError::InvalidAddress => SenderBackendError::InvalidAddress,
             BackendError::Canceled => SenderBackendError::Canceled,
             BackendError::InvalidState => SenderBackendError::InvalidState,
+            BackendError::Timeout => SenderBackendError::Timeout,
         }
     }
 
@@ -709,6 +736,7 @@ impl<T> SenderBackendError<T> {
             Self::Canceled => SenderBackendError::Canceled,
             Self::InvalidState => SenderBackendError::InvalidState,
             Self::Retry(task) => SenderBackendError::Retry(f(task)),
+            Self::Timeout => SenderBackendError::Timeout,
         }
     }
 }
@@ -717,12 +745,13 @@ impl<T> fmt::Debug for SenderBackendError<T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Self::Io(io_err) => write!(f, "BackendError::Io({:?})", io_err),
-            Self::NodeNotFound => write!(f, "backendError::NodeNotFound"),
-            Self::InvalidProtocol => write!(f, "backendError::InvalidProtocol"),
-            Self::InvalidAddress => write!(f, "backendError::InvalidAddress"),
-            Self::Canceled => write!(f, "backendError::Canceled"),
-            Self::InvalidState => write!(f, "backendError::InvalidState"),
+            Self::NodeNotFound => write!(f, "BackendError::NodeNotFound"),
+            Self::InvalidProtocol => write!(f, "BackendError::InvalidProtocol"),
+            Self::InvalidAddress => write!(f, "BackendError::InvalidAddress"),
+            Self::Canceled => write!(f, "BackendError::Canceled"),
+            Self::InvalidState => write!(f, "BackendError::InvalidState"),
             Self::Retry(_) => write!(f, "BackendError::Retry"),
+            Self::Timeout => write!(f, "BackendError::Timeout"),
         }
     }
 }
@@ -874,6 +903,7 @@ mod tests {
             NonZeroUsize::new(1024).unwrap(),
             Duration::from_nanos(200_000),
             Duration::from_nanos(400_000),
+            Duration::from_secs(10),
         )
         .await;
 
@@ -943,6 +973,7 @@ mod tests {
             NonZeroUsize::new(1024).unwrap(),
             Duration::from_nanos(200_000),
             Duration::from_nanos(400_000),
+            Duration::from_secs(10),
         )
         .await;
 
@@ -957,4 +988,50 @@ mod tests {
             }
         }
     }
+
+    #[tokio::test]
+    async fn test_handle_conn_timeout() {
+        test_handle_conn_timeout_helper(BatchStrategy::Disabled).await;
+        test_handle_conn_timeout_helper(BatchStrategy::Fixed).await;
+        test_handle_conn_timeout_helper(BatchStrategy::Dynamic).await;
+    }
+
+    async fn test_handle_conn_timeout_helper(strategy: BatchStrategy) {
+        let (sender, _dummy_receiver) = mpsc::unbounded::<RespVec>();
+        let (_dummy_sender, receiver) = mpsc::unbounded::<RespVec>();
+        let writer: ConnSink<RespVec> = Box::pin(sender.sink_map_err(|_| BackendError::Canceled));
+        let reader: ConnStream<RespVec> = Box::pin(receiver.map(Ok));
+
+        let (mut task_sender, mut task_receiver) = mpsc::unbounded();
+        for i in 0..3 {
+            let i: usize = i;
+            task_sender
+                .send(DummyCmdTask::new(RespVec::Simple(
+                    i.to_string().into_bytes(),
+                )))
+                .await
+                .unwrap();
+        }
+
+        let handler = Arc::new(CounterHandler::new());
+        let batch_stats = Arc::new(BatchStats::default());
+        let res = handle_conn(
+            writer,
+            reader,
+            &mut task_receiver,
+            handler.clone(),
+            None,
+            batch_stats,
+            strategy,
+            NonZeroUsize::new(1024).unwrap(),
+            Duration::from_nanos(200_000),
+            Duration::from_nanos(400_000),
+            Duration::from_secs(3),
+        )
+        .await;
+
+        let (err, _) = res.unwrap_err();
+        assert!(matches!(err, BackendError::Timeout));
+        assert!(handler.get_replies().is_empty());
+    }
 }

src/proxy/service.rs

@@ -33,6 +33,7 @@ pub struct ServerProxyConfig {
     pub backend_flush_size: NonZeroUsize,
     pub backend_low_flush_interval: Duration,
     pub backend_high_flush_interval: Duration,
+    pub backend_timeout: Duration,
 }
 
 impl ServerProxyConfig {

tests/proxy_manager_test.rs

@@ -63,6 +63,7 @@ mod tests {
             backend_flush_size: NonZeroUsize::new(1024).unwrap(),
             backend_low_flush_interval: Duration::from_nanos(200_000),
             backend_high_flush_interval: Duration::from_nanos(800_000),
+            backend_timeout: Duration::from_secs(3),
         }
     }