From 50a944a77b45fbd1dab945f571cd2ceb532771e1 Mon Sep 17 00:00:00 2001
From: Shachar Langbeheim <98546660+shachlanAmazon@users.noreply.github.com>
Date: Wed, 20 Dec 2023 16:04:48 +0200
Subject: [PATCH] Add timeouts to standalone pipeline requests. (#706)

Co-authored-by: nihohit <nihohit@gmail.com>
---
 babushka-core/src/client/mod.rs    |  9 +++------
 babushka-core/tests/test_client.rs | 26 ++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/babushka-core/src/client/mod.rs b/babushka-core/src/client/mod.rs
index 7e90cb43f7..afdd45f7ba 100644
--- a/babushka-core/src/client/mod.rs
+++ b/babushka-core/src/client/mod.rs
@@ -135,7 +135,7 @@ impl Client {
         count: usize,
         routing: Option<RoutingInfo>,
     ) -> redis::RedisFuture<'a, Vec<redis::Value>> {
-        (async move {
+        run_with_timeout(self.request_timeout, async move {
             match self.internal_client {
                 ClientWrapper::Standalone(ref mut client) => {
                     client.send_packed_commands(cmd, offset, count).await
@@ -146,11 +146,8 @@ impl Client {
                         Some(RoutingInfo::SingleNode(route)) => route,
                         _ => SingleNodeRoutingInfo::Random,
                     };
-                    run_with_timeout(
-                        self.request_timeout,
-                        client.route_pipeline(cmd, offset, count, route),
-                    )
-                    .await
+
+                    client.route_pipeline(cmd, offset, count, route).await
                 }
             }
         })
diff --git a/babushka-core/tests/test_client.rs b/babushka-core/tests/test_client.rs
index 1ba84cfee9..465711d103 100644
--- a/babushka-core/tests/test_client.rs
+++ b/babushka-core/tests/test_client.rs
@@ -257,4 +257,30 @@ mod shared_client_tests {
             assert!(err.is_timeout(), "{err}");
         });
     }
+
+    #[rstest]
+    #[timeout(SHORT_CLUSTER_TEST_TIMEOUT)]
+    fn test_request_pipeline_timeout(#[values(false, true)] use_cluster: bool) {
+        block_on_all(async {
+            let mut test_basics = setup_test_basics(
+                use_cluster,
+                TestConfiguration {
+                    request_timeout: Some(1),
+                    shared_server: true,
+                    ..Default::default()
+                },
+            )
+            .await;
+
+            let mut pipeline = redis::pipe();
+            pipeline.blpop("foo", 0.0); // 0 timeout blocks indefinitely
+            let result = test_basics
+                .client
+                .req_packed_commands(&pipeline, 0, 1, None)
+                .await;
+            assert!(result.is_err());
+            let err = result.unwrap_err();
+            assert!(err.is_timeout(), "{err}");
+        });
+    }
 }