fix(server):set of existing object with expiry (issue#607) (#609)

fix(server):set of existing object with expriy (issue#607)

Signed-off-by: ashotland <ari@dragonflydb.io>

Signed-off-by: ashotland <ari@dragonflydb.io>

src/server/db_slice.cc

@@ -506,21 +506,29 @@ void DbSlice::FlushDb(DbIndex db_ind) {
   }).detach();
 }
 
+void DbSlice::AddExpire(DbIndex db_ind, PrimeIterator main_it, uint64_t at) {
+  uint64_t delta = at - expire_base_[0];  // TODO: employ multigen expire updates.
+  CHECK(db_arr_[db_ind]->expire.Insert(main_it->first.AsRef(), ExpirePeriod(delta)).second);
+  main_it->second.SetExpire(true);
+}
+
+bool DbSlice::RemoveExpire(DbIndex db_ind, PrimeIterator main_it) {
+  if (main_it->second.HasExpire()) {
+    CHECK_EQ(1u, db_arr_[db_ind]->expire.Erase(main_it->first));
+    main_it->second.SetExpire(false);
+    return true;
+  }
+  return false;
+}
+
 // Returns true if a state has changed, false otherwise.
 bool DbSlice::UpdateExpire(DbIndex db_ind, PrimeIterator it, uint64_t at) {
-  auto& db = *db_arr_[db_ind];
-  if (at == 0 && it->second.HasExpire()) {
-    CHECK_EQ(1u, db.expire.Erase(it->first));
-    it->second.SetExpire(false);
-
-    return true;
+  if (at == 0) {
+    return RemoveExpire(db_ind, it);
   }
 
   if (!it->second.HasExpire() && at) {
-    uint64_t delta = at - expire_base_[0];  // TODO: employ multigen expire updates.
-    CHECK(db.expire.Insert(it->first.AsRef(), ExpirePeriod(delta)).second);
-    it->second.SetExpire(true);
-
+    AddExpire(db_ind, it, at);
     return true;
   }
 

src/server/db_slice.h

@@ -174,6 +174,13 @@ class DbSlice {
   facade::OpStatus UpdateExpire(const Context& cntx, PrimeIterator prime_it, ExpireIterator exp_it,
                                 const ExpireParams& params);
 
+  // Adds expiry information.
+  void AddExpire(DbIndex db_ind, PrimeIterator main_it, uint64_t at);
+
+  // Removes the corresponing expiry information if exists.
+  // Returns true if expiry existed (and removed).
+  bool RemoveExpire(DbIndex db_ind, PrimeIterator main_it);
+
   // Either adds or removes (if at == 0) expiry. Returns true if a change was made.
   // Does not change expiry if at != 0 and expiry already exists.
   bool UpdateExpire(DbIndex db_ind, PrimeIterator main_it, uint64_t at);

src/server/dragonfly_test.cc

@@ -772,6 +772,21 @@ TEST_F(DflyEngineTest, Bug496) {
   });
 }
 
+TEST_F(DflyEngineTest, Issue706) {
+  // https://github.com/dragonflydb/dragonfly/issues/607
+
+  Run({"SET", "key", "value1"});
+  EXPECT_EQ(Run({"GET", "key"}), "value1");
+
+  Run({"SET", "key", "value2"});
+  EXPECT_EQ(Run({"GET", "key"}), "value2");
+
+  Run({"EXPIRE", "key", "1000"});
+
+  Run({"SET", "key", "value3"});
+  EXPECT_EQ(Run({"GET", "key"}), "value3");
+}
+
 TEST_F(DefragDflyEngineTest, TestDefragOption) {
   absl::SetFlag(&FLAGS_commit_use_threshold, 1.1);
   // Fill data into dragonfly and then check if we have

src/server/string_family.cc

@@ -313,16 +313,17 @@ OpResult<int64_t> OpIncrBy(const OpArgs& op_args, string_view key, int64_t incr,
   return new_val;
 }
 
-int64_t CalculateAbsTime(int64_t unix_time, bool as_milli) {
+int64_t AbsExpiryToTtl(int64_t abs_expiry_time, bool as_milli) {
   using std::chrono::duration_cast;
   using std::chrono::milliseconds;
   using std::chrono::seconds;
   using std::chrono::system_clock;
 
   if (as_milli) {
-    return unix_time - duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
+    return abs_expiry_time -
+           duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
   } else {
-    return unix_time - duration_cast<seconds>(system_clock::now().time_since_epoch()).count();
+    return abs_expiry_time - duration_cast<seconds>(system_clock::now().time_since_epoch()).count();
   }
 }
 
@@ -438,13 +439,19 @@ OpStatus SetCmd::SetExisting(const SetParams& params, PrimeIterator it, ExpireIt
   DbSlice& db_slice = shard->db_slice();
   uint64_t at_ms =
       params.expire_after_ms ? params.expire_after_ms + op_args_.db_cntx.time_now_ms : 0;
-  if (IsValid(e_it) && at_ms) {
-    e_it->second = db_slice.FromAbsoluteTime(at_ms);
-  } else if (!(params.flags & SET_KEEP_EXPIRE)) {
-    // We need to update expiry, or maybe erase the object if it was expired.
-    bool changed = db_slice.UpdateExpire(op_args_.db_cntx.db_index, it, at_ms);
-    if (changed && at_ms == 0)  // erased.
-      return OpStatus::OK;      // TODO: to update journal with deletion.
+
+  if (!(params.flags & SET_KEEP_EXPIRE)) {
+    if (at_ms) {  // Command has an expiry paramater.
+      if (IsValid(e_it)) {
+        // Updated exisitng expiry information.
+        e_it->second = db_slice.FromAbsoluteTime(at_ms);
+      } else {
+        // Add new expiry information.
+        db_slice.AddExpire(op_args_.db_cntx.db_index, it, at_ms);
+      }
+    } else {
+      db_slice.RemoveExpire(op_args_.db_cntx.db_index, it);
+    }
   }
 
   db_slice.PreUpdate(op_args_.db_cntx.db_index, it);
@@ -513,7 +520,7 @@ void StringFamily::Set(CmdArgList args, ConnectionContext* cntx) {
       // check here and if the time is in the past, return OK but don't set it
       // Note that the time pass here for PXAT is in milliseconds, we must not change it!
       if (cur_arg == "EXAT" || cur_arg == "PXAT") {
-        int_arg = CalculateAbsTime(int_arg, is_ms);
+        int_arg = AbsExpiryToTtl(int_arg, is_ms);
         if (int_arg < 0) {
           // this happened in the past, just return, for some reason Redis reports OK in this case
           return builder->SendStored();