Extract a rand sample helper and support negative sample count in "Set"

src/types/redis_set.cc

@@ -197,9 +197,14 @@ rocksdb::Status Set::MIsMember(const Slice &user_key, const std::vector<Slice> &
 }
 
 rocksdb::Status Set::Take(const Slice &user_key, std::vector<std::string> *members, int count, bool pop) {
-  int n = 0;
   members->clear();
-  if (count <= 0) return rocksdb::Status::OK();
+  if (count == 0) return rocksdb::Status::OK();
+  if (count < 0) {
+    DCHECK(!pop);
+    // NOTE: Currently, for SRANDMEMBER, we don't
+    // make duplicate members to be returned.
+    count = -count;
+  }
 
   std::string ns_key = AppendNamespacePrefix(user_key);
 
@@ -210,9 +215,11 @@ rocksdb::Status Set::Take(const Slice &user_key, std::vector<std::string> *membe
   rocksdb::Status s = GetMetadata(ns_key, &metadata);
   if (!s.ok()) return s.IsNotFound() ? rocksdb::Status::OK() : s;
 
-  auto batch = storage_->GetWriteBatchBase();
-  WriteBatchLogData log_data(kRedisSet);
-  batch->PutLogData(log_data.Encode());
+  ObserverOrUniquePtr<rocksdb::WriteBatchBase> batch = storage_->GetWriteBatchBase();
+  if (pop) {
+    WriteBatchLogData log_data(kRedisSet);
+    batch->PutLogData(log_data.Encode());
+  }
 
   std::string prefix = InternalKey(ns_key, "", metadata.version, storage_->IsSlotIdEncoded()).Encode();
   std::string next_version_prefix = InternalKey(ns_key, "", metadata.version + 1, storage_->IsSlotIdEncoded()).Encode();
@@ -229,11 +236,10 @@ rocksdb::Status Set::Take(const Slice &user_key, std::vector<std::string> *membe
   std::mt19937 gen(rd());
   auto iter = util::UniqueIterator(storage_, read_options);
   for (iter->Seek(prefix); iter->Valid() && iter->key().starts_with(prefix); iter->Next()) {
-    ++n;
-    if (n <= count) {
+    if (iter_keys.size() < count) {
       iter_keys.push_back(iter->key().ToString());
     } else {  // n > count
-      std::uniform_int_distribution<> distrib(0, n - 1);
+      std::uniform_int_distribution<> distrib(0, static_cast<int>(iter_keys.size()));
       int random = distrib(gen);  // [0,n-1]
       if (random < count) {
         iter_keys[random] = iter->key().ToString();
@@ -247,7 +253,10 @@ rocksdb::Status Set::Take(const Slice &user_key, std::vector<std::string> *membe
       batch->Delete(key);
     }
   }
-  if (pop && !iter_keys.empty()) {
+  if (!pop) {
+    return rocksdb::Status::OK();
+  }
+  if (!iter_keys.empty()) {
     metadata.size -= iter_keys.size();
     std::string bytes;
     metadata.Encode(&bytes);

tests/cppunit/types/set_test.cc

@@ -275,6 +275,9 @@ TEST_F(RedisSetTest, TakeWithoutPop) {
   s = set_->Take(key_, &members, int(fields_.size() - 1), false);
   EXPECT_TRUE(s.ok());
   EXPECT_EQ(members.size(), fields_.size() - 1);
+  s = set_->Take(key_, &members, -int(fields_.size() - 1), false);
+  EXPECT_TRUE(s.ok());
+  EXPECT_EQ(members.size(), fields_.size() - 1);
   s = set_->Remove(key_, fields_, &ret);
   EXPECT_TRUE(s.ok() && fields_.size() == ret);
   s = set_->Del(key_);