Use the rocksdb MultiGet in the MGet command (apache#331)

ShooterIT · Sep 9, 2021 · 2cdd69d · 2cdd69d
1 parent f92cc41
commit 2cdd69d
Show file tree

Hide file tree

Showing 8 changed files with 69 additions and 19 deletions.
diff --git a/src/redis_bitmap.cc b/src/redis_bitmap.cc
@@ -36,12 +36,12 @@ rocksdb::Status Bitmap::GetMetadata(const Slice &ns_key, BitmapMetadata *metadat
 
   if (metadata->Expired()) {
     metadata->Decode(old_metadata);
-    return rocksdb::Status::NotFound("the key was Expired");
+    return rocksdb::Status::NotFound(kErrMsgKeyExpired);
   }
   if (metadata->Type() == kRedisString) return s;
   if (metadata->Type() != kRedisBitmap && metadata->size > 0) {
     metadata->Decode(old_metadata);
-    return rocksdb::Status::InvalidArgument("WRONGTYPE Operation against a key holding the wrong kind of value");
+    return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
   }
   if (metadata->size == 0) {
     metadata->Decode(old_metadata);

diff --git a/src/redis_db.cc b/src/redis_db.cc
@@ -23,11 +23,11 @@ rocksdb::Status Database::GetMetadata(RedisType type, const Slice &ns_key, Metad
 
   if (metadata->Expired()) {
     metadata->Decode(old_metadata);
-    return rocksdb::Status::NotFound("the key was Expired");
+    return rocksdb::Status::NotFound(kErrMsgKeyExpired);
   }
   if (metadata->Type() != type && (metadata->size > 0 || metadata->Type() == kRedisString)) {
     metadata->Decode(old_metadata);
-    return rocksdb::Status::InvalidArgument("WRONGTYPE Operation against a key holding the wrong kind of value");
+    return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
   }
   if (metadata->size == 0) {
     metadata->Decode(old_metadata);
@@ -60,7 +60,7 @@ rocksdb::Status Database::Expire(const Slice &user_key, int timestamp) {
   if (!s.ok()) return s;
   metadata.Decode(value);
   if (metadata.Expired()) {
-    return rocksdb::Status::NotFound("the key was expired");
+    return rocksdb::Status::NotFound(kErrMsgKeyExpired);
   }
   if (metadata.Type() != kRedisString && metadata.size == 0) {
     return rocksdb::Status::NotFound("no elements");
@@ -91,7 +91,7 @@ rocksdb::Status Database::Del(const Slice &user_key) {
   Metadata metadata(kRedisNone, false);
   metadata.Decode(value);
   if (metadata.Expired()) {
-    return rocksdb::Status::NotFound("the key was expired");
+    return rocksdb::Status::NotFound(kErrMsgKeyExpired);
   }
   return storage_->Delete(rocksdb::WriteOptions(), metadata_cf_handle_, ns_key);
 }

diff --git a/src/redis_list.cc b/src/redis_list.cc
@@ -399,7 +399,7 @@ rocksdb::Status List::RPopLPush(const Slice &src, const Slice &dst, std::string
   rocksdb::Status s = Type(dst, &type);
   if (!s.ok()) return s;
   if (type != kRedisNone && type != kRedisList) {
-    return rocksdb::Status::InvalidArgument("WRONGTYPE Operation against a key holding the wrong kind of value");
+    return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
   }
 
   s = Pop(src, elem, false);

diff --git a/src/redis_metadata.h b/src/redis_metadata.h
@@ -36,6 +36,9 @@ const std::vector<std::string> RedisTypeNames = {
     "list", "set", "zset", "bitmap", "sortedint"
 };
 
+const char kErrMsgWrongType[] = "WRONGTYPE Operation against a key holding the wrong kind of value";
+const char kErrMsgKeyExpired[] = "the key was expired";
+
 using rocksdb::Slice;
 
 struct KeyNumStats {

diff --git a/src/redis_reply.cc b/src/redis_reply.cc
@@ -39,7 +39,7 @@ std::string MultiBulkString(std::vector<std::string> values, bool output_nil_for
 
 std::string MultiBulkString(std::vector<std::string> values, const std::vector<rocksdb::Status> &statuses) {
   for (size_t i = 0; i < values.size(); i++) {
-    if (i < statuses.size() && statuses[i].IsNotFound()) {
+    if (i < statuses.size() && !statuses[i].ok()) {
       values[i] = NilString();
     } else {
       values[i] = BulkString(values[i]);

diff --git a/src/redis_set.cc b/src/redis_set.cc
@@ -209,7 +209,7 @@ rocksdb::Status Set::Move(const Slice &src, const Slice &dst, const Slice &membe
   rocksdb::Status s = Type(dst, &type);
   if (!s.ok()) return s;
   if (type != kRedisNone && type != kRedisSet) {
-    return rocksdb::Status::InvalidArgument("WRONGTYPE Operation against a key holding the wrong kind of value");
+    return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
   }
 
   std::vector<Slice> members{member};

diff --git a/src/redis_string.cc b/src/redis_string.cc
@@ -6,6 +6,34 @@
 
 namespace Redis {
 
+std::vector<rocksdb::Status> String::getRawValues(
+    const std::vector<Slice> &keys,
+    std::vector<std::string> *raw_values) {
+  raw_values->clear();
+
+  rocksdb::ReadOptions read_options;
+  LatestSnapShot ss(db_);
+  read_options.snapshot = ss.GetSnapShot();
+  std::vector<rocksdb::ColumnFamilyHandle*> cfs(keys.size(), metadata_cf_handle_);
+  auto statuses = db_->MultiGet(read_options, cfs, keys, raw_values);
+  for (size_t i = 0; i < keys.size(); i++) {
+    if (!statuses[i].ok()) continue;
+    Metadata metadata(kRedisNone, false);
+    metadata.Decode((*raw_values)[i]);
+    if (metadata.Expired()) {
+      (*raw_values)[i].clear();
+      statuses[i] = rocksdb::Status::NotFound(kErrMsgKeyExpired);
+      continue;
+    }
+    if (metadata.Type() != kRedisString && metadata.size > 0) {
+      (*raw_values)[i].clear();
+      statuses[i] = rocksdb::Status::InvalidArgument(kErrMsgWrongType);
+      continue;
+    }
+  }
+  return statuses;
+}
+
 rocksdb::Status String::getRawValue(const std::string &ns_key, std::string *raw_value) {
   raw_value->clear();
 
@@ -19,23 +47,33 @@ rocksdb::Status String::getRawValue(const std::string &ns_key, std::string *raw_
   metadata.Decode(*raw_value);
   if (metadata.Expired()) {
     raw_value->clear();
-    return rocksdb::Status::NotFound("the key was expired");
+    return rocksdb::Status::NotFound(kErrMsgKeyExpired);
   }
   if (metadata.Type() != kRedisString && metadata.size > 0) {
-    return rocksdb::Status::InvalidArgument("WRONGTYPE Operation against a key holding the wrong kind of value");
+    return rocksdb::Status::InvalidArgument(kErrMsgWrongType);
   }
   return rocksdb::Status::OK();
 }
 
 rocksdb::Status String::getValue(const std::string &ns_key, std::string *value) {
   value->clear();
+
   std::string raw_value;
   auto s = getRawValue(ns_key, &raw_value);
   if (!s.ok()) return s;
   *value = raw_value.substr(STRING_HDR_SIZE, raw_value.size()-STRING_HDR_SIZE);
   return rocksdb::Status::OK();
 }
 
+std::vector<rocksdb::Status> String::getValues(const std::vector<Slice> &ns_keys, std::vector<std::string> *values) {
+  auto statuses = getRawValues(ns_keys, values);
+  for (size_t i = 0; i < ns_keys.size(); i++) {
+    if (!statuses[i].ok()) continue;
+    (*values)[i] = (*values)[i].substr(STRING_HDR_SIZE, (*values)[i].size()-STRING_HDR_SIZE);
+  }
+  return statuses;
+}
+
 rocksdb::Status String::updateRawValue(const std::string &ns_key, const std::string &raw_value) {
   rocksdb::WriteBatch batch;
   WriteBatchLogData log_data(kRedisString);
@@ -63,15 +101,22 @@ rocksdb::Status String::Append(const std::string &user_key, const std::string &v
 }
 
 std::vector<rocksdb::Status> String::MGet(const std::vector<Slice> &keys, std::vector<std::string> *values) {
-  std::string ns_key;
-  std::string value;
-  std::vector<rocksdb::Status> statuses;
-  for (size_t i = 0; i < keys.size(); i++) {
-    AppendNamespacePrefix(keys[i], &ns_key);
-    statuses.emplace_back(getValue(ns_key, &value));
-    values->emplace_back(value);
+  std::vector<std::string> ns_keys;
+  ns_keys.reserve(keys.size());
+  for (const auto &key : keys) {
+    std::string ns_key;
+    AppendNamespacePrefix(key, &ns_key);
+    ns_keys.emplace_back(ns_key);
   }
-  return statuses;
+  std::vector<Slice> slice_keys;
+  slice_keys.reserve(ns_keys.size());
+  // don't use range-based for loop here, coz the slice member
+  // would refer the address instead of copy the value, and use
+  // range-based for loop may cause all members refer to the same addr
+  for (size_t i = 0; i < ns_keys.size(); i++) {
+    slice_keys.emplace_back(ns_keys[i]);
+  }
+  return getValues(slice_keys, values);
 }
 
 rocksdb::Status String::Get(const std::string &user_key, std::string *value) {

diff --git a/src/redis_string.h b/src/redis_string.h
@@ -34,7 +34,9 @@ class String : public Database {
 
  private:
   rocksdb::Status getValue(const std::string &ns_key, std::string *value);
+  std::vector<rocksdb::Status> getValues(const std::vector<Slice> &ns_keys, std::vector<std::string> *values);
   rocksdb::Status getRawValue(const std::string &ns_key, std::string *raw_value);
+  std::vector<rocksdb::Status> getRawValues(const std::vector<Slice> &keys, std::vector<std::string> *raw_values);
   rocksdb::Status updateRawValue(const std::string &ns_key, const std::string &raw_value);
 };