From 125add69dc49d8b8f80199537138beea14d19bf8 Mon Sep 17 00:00:00 2001
From: NegatioN <4037769+NegatioN@users.noreply.github.com>
Date: Mon, 23 Sep 2024 21:05:48 +0200
Subject: [PATCH] add dense_set.SetExpiryTime in preparation for fieldexpire

---
 src/core/dense_set.cc       | 26 ++++++++++++++
 src/core/dense_set.h        | 15 ++++++++
 src/core/score_map.cc       | 10 ++++++
 src/core/score_map.h        |  2 ++
 src/core/string_map.cc      | 29 +++++++++++++++
 src/core/string_map.h       |  5 ++-
 src/core/string_map_test.cc | 19 ++++++++++
 src/core/string_set.cc      | 71 ++++++++++++++++++++++++++++---------
 src/core/string_set.h       |  5 +++
 src/core/string_set_test.cc | 19 ++++++++++
 10 files changed, 183 insertions(+), 18 deletions(-)

diff --git a/src/core/dense_set.cc b/src/core/dense_set.cc
index 05df23938b15..3525f547b3f6 100644
--- a/src/core/dense_set.cc
+++ b/src/core/dense_set.cc
@@ -337,6 +337,13 @@ void DenseSet::Grow(size_t prev_size) {
   }
 }
 
+auto DenseSet::FindDense(void* ptr) -> DensePtr* {
+  uint64_t hc = Hash(ptr, 0);
+  uint32_t bucket_id = BucketId(hc);
+  DensePtr* dptr = Find(ptr, bucket_id, 0).second;
+  return dptr;
+}
+
 auto DenseSet::AddOrFindDense(void* ptr, bool has_ttl) -> DensePtr* {
   uint64_t hc = Hash(ptr, 0);
 
@@ -544,6 +551,25 @@ void* DenseSet::PopInternal() {
   return ret;
 }
 
+void* DenseSet::ReplaceObj(void* obj, bool has_ttl) {
+  DensePtr* ptr = FindDense(obj);
+  if (!ptr)
+    return nullptr;
+
+  if (ptr->IsLink()) {
+    ptr = ptr->AsLink();
+  }
+
+  void* res = ptr->Raw();
+  obj_malloc_used_ -= ObjectAllocSize(res);
+  obj_malloc_used_ += ObjectAllocSize(obj);
+
+  ptr->SetObject(obj);
+  ptr->SetTtl(has_ttl);
+
+  return res;
+}
+
 void* DenseSet::AddOrReplaceObj(void* obj, bool has_ttl) {
   DensePtr* ptr = AddOrFindDense(obj, has_ttl);
   if (!ptr)
diff --git a/src/core/dense_set.h b/src/core/dense_set.h
index 1486c16bba56..77f1c9e28fd4 100644
--- a/src/core/dense_set.h
+++ b/src/core/dense_set.h
@@ -187,6 +187,14 @@ class DenseSet {
       return curr_entry_->HasTtl() ? owner_->ObjExpireTime(curr_entry_->GetObject()) : UINT32_MAX;
     }
 
+    void SetExpiryTime(uint32_t ttl_sec) {
+      if (HasExpiry()) {
+        owner_->ObjUpdateExpireTime(curr_entry_->GetObject(), ttl_sec);
+      } else {
+        owner_->ObjReplace(curr_entry_->GetObject(), ttl_sec);
+      }
+    }
+
     bool HasExpiry() const {
       return curr_entry_->HasTtl();
     }
@@ -263,6 +271,8 @@ class DenseSet {
   virtual bool ObjEqual(const void* left, const void* right, uint32_t right_cookie) const = 0;
   virtual size_t ObjectAllocSize(const void* obj) const = 0;
   virtual uint32_t ObjExpireTime(const void* obj) const = 0;
+  virtual uint32_t ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) = 0;
+  virtual uint32_t ObjReplace(const void* obj, uint32_t ttl_sec) = 0;
   virtual void ObjDelete(void* obj, bool has_ttl) const = 0;
 
   void CollectExpired();
@@ -314,6 +324,7 @@ class DenseSet {
   // Returns the previous object if it has been replaced.
   // nullptr, if obj was added.
   void* AddOrReplaceObj(void* obj, bool has_ttl);
+  void* ReplaceObj(void* obj, bool has_ttl);
 
   // Assumes that the object does not exist in the set.
   void AddUnique(void* obj, bool has_ttl, uint64_t hashcode);
@@ -356,6 +367,10 @@ class DenseSet {
   // Returns null if obj was added.
   DensePtr* AddOrFindDense(void* obj, bool has_ttl);
 
+  // Returns DensePtr if the object with such key exists,
+  // Returns null if key does not exist.
+  DensePtr* FindDense(void* obj);
+
   // ============ Pseudo Linked List in DenseSet end ==================
 
   // returns (prev, item) pair. If item is root, then prev is null.
diff --git a/src/core/score_map.cc b/src/core/score_map.cc
index 505e31f8aa5e..fb44d63bbd08 100644
--- a/src/core/score_map.cc
+++ b/src/core/score_map.cc
@@ -128,6 +128,16 @@ uint32_t ScoreMap::ObjExpireTime(const void* obj) const {
   return UINT32_MAX;
 }
 
+uint32_t ScoreMap::ObjReplace(const void* obj, uint32_t ttl_sec) {
+  // Should not reach.
+  return UINT32_MAX;
+}
+
+uint32_t ScoreMap::ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) {
+  // Should not reach.
+  return UINT32_MAX;
+}
+
 void ScoreMap::ObjDelete(void* obj, bool has_ttl) const {
   sds s1 = (sds)obj;
   sdsfree(s1);
diff --git a/src/core/score_map.h b/src/core/score_map.h
index 36b80d5ffbc9..b97adc2797a8 100644
--- a/src/core/score_map.h
+++ b/src/core/score_map.h
@@ -122,6 +122,8 @@ class ScoreMap : public DenseSet {
   bool ObjEqual(const void* left, const void* right, uint32_t right_cookie) const final;
   size_t ObjectAllocSize(const void* obj) const final;
   uint32_t ObjExpireTime(const void* obj) const final;
+  uint32_t ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) override;
+  uint32_t ObjReplace(const void* obj, uint32_t ttl_sec) override;
   void ObjDelete(void* obj, bool has_ttl) const final;
 };
 
diff --git a/src/core/string_map.cc b/src/core/string_map.cc
index ef7e12b49d28..c9c69e8507b7 100644
--- a/src/core/string_map.cc
+++ b/src/core/string_map.cc
@@ -276,6 +276,35 @@ uint32_t StringMap::ObjExpireTime(const void* obj) const {
   return UINT32_MAX;
 }
 
+uint32_t StringMap::ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) {
+  sds str = (sds)obj;
+  char* valptr = str + sdslen(str) + 1;
+
+  uint32_t at = time_now() + ttl_sec;
+  uint64_t val = absl::little_endian::Load64(valptr);
+
+  absl::little_endian::Store32(valptr + 8, at);
+  // TODO
+  return 1;
+}
+
+uint32_t StringMap::ObjReplace(const void* obj, uint32_t ttl_sec) {
+  sds str = (sds)obj;
+  char* valptr = str + sdslen(str) + 1;
+
+  uint32_t at = time_now() + ttl_sec;
+  uint64_t val = absl::little_endian::Load64(valptr);
+
+  auto pair = detail::SdsPair(str, GetValue(str));
+  auto [newkey, sdsval_tag] = CreateEntry(pair->first, pair->second, time_now(), ttl_sec);
+
+  sds prev_entry = (sds)ReplaceObj(newkey, sdsval_tag & kValTtlBit);
+  if (prev_entry) {
+    ObjDelete(prev_entry, false);
+  }
+  return 1;
+}
+
 void StringMap::ObjDelete(void* obj, bool has_ttl) const {
   sds s1 = (sds)obj;
   sds value = GetValue(s1);
diff --git a/src/core/string_map.h b/src/core/string_map.h
index 018d206be3e2..1961a3175e05 100644
--- a/src/core/string_map.h
+++ b/src/core/string_map.h
@@ -100,6 +100,7 @@ class StringMap : public DenseSet {
 
     using IteratorBase::ExpiryTime;
     using IteratorBase::HasExpiry;
+    using IteratorBase::SetExpiryTime;
   };
 
   // Returns true if field was added
@@ -114,7 +115,7 @@ class StringMap : public DenseSet {
 
   bool Contains(std::string_view s1) const;
 
-  /// @brief  Returns value of the key or nullptr if key not found.
+  /// @brief  Returns value of the key or an empty iterator if key not found.
   /// @param key
   /// @return sds
   iterator Find(std::string_view member) {
@@ -157,6 +158,8 @@ class StringMap : public DenseSet {
   bool ObjEqual(const void* left, const void* right, uint32_t right_cookie) const final;
   size_t ObjectAllocSize(const void* obj) const final;
   uint32_t ObjExpireTime(const void* obj) const final;
+  uint32_t ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) override;
+  uint32_t ObjReplace(const void* obj, uint32_t ttl_sec) override;
   void ObjDelete(void* obj, bool has_ttl) const final;
 };
 
diff --git a/src/core/string_map_test.cc b/src/core/string_map_test.cc
index 904ea3590086..e59efc1a8be8 100644
--- a/src/core/string_map_test.cc
+++ b/src/core/string_map_test.cc
@@ -127,6 +127,25 @@ TEST_F(StringMapTest, IterateExpired) {
   EXPECT_EQ(it, sm_->end());
 }
 
+TEST_F(StringMapTest, SetFieldExpireHasExpiry) {
+  EXPECT_TRUE(sm_->AddOrUpdate("k1", "v1", 5));
+  auto k = sm_->Find("k1");
+  EXPECT_TRUE(k.HasExpiry());
+  EXPECT_EQ(k.ExpiryTime(), 5);
+  k.SetExpiryTime(1);
+  EXPECT_TRUE(k.HasExpiry());
+  EXPECT_EQ(k.ExpiryTime(), 1);
+}
+
+TEST_F(StringMapTest, SetFieldExpireNoHasExpiry) {
+  EXPECT_TRUE(sm_->AddOrUpdate("k1", "v1"));
+  auto k = sm_->Find("k1");
+  EXPECT_FALSE(k.HasExpiry());
+  k.SetExpiryTime(1);
+  EXPECT_TRUE(k.HasExpiry());
+  EXPECT_EQ(k.ExpiryTime(), 1);
+}
+
 unsigned total_wasted_memory = 0;
 
 TEST_F(StringMapTest, ReallocIfNeeded) {
diff --git a/src/core/string_set.cc b/src/core/string_set.cc
index 087dfcfbce08..79363cdf11cb 100644
--- a/src/core/string_set.cc
+++ b/src/core/string_set.cc
@@ -20,6 +20,8 @@ namespace dfly {
 
 namespace {
 
+constexpr uint64_t kValTtlBit = 1ULL << 63;
+
 inline bool MayHaveTtl(sds s) {
   char* alloc_ptr = (char*)sdsAllocPtr(s);
   return sdslen(s) + 1 + 4 <= zmalloc_usable_size(alloc_ptr);
@@ -27,7 +29,7 @@ inline bool MayHaveTtl(sds s) {
 
 sds AllocImmutableWithTtl(uint32_t len, uint32_t at) {
   sds res = AllocSdsWithSpace(len, sizeof(at));
-  absl::little_endian::Store32(res + len + 1, at);
+  absl::little_endian::Store32(res + len + 1, at);  // Save TTL
 
   return res;
 }
@@ -43,22 +45,8 @@ bool StringSet::AddSds(sds s1) {
 }
 
 bool StringSet::Add(string_view src, uint32_t ttl_sec) {
-  DCHECK_GT(ttl_sec, 0u);  // ttl_sec == 0 would mean find and delete immediately
-
-  sds newsds = nullptr;
-  bool has_ttl = false;
-
-  if (ttl_sec == UINT32_MAX) {
-    newsds = sdsnewlen(src.data(), src.size());
-  } else {
-    uint32_t at = time_now() + ttl_sec;
-    DCHECK_LT(time_now(), at);
-
-    newsds = AllocImmutableWithTtl(src.size(), at);
-    if (!src.empty())
-      memcpy(newsds, src.data(), src.size());
-    has_ttl = true;
-  }
+  sds newsds = MakeSdsWithTtl(src, ttl_sec);
+  bool has_ttl = ttl_sec == UINT32_MAX ? false : true;
 
   if (AddOrFindObj(newsds, has_ttl) != nullptr) {
     sdsfree(newsds);
@@ -68,6 +56,21 @@ bool StringSet::Add(string_view src, uint32_t ttl_sec) {
   return true;
 }
 
+bool StringSet::Replace(string_view src, uint32_t ttl_sec) {
+  sds newsds = MakeSdsWithTtl(src, ttl_sec);
+  bool has_ttl = ttl_sec == UINT32_MAX ? false : true;
+
+  sds prev_entry = (sds)ReplaceObj(newsds, has_ttl);
+  if (prev_entry) {
+    sdsfree(prev_entry);
+    return true;
+  } else {
+    // No previous entry found, so we need to free the new entry.
+    sdsfree(newsds);
+    return false;
+  }
+}
+
 std::optional<std::string> StringSet::Pop() {
   sds str = (sds)PopInternal();
 
@@ -129,8 +132,42 @@ uint32_t StringSet::ObjExpireTime(const void* str) const {
   return absl::little_endian::Load32(ttlptr);
 }
 
+uint32_t StringSet::ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) {
+  DCHECK_GT(ttl_sec, 0u);  // ttl_sec == 0 would mean find and delete immediately
+  sds str = (sds)obj;
+  char* valptr = str + sdslen(str) + 1;
+
+  uint32_t at = time_now() + ttl_sec;
+  absl::little_endian::Store32(valptr, at);
+  return 1;
+}
+
+uint32_t StringSet::ObjReplace(const void* obj, uint32_t ttl_sec) {
+  sds str = (sds)obj;
+  Replace(str, ttl_sec);
+  return 1;
+}
+
 void StringSet::ObjDelete(void* obj, bool has_ttl) const {
   sdsfree((sds)obj);
 }
 
+sds StringSet::MakeSdsWithTtl(string_view src, uint32_t ttl_sec) {
+  DCHECK_GT(ttl_sec, 0u);  // ttl_sec == 0 would mean find and delete immediately
+  sds newsds = nullptr;
+
+  if (ttl_sec == UINT32_MAX) {
+    newsds = sdsnewlen(src.data(), src.size());
+  } else {
+    uint32_t at = time_now() + ttl_sec;
+    DCHECK_LT(time_now(), at);
+
+    newsds = AllocImmutableWithTtl(src.size(), at);
+    if (!src.empty())
+      memcpy(newsds, src.data(), src.size());
+  }
+
+  return newsds;
+}
+
 }  // namespace dfly
diff --git a/src/core/string_set.h b/src/core/string_set.h
index 93fbcfa783d4..661f2f679073 100644
--- a/src/core/string_set.h
+++ b/src/core/string_set.h
@@ -27,6 +27,7 @@ class StringSet : public DenseSet {
 
   // Returns true if elem was added.
   bool Add(std::string_view s1, uint32_t ttl_sec = UINT32_MAX);
+  bool Replace(std::string_view s1, uint32_t ttl_sec = UINT32_MAX);
 
   // Used currently by rdb_load. Returns true if elem was added.
   bool AddSds(sds elem);
@@ -88,6 +89,7 @@ class StringSet : public DenseSet {
 
     using IteratorBase::ExpiryTime;
     using IteratorBase::HasExpiry;
+    using IteratorBase::SetExpiryTime;
   };
 
   iterator begin() {
@@ -111,7 +113,10 @@ class StringSet : public DenseSet {
 
   size_t ObjectAllocSize(const void* s1) const override;
   uint32_t ObjExpireTime(const void* obj) const override;
+  uint32_t ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) override;
+  uint32_t ObjReplace(const void* obj, uint32_t ttl_sec) override;
   void ObjDelete(void* obj, bool has_ttl) const override;
+  sds MakeSdsWithTtl(std::string_view src, uint32_t ttl_sec);
 };
 
 }  // end namespace dfly
diff --git a/src/core/string_set_test.cc b/src/core/string_set_test.cc
index 58e3ba656a60..c2272e389346 100644
--- a/src/core/string_set_test.cc
+++ b/src/core/string_set_test.cc
@@ -386,6 +386,25 @@ TEST_F(StringSetTest, Iteration) {
   EXPECT_EQ(to_insert.size(), 0);
 }
 
+TEST_F(StringSetTest, SetFieldExpireHasExpiry) {
+  EXPECT_TRUE(ss_->Add("k1", 100));
+  auto k = ss_->Find("k1");
+  EXPECT_TRUE(k.HasExpiry());
+  EXPECT_EQ(k.ExpiryTime(), 100);
+  k.SetExpiryTime(1);
+  EXPECT_TRUE(k.HasExpiry());
+  EXPECT_EQ(k.ExpiryTime(), 1);
+}
+
+TEST_F(StringSetTest, SetFieldExpireNoHasExpiry) {
+  EXPECT_TRUE(ss_->Add("k1"));
+  auto k = ss_->Find("k1");
+  EXPECT_FALSE(k.HasExpiry());
+  k.SetExpiryTime(10);
+  EXPECT_TRUE(k.HasExpiry());
+  EXPECT_EQ(k.ExpiryTime(), 10);
+}
+
 TEST_F(StringSetTest, Ttl) {
   EXPECT_TRUE(ss_->Add("bla"sv, 1));
   EXPECT_FALSE(ss_->Add("bla"sv, 1));