From 0db38d46239fbaffd89847fc4838dec714989167 Mon Sep 17 00:00:00 2001
From: Roman Gershman <roman@dragonflydb.io>
Date: Thu, 24 Oct 2024 19:35:03 +0300
Subject: [PATCH] fix: properly set ttl bit during object replacement

fixes #3984

Signed-off-by: Roman Gershman <roman@dragonflydb.io>
---
 src/core/dense_set.cc       | 11 ++++++++---
 src/core/string_map_test.cc | 16 ++++++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/core/dense_set.cc b/src/core/dense_set.cc
index c4be03b24429..f03d2cbabb00 100644
--- a/src/core/dense_set.cc
+++ b/src/core/dense_set.cc
@@ -52,6 +52,8 @@ void DenseSet::IteratorBase::SetExpiryTime(uint32_t ttl_sec) {
   if (!HasExpiry()) {
     void* new_obj = owner_->ObjectClone(src, false, true);
     ptr->SetObject(new_obj);
+
+    // Important: we set the ttl bit on the wrapping pointer.
     curr_entry_->SetTtl(true);
     owner_->ObjDelete(src, false);
     src = new_obj;
@@ -678,8 +680,12 @@ 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;
 
-  if (dptr) {  // replace
-    if (dptr->IsLink())
+  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.
+    dptr->SetTtl(has_ttl);
+
+    if (dptr->IsLink())  // unwrap the pointer.
       dptr = dptr->AsLink();
 
     void* res = dptr->Raw();
@@ -687,7 +693,6 @@ void* DenseSet::AddOrReplaceObj(void* obj, bool has_ttl) {
     obj_malloc_used_ += ObjectAllocSize(obj);
 
     dptr->SetObject(obj);
-    dptr->SetTtl(has_ttl);
 
     return res;
   }
diff --git a/src/core/string_map_test.cc b/src/core/string_map_test.cc
index 3feaef03d4d1..4f59fa0d3fe8 100644
--- a/src/core/string_map_test.cc
+++ b/src/core/string_map_test.cc
@@ -168,6 +168,22 @@ TEST_F(StringMapTest, Bug3973) {
   }
 }
 
+TEST_F(StringMapTest, Bug3984) {
+  for (unsigned i = 0; i < 6; i++) {
+    EXPECT_TRUE(sm_->AddOrUpdate(to_string(i), "val"));
+  }
+  for (unsigned i = 0; i < 6; i++) {
+    auto k = sm_->Find(to_string(i));
+    ASSERT_FALSE(k.HasExpiry());
+    k.SetExpiryTime(1);
+    EXPECT_EQ(k.ExpiryTime(), 1);
+  }
+
+  for (unsigned i = 0; i < 6; i++) {
+    EXPECT_FALSE(sm_->AddOrUpdate(to_string(i), "val"));
+  }
+}
+
 unsigned total_wasted_memory = 0;
 
 TEST_F(StringMapTest, ReallocIfNeeded) {