feat: support count min sketch data structure and most commands

src/commands/cmd_cms.cc

@@ -30,6 +30,9 @@
 namespace redis {
 
 /// CMS.INCRBY key item increment [item increment ...]
+///
+/// The `key` should be an existing Count-Min Sketch key,
+/// otherwise, the command will return an error.
 class CommandCMSIncrBy final : public Commander {
  public:
   Status Execute(Server *srv, Connection *conn, std::string *output) override {
@@ -39,25 +42,28 @@ class CommandCMSIncrBy final : public Commander {
     redis::CMS cms(srv->storage, conn->GetNamespace());
     engine::Context ctx(srv->storage);
     rocksdb::Status s;
-    std::unordered_map<std::string, uint64_t> elements;
+    // <item, increment> pairs
+    std::vector<redis::CMS::IncrByPair> elements;
+    elements.reserve((args_.size() - 2) / 2);
     for (size_t i = 2; i < args_.size(); i += 2) {
-      std::string key = args_[i];
-      auto parse_result = ParseInt<uint64_t>(args_[i + 1]);
+      std::string_view key = args_[i];
+      auto parse_result = ParseInt<int64_t>(args_[i + 1]);
       if (!parse_result) {
         return {Status::RedisParseErr, errValueNotInteger};
       }
-      uint64_t value = *parse_result;
-      elements[key] = value;
+      int64_t value = *parse_result;
+      elements.emplace_back(redis::CMS::IncrByPair{key, value});
     }
-
-    s = cms.IncrBy(ctx, args_[1], elements);
+    std::vector<uint32_t> counters;
+    s = cms.IncrBy(ctx, args_[1], elements, &counters);
     if (s.IsNotFound()) {
       return {Status::RedisExecErr, "Key not found"};
     }
     if (!s.ok()) {
       return {Status::RedisExecErr, s.ToString()};
     }
 
+    // TODO(mwish): adjust the output
     *output = redis::SimpleString("OK");
     return Status::OK();
   }
@@ -78,7 +84,7 @@ class CommandCMSInfo final : public Commander {
       return {Status::RedisExecErr, "Key not found"};
     }
 
-    if (!s.ok() && !s.IsNotFound()) {
+    if (!s.ok()) {
       return {Status::RedisExecErr, s.ToString()};
     }
 

src/types/cms.cc

@@ -27,36 +27,41 @@
 
 #include "xxhash.h"
 
+uint64_t CMSketch::CountMinSketchHash(std::string_view item, uint64_t seed) {
+  return XXH64(item.data(), item.size(), seed);
+}
+
 CMSketch::CMSketchDimensions CMSketch::CMSDimFromProb(double error, double delta) {
   CMSketchDimensions dims;
   dims.width = std::ceil(2 / error);
   dims.depth = std::ceil(std::log10(delta) / std::log10(0.5));
   return dims;
 }
 
-size_t CMSketch::IncrBy(std::string_view item, uint32_t value) {
-  size_t min_count = std::numeric_limits<size_t>::max();
+uint32_t CMSketch::IncrBy(std::string_view item, uint32_t value) {
+  uint32_t min_count = std::numeric_limits<uint32_t>::max();
 
   for (size_t i = 0; i < depth_; ++i) {
-    uint64_t hash = XXH32(item.data(), static_cast<int>(item.size()), i);
+    uint64_t hash = CountMinSketchHash(item, /*seed=*/i);
     size_t loc = GetLocationForHash(hash, i);
-    if (array_[loc] > UINT32_MAX - value) {
-      array_[loc] = UINT32_MAX;
+    // Do overflow check
+    if (array_[loc] > std::numeric_limits<uint32_t>::max() - value) {
+      array_[loc] = std::numeric_limits<uint32_t>::max();
     } else {
       array_[loc] += value;
     }
-    min_count = std::min(min_count, static_cast<size_t>(array_[loc]));
+    min_count = std::min(min_count, array_[loc]);
   }
   counter_ += value;
   return min_count;
 }
 
-size_t CMSketch::Query(std::string_view item) const {
-  size_t min_count = std::numeric_limits<size_t>::max();
+uint32_t CMSketch::Query(std::string_view item) const {
+  uint32_t min_count = std::numeric_limits<uint32_t>::max();
 
   for (size_t i = 0; i < depth_; ++i) {
-    uint64_t hash = XXH32(item.data(), static_cast<int>(item.size()), i);
-    min_count = std::min(min_count, static_cast<size_t>(array_[GetLocationForHash(hash, i)]));
+    uint64_t hash = CountMinSketchHash(item, /*seed=*/i);
+    min_count = std::min(min_count, array_[GetLocationForHash(hash, i)]);
   }
   return min_count;
 }

src/types/cms.h

@@ -46,9 +46,13 @@ class CMSketch {
 
   static CMSketchDimensions CMSDimFromProb(double error, double delta);
 
-  size_t IncrBy(std::string_view item, uint32_t value);
+  /// Increment the counter of the given item by the specified increment.
+  ///
+  /// \param item The item to increment. Returns UINT32_MAX if the
+  ///             counter overflows.
+  uint32_t IncrBy(std::string_view item, uint32_t value);
 
-  size_t Query(std::string_view item) const;
+  uint32_t Query(std::string_view item) const;
 
   static Status Merge(CMSketch* dest, size_t num_keys, std::vector<const CMSketch*> cms_array,
                       std::vector<uint32_t> weights);
@@ -67,6 +71,9 @@ class CMSketch {
   static int CheckOverflow(CMSketch* dest, size_t quantity, const std::vector<const CMSketch*>& src,
                            const std::vector<uint32_t>& weights);
 
+ private:
+  static uint64_t CountMinSketchHash(std::string_view item, uint64_t seed);
+
  private:
   uint32_t width_;
   uint32_t depth_;

src/types/redis_cms.cc

@@ -31,8 +31,8 @@ rocksdb::Status CMS::GetMetadata(engine::Context &ctx, const Slice &ns_key, Coun
   return Database::GetMetadata(ctx, {kRedisCountMinSketch}, ns_key, metadata);
 }
 
-rocksdb::Status CMS::IncrBy(engine::Context &ctx, const Slice &user_key,
-                            const std::unordered_map<std::string, uint64_t> &elements) {
+rocksdb::Status CMS::IncrBy(engine::Context &ctx, const Slice &user_key, const std::vector<IncrByPair> &elements,
+                            std::vector<uint32_t> *counters) {
   std::string ns_key = AppendNamespacePrefix(user_key);
 
   LockGuard guard(storage_->GetLockManager(), ns_key);
@@ -52,12 +52,12 @@ rocksdb::Status CMS::IncrBy(engine::Context &ctx, const Slice &user_key,
 
   CMSketch cms(metadata.width, metadata.depth, metadata.counter, metadata.array);
 
-  for (auto &element : elements) {
-    if (element.second > 0 && metadata.counter > std::numeric_limits<uint64_t>::max() - element.second) {
+  for (const auto &element : elements) {
+    if (element.value > 0 && metadata.counter > std::numeric_limits<int64_t>::max() - element.value) {
       return rocksdb::Status::InvalidArgument("Overflow error: IncrBy would result in counter overflow");
     }
-    cms.IncrBy(element.first, element.second);
-    metadata.counter += element.second;
+    uint32_t local_counter = cms.IncrBy(element.key, element.value);
+    metadata.counter += element.value;
   }
 
   metadata.array = std::move(cms.GetArray());
@@ -83,20 +83,20 @@ rocksdb::Status CMS::Info(engine::Context &ctx, const Slice &user_key, CMSketch:
   ret->width = metadata.width;
   ret->depth = metadata.depth;
   ret->count = metadata.counter;
-
   return rocksdb::Status::OK();
-};
+}
 
 rocksdb::Status CMS::InitByDim(engine::Context &ctx, const Slice &user_key, uint32_t width, uint32_t depth) {
   std::string ns_key = AppendNamespacePrefix(user_key);
 
   size_t memory_used = width * depth * sizeof(uint32_t);
-  const size_t max_memory = 1 * 1024 * 1024;
+  // We firstly limit the memory usage to 1MB.
+  constexpr size_t kMaxMemory = 1 * 1024 * 1024;
 
   if (memory_used == 0) {
     return rocksdb::Status::InvalidArgument("Memory usage must be greater than 0.");
   }
-  if (memory_used > max_memory) {
+  if (memory_used > kMaxMemory) {
     return rocksdb::Status::InvalidArgument("Memory usage exceeds 1MB.");
   }
 
@@ -137,21 +137,7 @@ rocksdb::Status CMS::InitByProb(engine::Context &ctx, const Slice &user_key, dou
     return rocksdb::Status::InvalidArgument("Delta must be between 0 and 1 (exclusive).");
   }
   CMSketch::CMSketchDimensions dim = CMSketch::CMSDimFromProb(error, delta);
-  auto s = InitByDim(ctx, user_key, dim.width, dim.depth);
-  if (!s.ok()) {
-    return s;
-  }
-
-  size_t memory_used = dim.width * dim.depth * sizeof(uint32_t);
-  const size_t max_memory = 1 * 1024 * 1024;
-
-  if (memory_used == 0) {
-    return rocksdb::Status::InvalidArgument("Memory usage must be greater than 0.");
-  }
-  if (memory_used > max_memory) {
-    return rocksdb::Status::InvalidArgument("Memory usage exceeds 1MB.");
-  }
-  return rocksdb::Status::OK();
+  return InitByDim(ctx, user_key, dim.width, dim.depth);
 };
 
 rocksdb::Status CMS::MergeUserKeys(engine::Context &ctx, const Slice &user_key, const std::vector<Slice> &src_keys,
@@ -165,33 +151,34 @@ rocksdb::Status CMS::MergeUserKeys(engine::Context &ctx, const Slice &user_key,
   }
 
   std::string dest_ns_key = AppendNamespacePrefix(user_key);
-  LockGuard guard(storage_->GetLockManager(), dest_ns_key);
+  std::vector<std::string> ns_keys{dest_ns_key};
+  for (const auto &src_key : src_keys) {
+    ns_keys.emplace_back(AppendNamespacePrefix(src_key));
+  }
+  MultiLockGuard guard(storage_->GetLockManager(), ns_keys);
 
   CountMinSketchMetadata dest_metadata{};
   rocksdb::Status dest_status = GetMetadata(ctx, dest_ns_key, &dest_metadata);
+  if (dest_status.IsNotFound()) {
+    return rocksdb::Status::InvalidArgument("Destination CMS does not exist.");
+  }
   if (!dest_status.ok()) {
-    if (dest_status.IsNotFound()) {
-      return rocksdb::Status::InvalidArgument("Destination CMS does not exist.");
-    }
     return dest_status;
   }
 
   CMSketch dest_cms(dest_metadata.width, dest_metadata.depth, dest_metadata.counter, dest_metadata.array);
 
   std::vector<CMSketch> src_cms_objects;
   src_cms_objects.reserve(num_sources);
-  std::vector<const CMSketch *> src_cms_pointers;
-  src_cms_pointers.reserve(num_sources);
   std::vector<uint32_t> weights_long;
   weights_long.reserve(num_sources);
 
   for (size_t i = 0; i < num_sources; ++i) {
-    std::string src_ns_key = AppendNamespacePrefix(src_keys[i]);
-    LockGuard guard(storage_->GetLockManager(), src_ns_key);
-
+    const auto &src_ns_key = ns_keys[i + 1];
     CountMinSketchMetadata src_metadata{};
     rocksdb::Status src_status = GetMetadata(ctx, src_ns_key, &src_metadata);
     if (!src_status.ok()) {
+      // TODO(mwish): check the not found syntax here.
       if (src_status.IsNotFound()) {
         return rocksdb::Status::InvalidArgument("Source CMS key not found.");
       }
@@ -204,11 +191,15 @@ rocksdb::Status CMS::MergeUserKeys(engine::Context &ctx, const Slice &user_key,
 
     CMSketch src_cms(src_metadata.width, src_metadata.depth, src_metadata.counter, src_metadata.array);
     src_cms_objects.emplace_back(std::move(src_cms));
-    src_cms_pointers.push_back(&src_cms_objects.back());
 
     weights_long.push_back(static_cast<uint32_t>(src_weights[i]));
   }
-
+  // Initialize the destination CMS with the source CMSes after initializations
+  // since vector might resize and reallocate memory.
+  std::vector<const CMSketch *> src_cms_pointers(num_sources);
+  for (size_t i = 0; i < num_sources; ++i) {
+    src_cms_pointers[i] = &src_cms_objects[i];
+  }
   auto merge_result = CMSketch::Merge(&dest_cms, num_sources, src_cms_pointers, weights_long);
   if (!merge_result.IsOK()) {
     return rocksdb::Status::InvalidArgument("Merge operation failed due to overflow or invalid dimensions.");
@@ -240,7 +231,8 @@ rocksdb::Status CMS::Query(engine::Context &ctx, const Slice &user_key, const st
 
   if (s.IsNotFound()) {
     return rocksdb::Status::NotFound();
-  } else if (!s.ok()) {
+  }
+  if (!s.ok()) {
     return s;
   }
 
@@ -251,6 +243,6 @@ rocksdb::Status CMS::Query(engine::Context &ctx, const Slice &user_key, const st
   }
 
   return rocksdb::Status::OK();
-};
+}
 
 }  // namespace redis

src/types/redis_cms.h

@@ -20,6 +20,8 @@
 
 #pragma once
 
+#include <vector>
+
 #include "cms.h"
 #include "storage/redis_db.h"
 #include "storage/redis_metadata.h"
@@ -30,8 +32,17 @@ class CMS : public Database {
  public:
   explicit CMS(engine::Storage *storage, const std::string &ns) : Database(storage, ns) {}
 
-  rocksdb::Status IncrBy(engine::Context &ctx, const Slice &user_key,
-                         const std::unordered_map<std::string, uint64_t> &elements);
+  struct IncrByPair {
+    std::string_view key;
+    int64_t value;
+  };
+
+  /// Increment the counter of the given item(s) by the specified increment(s).
+  ///
+  /// \param[out] counters The counter values after the increment, if the value is UINT32_MAX,
+  ///                      it means the item does overflow.
+  rocksdb::Status IncrBy(engine::Context &ctx, const Slice &user_key, const std::vector<IncrByPair> &elements,
+                         std::vector<uint32_t> *counters);
   rocksdb::Status Info(engine::Context &ctx, const Slice &user_key, CMSketch::CMSInfo *ret);
   rocksdb::Status InitByDim(engine::Context &ctx, const Slice &user_key, uint32_t width, uint32_t depth);
   rocksdb::Status InitByProb(engine::Context &ctx, const Slice &user_key, double error, double delta);