Avoid to fetch old value in SET command if NX/XX/GET/KEEPTTL is not s…

…et (#1968)
apache · Dec 25, 2023 · 3d83f2e · 3d83f2e
1 parent b0b1b4e
commit 3d83f2e
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 38 deletions.
diff --git a/src/types/redis_string.cc b/src/types/redis_string.cc
@@ -89,17 +89,29 @@ rocksdb::Status String::getRawValue(const std::string &ns_key, std::string *raw_
   return rocksdb::Status::OK();
 }
 
-rocksdb::Status String::getValue(const std::string &ns_key, std::string *value) {
+rocksdb::Status String::getValueAndExpire(const std::string &ns_key, std::string *value, uint64_t *expire) {
   value->clear();
 
   std::string raw_value;
   auto s = getRawValue(ns_key, &raw_value);
   if (!s.ok()) return s;
+
   size_t offset = Metadata::GetOffsetAfterExpire(raw_value[0]);
   *value = raw_value.substr(offset);
+
+  if (expire) {
+    Metadata metadata(kRedisString, false);
+    s = metadata.Decode(raw_value);
+    if (!s.ok()) return s;
+    *expire = metadata.expire;
+  }
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status String::getValue(const std::string &ns_key, std::string *value) {
+  return getValueAndExpire(ns_key, value, nullptr);
+}
+
 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++) {
@@ -208,57 +220,55 @@ rocksdb::Status String::Set(const std::string &user_key, const std::string &valu
 
 rocksdb::Status String::Set(const std::string &user_key, const std::string &value, StringSetArgs args,
                             std::optional<std::string> &ret) {
+  uint64_t expire = 0;
   std::string ns_key = AppendNamespacePrefix(user_key);
 
   LockGuard guard(storage_->GetLockManager(), ns_key);
-
-  // Get old value for NX/XX/GET/KEEPTTL option
-  std::string old_raw_value;
-  auto s = getRawValue(ns_key, &old_raw_value);
-  if (!s.ok() && !s.IsNotFound() && !s.IsInvalidArgument()) return s;
-  auto old_key_found = !s.IsNotFound();
-  // The reply following Redis doc: https://redis.io/commands/set/
-  // Handle GET option
-  if (args.get) {
-    if (s.IsInvalidArgument()) {
-      return s;
+  bool need_old_value = args.type != StringSetType::NONE || args.get || args.keep_ttl;
+  if (need_old_value) {
+    std::string old_value;
+    uint64_t old_expire = 0;
+    auto s = getValueAndExpire(ns_key, &old_value, &old_expire);
+    if (!s.ok() && !s.IsNotFound() && !s.IsInvalidArgument()) return s;
+    // GET option
+    if (args.get) {
+      if (s.IsInvalidArgument()) {
+        return s;
+      }
+      if (s.IsNotFound()) {
+        // if GET option given, the key didn't exist before: return nil
+        ret = std::nullopt;
+      } else {
+        // if GET option given: return The previous value of the key.
+        ret = std::make_optional(old_value);
+      }
     }
-    if (old_key_found) {
-      // if GET option given: return The previous value of the key.
-      auto offset = Metadata::GetOffsetAfterExpire(old_raw_value[0]);
-      ret = std::make_optional(old_raw_value.substr(offset));
+    // NX/XX option
+    if (args.type == StringSetType::NX && s.ok()) {
+      // if NX option given, the key already exist: return nil
+      if (!args.get) ret = std::nullopt;
+      return rocksdb::Status::OK();
+    } else if (args.type == StringSetType::XX && s.IsNotFound()) {
+      // if XX option given, the key didn't exist before: return nil
+      if (!args.get) ret = std::nullopt;
+      return rocksdb::Status::OK();
     } else {
-      // if GET option given, the key didn't exist before: return nil
-      ret = std::nullopt;
+      // if GET option not given, make ret not nil
+      if (!args.get) ret = "";
+    }
+    if (s.ok() && args.keep_ttl) {
+      // if KEEPTTL option given, use the old ttl
+      expire = old_expire;
     }
-  }
-
-  // Handle NX/XX option
-  if (old_key_found && args.type == StringSetType::NX) {
-    // if GET option not given, operation aborted: return nil
-    if (!args.get) ret = std::nullopt;
-    return rocksdb::Status::OK();
-  } else if (!old_key_found && args.type == StringSetType::XX) {
-    // if GET option not given, operation aborted: return nil
-    if (!args.get) ret = std::nullopt;
-    return rocksdb::Status::OK();
   } else {
-    // if GET option not given, make ret not nil
+    // if no option given, make ret not nil
     if (!args.get) ret = "";
   }
 
   // Handle expire time
-  uint64_t expire = 0;
   if (args.ttl > 0) {
     uint64_t now = util::GetTimeStampMS();
     expire = now + args.ttl;
-  } else if (args.keep_ttl && old_key_found) {
-    Metadata metadata(kRedisString, false);
-    auto s = metadata.Decode(old_raw_value);
-    if (!s.ok()) {
-      return s;
-    }
-    expire = metadata.expire;
   }
 
   // Create new value

diff --git a/src/types/redis_string.h b/src/types/redis_string.h
@@ -71,6 +71,7 @@ class String : public Database {
 
  private:
   rocksdb::Status getValue(const std::string &ns_key, std::string *value);
+  rocksdb::Status getValueAndExpire(const std::string &ns_key, std::string *value, uint64_t *expire);
   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);