feat(hset_family): Add KEEPTTL support to HSetEx (#4730)

* feat(hset_family): Add support for KEEPTTL to HSETEX

The KEEPTTL option if specified makes sure that TTL is preserved for existing members.

Signed-off-by: Abhijat Malviya <[email protected]>

Author: abhijat

src/core/dense_set.cc

@@ -690,8 +690,8 @@ void* DenseSet::PopInternal() {
 
 void* DenseSet::AddOrReplaceObj(void* obj, bool has_ttl) {
   uint64_t hc = Hash(obj, 0);
-  DensePtr* dptr = entries_.empty() ? nullptr : Find(obj, BucketId(hc), 0).second;
 
+  DensePtr* dptr = entries_.empty() ? nullptr : Find(obj, BucketId(hc), 0).second;
   if (dptr) {  // replace existing object.
     // A bit confusing design: ttl bit is located on the wrapping pointer,
     // therefore we must set ttl bit before unrapping below.

src/core/string_map.cc

@@ -65,13 +65,18 @@ StringMap::~StringMap() {
   Clear();
 }
 
-bool StringMap::AddOrUpdate(string_view field, string_view value, uint32_t ttl_sec) {
+bool StringMap::AddOrUpdate(std::string_view field, std::string_view value, uint32_t ttl_sec,
+                            bool keepttl) {
   // 8 additional bytes for a pointer to value.
   auto [newkey, sdsval_tag] = CreateEntry(field, value, time_now(), ttl_sec);
 
   // Replace the whole entry.
-  sds prev_entry = (sds)AddOrReplaceObj(newkey, sdsval_tag & kValTtlBit);
-  if (prev_entry) {
+  if (sds prev_entry = (sds)AddOrReplaceObj(newkey, sdsval_tag & kValTtlBit); prev_entry) {
+    const bool prev_has_ttl =
+        absl::little_endian::Load64(prev_entry + sdslen(prev_entry) + 1) & kValTtlBit;
+    if (keepttl && prev_has_ttl) {
+      SdsUpdateExpireTime(newkey, ObjExpireTime(prev_entry), 8);
+    }
     ObjDelete(prev_entry, false);
     return false;
   }

src/core/string_map.h

@@ -104,7 +104,8 @@ class StringMap : public DenseSet {
   };
 
   // otherwise updates its value and returns false.
-  bool AddOrUpdate(std::string_view field, std::string_view value, uint32_t ttl_sec = UINT32_MAX);
+  bool AddOrUpdate(std::string_view field, std::string_view value, uint32_t ttl_sec = UINT32_MAX,
+                   bool keepttl = false);
 
   // Returns true if field was added
   // false, if already exists. In that case no update is done.

src/server/hset_family.cc

@@ -607,6 +607,7 @@ OpResult<size_t> OpStrLen(const OpArgs& op_args, string_view key, string_view fi
 struct OpSetParams {
   bool skip_if_exists = false;
   uint32_t ttl = UINT32_MAX;
+  bool keepttl = false;
 };
 
 OpResult<uint32_t> OpSet(const OpArgs& op_args, string_view key, CmdArgList values,
@@ -673,7 +674,7 @@ OpResult<uint32_t> OpSet(const OpArgs& op_args, string_view key, CmdArgList valu
       if (op_sp.skip_if_exists)
         added = sm->AddOrSkip(field, value, op_sp.ttl);
       else
-        added = sm->AddOrUpdate(field, value, op_sp.ttl);
+        added = sm->AddOrUpdate(field, value, op_sp.ttl, op_sp.keepttl);
 
       created += unsigned(added);
     }
@@ -720,19 +721,23 @@ OpResult<vector<long>> OpHExpire(const OpArgs& op_args, string_view key, uint32_
   return HSetFamily::SetFieldsExpireTime(op_args, ttl_sec, key, values, pv);
 }
 
-// HSETEX key [NX] tll_sec field value field value ...
+// HSETEX key [NX] [KEEPTTL] tll_sec field value field value ...
 void HSetEx(CmdArgList args, const CommandContext& cmd_cntx) {
   CmdArgParser parser{args};
 
   string_view key = parser.Next();
+  OpSetParams op_sp;
 
-  bool skip_if_exists = static_cast<bool>(parser.Check("NX"sv));
-  string_view ttl_str = parser.Next();
+  op_sp.skip_if_exists = parser.Check("NX");
+  op_sp.keepttl = parser.Check("KEEPTTL");
+  op_sp.ttl = parser.Next<uint32_t>();
 
-  uint32_t ttl_sec;
-  constexpr uint32_t kMaxTtl = (1UL << 26);
+  if (parser.HasError()) {
+    return cmd_cntx.rb->SendError(parser.Error()->MakeReply());
+  }
 
-  if (!absl::SimpleAtoi(ttl_str, &ttl_sec) || ttl_sec == 0 || ttl_sec > kMaxTtl) {
+  constexpr uint32_t kMaxTtl = (1UL << 26);
+  if (op_sp.ttl == 0 || op_sp.ttl > kMaxTtl) {
     return cmd_cntx.rb->SendError(kInvalidIntErr);
   }
 
@@ -743,8 +748,6 @@ void HSetEx(CmdArgList args, const CommandContext& cmd_cntx) {
                                   kSyntaxErrType);
   }
 
-  OpSetParams op_sp{skip_if_exists, ttl_sec};
-
   auto cb = [&](Transaction* t, EngineShard* shard) {
     return OpSet(t->GetOpArgs(shard), key, fields, op_sp);
   };

src/server/hset_family_test.cc

@@ -364,6 +364,47 @@ TEST_F(HSetFamilyTest, HSetEx) {
   resp = Run({"HGET", "k", "field4"});
   EXPECT_THAT(resp,
               "value");  // HSETEX with NX option; old expiration time was NOT replaced by a new one
+
+  // KEEPTTL related asserts
+  EXPECT_THAT(Run({"HSETEX", "k", long_time, "kttlfield", "value"}), IntArg(1));
+  EXPECT_EQ(Run({"HGET", "k", "kttlfield"}), "value");
+  EXPECT_EQ(CheckedInt({"FIELDTTL", "k", "kttlfield"}), 100);
+
+  // KEEPTTL resets value of kttlfield, but preserves its TTL. afield is added with TTL=1
+  EXPECT_THAT(Run({"HSETEX", "k", "KEEPTTL", "1", "kttlfield", "resetvalue", "afield", "aval"}),
+              IntArg(1));
+  EXPECT_EQ(CheckedInt({"FIELDTTL", "k", "kttlfield"}), 100);
+  EXPECT_EQ(Run({"FIELDTTL", "k", "afield"}).GetInt(), 1);
+  EXPECT_EQ(Run({"HGET", "k", "afield"}), "aval");
+  // make afield expire
+  AdvanceTime(1000);
+  EXPECT_THAT(Run({"HGET", "k", "afield"}), ArgType(RespExpr::NIL));
+
+  // kttlfield is still present although with updated value
+  EXPECT_EQ(Run({"HGET", "k", "kttlfield"}), "resetvalue");
+  EXPECT_EQ(Run({"FIELDTTL", "k", "kttlfield"}).GetInt(), 99);
+
+  // If NX is supplied, with or without KEEPTTL neither expiry nor value is updated
+  EXPECT_THAT(Run({"HSETEX", "k", "NX", "KEEPTTL", "1", "kttlfield", "value"}), IntArg(0));
+
+  // No updates
+  EXPECT_EQ(Run({"HGET", "k", "kttlfield"}), "resetvalue");
+  EXPECT_EQ(Run({"FIELDTTL", "k", "kttlfield"}).GetInt(), 99);
+
+  EXPECT_THAT(Run({"HSETEX", "k", "NX", "1", "kttlfield", "value"}), IntArg(0));
+  // No updates
+  EXPECT_EQ(Run({"HGET", "k", "kttlfield"}), "resetvalue");
+  EXPECT_EQ(Run({"FIELDTTL", "k", "kttlfield"}).GetInt(), 99);
+
+  // Invalid TTL handling
+  EXPECT_THAT(Run({"HSETEX", "k", "NX", "zero", "kttlfield", "value"}),
+              ErrArg("ERR value is not an integer or out of range"));
+
+  // Exercise the code path where a field is added without TTL, but then we set a new expiration AND
+  // provide KEEPTTL. Since there was no old expiry, the new TTL should be applied.
+  EXPECT_EQ(Run({"HSET", "k", "nottl", "val"}), 1);
+  EXPECT_EQ(Run({"HSETEX", "k", "KEEPTTL", long_time, "nottl", "newval"}), 0);
+  EXPECT_EQ(Run({"FIELDTTL", "k", "nottl"}).GetInt(), 100);
 }
 
 TEST_F(HSetFamilyTest, TriggerConvertToStrMap) {