From 0c66e041b5e43f02ff1e2c68bfb86bf0e3b679fe Mon Sep 17 00:00:00 2001
From: #Suyghur <suyghurmjp@outlook.com>
Date: Sat, 1 Oct 2022 20:53:54 +0800
Subject: [PATCH] feat(redis):add timeout method to extend blpop (#2472)

---
 core/stores/redis/redis.go      | 25 +++++++++++++++++++++++++
 core/stores/redis/redis_test.go | 11 +++++++++++
 2 files changed, 36 insertions(+)

diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go
index c4d9c401..535f523a 100644
--- a/core/stores/redis/redis.go
+++ b/core/stores/redis/redis.go
@@ -267,6 +267,31 @@ func (s *Redis) BlpopExCtx(ctx context.Context, node RedisNode, key string) (str
 	return vals[1], true, nil
 }
 
+// BlpopWithTimeout uses passed in redis connection to execute blpop command.
+// Control blocking query timeout
+func (s *Redis) BlpopWithTimeout(node RedisNode, timeout time.Duration, key string) (string, error) {
+	return s.BlpopWithTimeoutCtx(context.Background(), node, timeout, key)
+}
+
+// BlpopWithTimeoutCtx uses passed in redis connection to execute blpop command.
+// Control blocking query timeout
+func (s *Redis) BlpopWithTimeoutCtx(ctx context.Context, node RedisNode, timeout time.Duration, key string) (string, error) {
+	if node == nil {
+		return "", ErrNilNode
+	}
+
+	vals, err := node.BLPop(ctx, timeout, key).Result()
+	if err != nil {
+		return "", err
+	}
+
+	if len(vals) < 2 {
+		return "", fmt.Errorf("no value on key: %s", key)
+	}
+
+	return vals[1], nil
+}
+
 // Decr is the implementation of redis decr command.
 func (s *Redis) Decr(key string) (int64, error) {
 	return s.DecrCtx(context.Background(), key)
diff --git a/core/stores/redis/redis_test.go b/core/stores/redis/redis_test.go
index 4f960909..3dd2a576 100644
--- a/core/stores/redis/redis_test.go
+++ b/core/stores/redis/redis_test.go
@@ -1117,6 +1117,17 @@ func TestRedisBlpopEx(t *testing.T) {
 	})
 }
 
+func TestRedisBlpopWithTimeout(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		var node mockedNode
+		_, err := client.BlpopWithTimeout(nil, 10*time.Second, "foo")
+		assert.NotNil(t, err)
+		_, err = client.BlpopWithTimeout(node, 10*time.Second, "foo")
+		assert.NotNil(t, err)
+	})
+}
+
 func TestRedisGeo(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		client.Ping()