(kvrocks) branch unstable updated: Fix rand member didn't work in the command SRANDMEMBER (#2032)

[hulk]

This is an automated email from the ASF dual-hosted git repository.

hulk pushed a commit to branch unstable
in repository https://gitbox.apache.org/repos/asf/kvrocks.git


The following commit(s) were added to refs/heads/unstable by this push:
     new 667c1458 Fix rand member didn't work in the command SRANDMEMBER (#2032)
667c1458 is described below

commit 667c1458c23db6938cb2a001bbc67c78eb21485e
Author: Lee <wse1714401...@gmail.com>
AuthorDate: Tue Jan 23 14:48:06 2024 +0800

    Fix rand member didn't work in the command SRANDMEMBER (#2032)
---
 src/types/redis_set.cc | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/types/redis_set.cc b/src/types/redis_set.cc
index 32c06ef9..98677e27 100644
--- a/src/types/redis_set.cc
+++ b/src/types/redis_set.cc
@@ -23,6 +23,7 @@
 #include <map>
 #include <memory>
 #include <optional>
+#include <random>
 
 #include "db_util.h"
 
@@ -222,15 +223,32 @@ rocksdb::Status Set::Take(const Slice &user_key, 
std::vector<std::string> *membe
   rocksdb::Slice upper_bound(next_version_prefix);
   read_options.iterate_upper_bound = &upper_bound;
 
+  std::vector<std::string> iter_keys;
+  iter_keys.reserve(count);
+  std::random_device rd;
+  std::mt19937 gen(rd());
   auto iter = util::UniqueIterator(storage_, read_options);
   for (iter->Seek(prefix); iter->Valid() && iter->key().starts_with(prefix); 
iter->Next()) {
-    InternalKey ikey(iter->key(), storage_->IsSlotIdEncoded());
+    ++n;
+    if (n <= count) {
+      iter_keys.push_back(iter->key().ToString());
+    } else {  // n > count
+      std::uniform_int_distribution<> distrib(0, n - 1);
+      int random = distrib(gen);  // [0,n-1]
+      if (random < count) {
+        iter_keys[random] = iter->key().ToString();
+      }
+    }
+  }
+  for (Slice key : iter_keys) {
+    InternalKey ikey(key, storage_->IsSlotIdEncoded());
     members->emplace_back(ikey.GetSubKey().ToString());
-    if (pop) batch->Delete(iter->key());
-    if (++n >= count) break;
+    if (pop) {
+      batch->Delete(key);
+    }
   }
-  if (pop && n > 0) {
-    metadata.size -= n;
+  if (pop && !iter_keys.empty()) {
+    metadata.size -= iter_keys.size();
     std::string bytes;
     metadata.Encode(&bytes);
     batch->Put(metadata_cf_handle_, ns_key, bytes);