From faf054d5800b6f38eccb6ad31d6043d1b7d4d98e Mon Sep 17 00:00:00 2001
From: Jerrod Carpenter <4128301+JerrodCarpenter@users.noreply.github.com>
Date: Mon, 22 May 2023 07:56:59 -0500
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20BZMPOP?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 lib/redis/commands/sorted_sets.rb | 40 +++++++++++++++++++++++++++----
 lib/redis/distributed.rb          |  7 ++++++
 test/lint/sorted_sets.rb          | 14 +++++++++++
 3 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/lib/redis/commands/sorted_sets.rb b/lib/redis/commands/sorted_sets.rb
index 1f8f24532..f07f95275 100644
--- a/lib/redis/commands/sorted_sets.rb
+++ b/lib/redis/commands/sorted_sets.rb
@@ -167,6 +167,40 @@ def zpopmin(key, count = nil)
         end
       end
 
+      # Removes and returns up to count members with scores in the sorted set stored at key.
+      #
+      # @example Popping a member
+      #   redis.bzmpop('zset')
+      #   #=> ['zset', ['a', 1.0]]
+      # @example With count option
+      #   redis.bzmpop('zset', count: 2)
+      #   #=> ['zset', [['a', 1.0], ['b', 2.0]]
+      #
+      # @params timeout [Float] a float value specifying the maximum number of seconds to block) elapses.
+      #   A timeout of zero can be used to block indefinitely.
+      # @params key [String, Array<String>] one or more keys with sorted sets
+      # @params modifier [String]
+      #  - when `"MIN"` - the elements popped are those with lowest scores
+      #  - when `"MAX"` - the elements popped are those with the highest scores
+      # @params count [Integer] a number of members to pop
+      #
+      # @return [Array<String, Array<String, Float>>] list of popped elements and scores
+      def bzmpop(timeout, *keys, modifier: "MIN", count: nil)
+        raise ArgumentError, "Pick either MIN or MAX" unless modifier == "MIN" || modifier == "MAX"
+
+        args = [:bzmpop, timeout, keys.size, *keys, modifier]
+        args << "COUNT" << Integer(count) if count
+
+        send_blocking_command(args, timeout) do |response|
+          response&.map do |entry|
+            case entry
+            when String then entry
+            when Array then entry.map { |pair| FloatifyPairs.call(pair) }.flatten(1)
+            end
+          end
+        end
+      end
+
       # Removes and returns up to count members with scores in the sorted set stored at key.
       #
       # @example Popping a member
@@ -187,11 +221,7 @@ def zmpop(*keys, modifier: "MIN", count: nil)
         raise ArgumentError, "Pick either MIN or MAX" unless modifier == "MIN" || modifier == "MAX"
 
         args = [:zmpop, keys.size, *keys, modifier]
-
-        if count
-          args << "COUNT"
-          args << Integer(count)
-        end
+        args << "COUNT" << Integer(count) if count
 
         send_command(args) do |response|
           response&.map do |entry|
diff --git a/lib/redis/distributed.rb b/lib/redis/distributed.rb
index d26e5cf11..75d33cc80 100644
--- a/lib/redis/distributed.rb
+++ b/lib/redis/distributed.rb
@@ -694,6 +694,13 @@ def zmscore(key, *members)
       node_for(key).zmscore(key, *members)
     end
 
+    # Iterate over keys, blocking and removing members from the first non empty sorted set found.
+    def bzmpop(timeout, *keys, modifier: "MIN", count: nil)
+      ensure_same_node(:bzmpop, keys) do |node|
+        node.bzmpop(timeout, *keys, modifier: modifier, count: count)
+      end
+    end
+
     # Iterate over keys, removing members from the first non empty sorted set found.
     def zmpop(*keys, modifier: "MIN", count: nil)
       ensure_same_node(:zmpop, keys) do |node|
diff --git a/test/lint/sorted_sets.rb b/test/lint/sorted_sets.rb
index 2f5b934e9..10f8ce5e9 100644
--- a/test/lint/sorted_sets.rb
+++ b/test/lint/sorted_sets.rb
@@ -479,6 +479,20 @@ def test_zpopmin
       assert_equal [['d', 3.0]], r.zrange('foo', 0, -1, with_scores: true)
     end
 
+    def test_bzmpop
+      target_version('7.0') do
+        assert_nil r.bzmpop(1.0, '{1}foo')
+
+        r.zadd('{1}foo', %w[0 a 1 b 2 c 3 d])
+        assert_equal ['{1}foo', [['a', 0.0]]], r.bzmpop(1.0, '{1}foo')
+        assert_equal ['{1}foo', [['b', 1.0], ['c', 2.0], ['d', 3.0]]], r.bzmpop(1.0, '{1}foo', count: 4)
+
+        r.zadd('{1}foo', %w[0 a 1 b 2 c 3 d])
+        r.zadd('{1}foo2', %w[0 a 1 b 2 c 3 d])
+        assert_equal ['{1}foo', [['d', 3.0]]], r.bzmpop(1.0, '{1}foo', '{1}foo2', modifier: "MAX")
+      end
+    end
+
     def test_zmpop
       target_version('7.0') do
         assert_nil r.zmpop('{1}foo')