Fix blocking commands tests in wrappers

barshaul · barshaul · commit a5dc2d28645f · 2024-04-16T08:48:11.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@
 
 #### Fixes
 * Python: Fix typing error "‘type’ object is not subscriptable" ([#1203](https://github.com/aws/glide-for-redis/pull/1203))
+* Core: Fixed blocking commands to use the specified timeout from the command argument
 
 ## 0.3.3 (2024-03-28)
 
diff --git a/glide-core/src/client/mod.rs b/glide-core/src/client/mod.rs
@@ -8,7 +8,7 @@ use futures::FutureExt;
 use logger_core::log_info;
 use redis::aio::ConnectionLike;
 use redis::cluster_async::ClusterConnection;
-use redis::cluster_routing::{RoutingInfo, SingleNodeRoutingInfo};
+use redis::cluster_routing::{Routable, RoutingInfo, SingleNodeRoutingInfo};
 use redis::RedisResult;
 use redis::{Cmd, ErrorKind, Value};
 pub use standalone_client::StandaloneClient;
@@ -104,14 +104,57 @@ async fn run_with_timeout<T>(
         .and_then(|res| res)
 }
 
+const BLOCKING_CMD_TIMEOUT_BUFFER: f64 = 0.5; // seconds
+
+enum TimeUnit {
+    Milliseconds = 1000,
+    Seconds = 1,
+}
+
+// Attempts to get the timeout duration from the command argument at `timeout_idx`.
+// If the argument can be parsed into a duration, it returns the duration in seconds. Otherwise, it returns None.
+fn try_get_timeout_from_cmd_arg(
+    cmd: &Cmd,
+    timeout_idx: usize,
+    time_unit: TimeUnit,
+) -> Option<Duration> {
+    cmd.arg_idx(timeout_idx).and_then(|timeout_bytes| {
+        std::str::from_utf8(timeout_bytes)
+            .ok()
+            .and_then(|timeout_str| {
+                timeout_str.parse::<f64>().ok().map(|timeout| {
+                    let timeout_secs = timeout / ((time_unit as i32) as f64);
+                    // We add a buffer to the request timeout to ensure we receive a response from the server
+                    Duration::from_secs_f64(timeout_secs + BLOCKING_CMD_TIMEOUT_BUFFER)
+                })
+            })
+    })
+}
+
+fn get_request_timeout(cmd: &Cmd, default_timeout: Duration) -> Duration {
+    let command = cmd.command().unwrap_or_default();
+    let blocking_timeout = match command.as_slice() {
+        b"BLPOP" | b"BRPOP" | b"BLMOVE" | b"BZPOPMAX" | b"BZPOPMIN" | b"BRPOPLPUSH" => {
+            try_get_timeout_from_cmd_arg(cmd, cmd.args_iter().len() - 1, TimeUnit::Seconds)
+        }
+        b"BLMPOP" | b"BZMPOP" => try_get_timeout_from_cmd_arg(cmd, 1, TimeUnit::Seconds),
+        b"XREAD" | b"XREADGROUP" => cmd
+            .position(b"BLOCK")
+            .and_then(|idx| try_get_timeout_from_cmd_arg(cmd, idx + 1, TimeUnit::Milliseconds)),
+        _ => None,
+    };
+
+    blocking_timeout.unwrap_or(default_timeout)
+}
+
 impl Client {
     pub fn send_command<'a>(
         &'a mut self,
         cmd: &'a Cmd,
         routing: Option<RoutingInfo>,
     ) -> redis::RedisFuture<'a, Value> {
         let expected_type = expected_type_for_cmd(cmd);
-        run_with_timeout(self.request_timeout, async move {
+        run_with_timeout(get_request_timeout(cmd, self.request_timeout), async move {
             match self.internal_client {
                 ClientWrapper::Standalone(ref mut client) => client.send_command(cmd).await,
 
@@ -472,3 +515,95 @@ impl GlideClientForTests for StandaloneClient {
         self.send_command(cmd).boxed()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use redis::Cmd;
+
+    use crate::client::{get_request_timeout, TimeUnit, BLOCKING_CMD_TIMEOUT_BUFFER};
+
+    use super::try_get_timeout_from_cmd_arg;
+
+    #[test]
+    fn test_get_timeout_from_cmd_returns_correct_duration_int() {
+        let mut cmd = Cmd::new();
+        cmd.arg("BLPOP").arg("key1").arg("key2").arg("5");
+        assert_eq!(
+            try_get_timeout_from_cmd_arg(&cmd, cmd.args_iter().len() - 1, TimeUnit::Seconds),
+            Some(Duration::from_secs_f64(5.0 + BLOCKING_CMD_TIMEOUT_BUFFER))
+        );
+    }
+
+    #[test]
+    fn test_get_timeout_from_cmd_returns_correct_duration_float() {
+        let mut cmd = Cmd::new();
+        cmd.arg("BLPOP").arg("key1").arg("key2").arg(0.5);
+        assert_eq!(
+            try_get_timeout_from_cmd_arg(&cmd, cmd.args_iter().len() - 1, TimeUnit::Seconds),
+            Some(Duration::from_secs_f64(0.5 + BLOCKING_CMD_TIMEOUT_BUFFER))
+        );
+    }
+
+    #[test]
+    fn test_get_timeout_from_cmd_returns_correct_duration_milliseconds() {
+        let mut cmd = Cmd::new();
+        cmd.arg("XREAD").arg("BLOCK").arg("500").arg("key");
+        assert_eq!(
+            try_get_timeout_from_cmd_arg(&cmd, 2, TimeUnit::Milliseconds),
+            Some(Duration::from_secs_f64(0.5 + BLOCKING_CMD_TIMEOUT_BUFFER))
+        );
+    }
+
+    #[test]
+    fn test_get_timeout_from_cmd_returns_none_when_timeout_isnt_passed() {
+        let mut cmd = Cmd::new();
+        cmd.arg("BLPOP").arg("key1").arg("key2").arg("key3");
+        assert_eq!(
+            try_get_timeout_from_cmd_arg(&cmd, cmd.args_iter().len() - 1, TimeUnit::Seconds),
+            None,
+        );
+    }
+
+    #[test]
+    fn test_get_request_timeout_with_blocking_command_returns_cmd_arg_timeout() {
+        let mut cmd = Cmd::new();
+        cmd.arg("BLPOP").arg("key1").arg("key2").arg("500");
+        assert_eq!(
+            get_request_timeout(&cmd, Duration::from_millis(100)),
+            Duration::from_secs_f64(500.0 + BLOCKING_CMD_TIMEOUT_BUFFER)
+        );
+
+        let mut cmd = Cmd::new();
+        cmd.arg("XREADGROUP").arg("BLOCK").arg("500").arg("key");
+        assert_eq!(
+            get_request_timeout(&cmd, Duration::from_millis(100)),
+            Duration::from_secs_f64(0.5 + BLOCKING_CMD_TIMEOUT_BUFFER)
+        );
+
+        let mut cmd = Cmd::new();
+        cmd.arg("BLMPOP").arg("0.857").arg("key");
+        assert_eq!(
+            get_request_timeout(&cmd, Duration::from_millis(100)),
+            Duration::from_secs_f64(0.857 + BLOCKING_CMD_TIMEOUT_BUFFER)
+        );
+    }
+
+    #[test]
+    fn test_get_request_timeout_non_blocking_command_returns_default_timeout() {
+        let mut cmd = Cmd::new();
+        cmd.arg("SET").arg("key").arg("value").arg("PX").arg("500");
+        assert_eq!(
+            get_request_timeout(&cmd, Duration::from_millis(100)),
+            Duration::from_millis(100)
+        );
+
+        let mut cmd = Cmd::new();
+        cmd.arg("XREADGROUP").arg("key");
+        assert_eq!(
+            get_request_timeout(&cmd, Duration::from_millis(100)),
+            Duration::from_millis(100)
+        );
+    }
+}
diff --git a/glide-core/tests/test_client.rs b/glide-core/tests/test_client.rs
@@ -321,21 +321,58 @@ pub(crate) mod shared_client_tests {
                 use_cluster,
                 TestConfiguration {
                     request_timeout: Some(1),
-                    shared_server: true,
+                    shared_server: false,
                     ..Default::default()
                 },
             )
             .await;
-
             let mut cmd = redis::Cmd::new();
-            cmd.arg("BLPOP").arg("foo").arg(0); // 0 timeout blocks indefinitely
+            cmd.arg("EVAL")
+                .arg(
+                    r#"
+                    local i = 0
+                    while (true)
+                    do
+                    redis.call('ping')
+                    i = i + 1
+                    end
+                "#,
+                )
+                .arg("0"); // Create a long running command to ensure we get into timeout
             let result = test_basics.client.send_command(&cmd, None).await;
             assert!(result.is_err());
             let err = result.unwrap_err();
             assert!(err.is_timeout(), "{err}");
         });
     }
 
+    #[rstest]
+    #[timeout(SHORT_CLUSTER_TEST_TIMEOUT)]
+    fn test_blocking_command_doesnt_raise_timeout_error(#[values(false, true)] use_cluster: bool) {
+        // We test that the request timeout is based on the value specified in the blocking command argument,
+        // and not on the one set in the client configuration. To achieve this, we execute a command designed to
+        // be blocked until it reaches the specified command timeout. We set the client's request timeout to
+        // a shorter duration than the blocking command's timeout. Subsequently, we confirm that we receive
+        // a response from the server instead of encountering a timeout error.
+        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 cmd = redis::Cmd::new();
+            cmd.arg("BLPOP").arg("foo").arg(0.3); // command should be timed out after 300 millisecond
+            let result = test_basics.client.send_command(&cmd, None).await;
+            assert!(result.is_ok());
+            assert_eq!(result.unwrap(), Value::Nil);
+        });
+    }
+
     #[rstest]
     #[timeout(SHORT_CLUSTER_TEST_TIMEOUT)]
     fn test_request_transaction_timeout(#[values(false, true)] use_cluster: bool) {
diff --git a/java/integTest/src/test/java/glide/SharedCommandTests.java b/java/integTest/src/test/java/glide/SharedCommandTests.java
@@ -1814,14 +1814,14 @@ public void brpop(BaseClient client) {
         // nothing popped out
         assertNull(
                 client
-                        .brpop(new String[] {listKey2}, REDIS_VERSION.isLowerThan("7.0.0") ? 1. : 0.001)
+                        .brpop(new String[] {listKey2}, REDIS_VERSION.isLowerThan("7.0.0") ? 1. : 0.5)
                         .get());
 
         // Key exists, but it is not a list
         assertEquals(OK, client.set("foo", "bar").get());
         ExecutionException executionException =
                 assertThrows(
-                        ExecutionException.class, () -> client.brpop(new String[] {"foo"}, .0001).get());
+                        ExecutionException.class, () -> client.brpop(new String[] {"foo"}, 0.5).get());
         assertTrue(executionException.getCause() instanceof RequestException);
     }
 
@@ -1867,14 +1867,14 @@ public void blpop(BaseClient client) {
         // nothing popped out
         assertNull(
                 client
-                        .blpop(new String[] {listKey2}, REDIS_VERSION.isLowerThan("7.0.0") ? 1. : 0.001)
+                        .blpop(new String[] {listKey2}, REDIS_VERSION.isLowerThan("7.0.0") ? 1. : 0.5)
                         .get());
 
         // Key exists, but it is not a list
         assertEquals(OK, client.set("foo", "bar").get());
         ExecutionException executionException =
                 assertThrows(
-                        ExecutionException.class, () -> client.blpop(new String[] {"foo"}, .0001).get());
+                        ExecutionException.class, () -> client.blpop(new String[] {"foo"}, 0.5).get());
         assertTrue(executionException.getCause() instanceof RequestException);
     }
 
diff --git a/node/tests/SharedTests.ts b/node/tests/SharedTests.ts
@@ -2041,14 +2041,14 @@ export function runBaseTests<Context>(config: {
                     await client.rpush("brpop-test", ["foo", "bar", "baz"]),
                 ).toEqual(3);
                 // Test basic usage
-                expect(await client.brpop(["brpop-test"], 0.1)).toEqual([
+                expect(await client.brpop(["brpop-test"], 0.5)).toEqual([
                     "brpop-test",
                     "baz",
                 ]);
                 // Delete all values from list
                 expect(await client.del(["brpop-test"])).toEqual(1);
                 // Test null return when key doesn't exist
-                expect(await client.brpop(["brpop-test"], 0.1)).toEqual(null);
+                expect(await client.brpop(["brpop-test"], 0.5)).toEqual(null);
             }, protocol);
         },
         config.timeout,
diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py
@@ -1945,7 +1945,7 @@ class TestExceptions:
     async def test_timeout_exception_with_blpop(self, redis_client: TRedisClient):
         key = get_random_string(10)
         with pytest.raises(TimeoutError):
-            await redis_client.custom_command(["BLPOP", key, "1"])
+            await redis_client.custom_command(["BLPOP", key, "0.1"])
 
 
 @pytest.mark.asyncio
