From 5904289e8aba1dc1f24a0723763205887ad50957 Mon Sep 17 00:00:00 2001
From: Jerrod Carpenter <4128301+JerrodCarpenter@users.noreply.github.com>
Date: Tue, 9 May 2023 08:04:54 -0500
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20ZMPOP?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fix: https://github.com/redis/redis-rb/pull/1189
---
 lib/redis/commands/sorted_sets.rb | 36 +++++++++++++++++++++++++++++++
 lib/redis/distributed.rb          |  7 ++++++
 test/lint/sorted_sets.rb          | 14 ++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/lib/redis/commands/sorted_sets.rb b/lib/redis/commands/sorted_sets.rb
index d74230503..1f8f24532 100644
--- a/lib/redis/commands/sorted_sets.rb
+++ b/lib/redis/commands/sorted_sets.rb
@@ -167,6 +167,42 @@ 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.zmpop('zset')
+      #   #=> ['zset', ['a', 1.0]]
+      # @example With count option
+      #   redis.zmpop('zset', count: 2)
+      #   #=> ['zset', [['a', 1.0], ['b', 2.0]]
+      #
+      # @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 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
+
+        send_command(args) 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 the highest scores in the sorted set stored at keys,
       #   or block until one is available.
       #
diff --git a/lib/redis/distributed.rb b/lib/redis/distributed.rb
index c238d2939..d26e5cf11 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, removing members from the first non empty sorted set found.
+    def zmpop(*keys, modifier: "MIN", count: nil)
+      ensure_same_node(:zmpop, keys) do |node|
+        node.zmpop(*keys, modifier: modifier, count: count)
+      end
+    end
+
     # Return a range of members in a sorted set, by index, score or lexicographical ordering.
     def zrange(key, start, stop, **options)
       node_for(key).zrange(key, start, stop, **options)
diff --git a/test/lint/sorted_sets.rb b/test/lint/sorted_sets.rb
index c523dc3f1..2f5b934e9 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_zmpop
+      target_version('7.0') do
+        assert_nil r.zmpop('{1}foo')
+
+        r.zadd('{1}foo', %w[0 a 1 b 2 c 3 d])
+        assert_equal ['{1}foo', [['a', 0.0]]], r.zmpop('{1}foo')
+        assert_equal ['{1}foo', [['b', 1.0], ['c', 2.0], ['d', 3.0]]], r.zmpop('{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.zmpop('{1}foo', '{1}foo2', modifier: "MAX")
+      end
+    end
+
     def test_zremrangebylex
       r.zadd('foo', %w[0 a 0 b 0 c 0 d 0 e 0 f 0 g])
       assert_equal 5, r.zremrangebylex('foo', '(b', '[g')