From 7f2dbced9c1d82681237843c97c037aa26ecbb1a 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 1/2] feat: add DenseSet::IteratorBase::SetExpiryTime

This commit is in preparation for adding FIELDEXPIRE and HEXPIRE.
---
 src/core/dense_set.cc       | 13 +++++++++-
 src/core/dense_set.h        |  5 +++-
 src/core/score_map.cc       |  6 ++++-
 src/core/score_map.h        |  5 ++--
 src/core/sds_utils.cc       |  8 ++++++
 src/core/sds_utils.h        |  3 +++
 src/core/string_map.cc      | 13 ++++++++--
 src/core/string_map.h       |  9 ++++---
 src/core/string_map_test.cc | 19 ++++++++++++++
 src/core/string_set.cc      | 50 ++++++++++++++++++-------------------
 src/core/string_set.h       |  5 +++-
 src/core/string_set_test.cc | 19 ++++++++++++++
 12 files changed, 117 insertions(+), 38 deletions(-)

diff --git a/src/core/dense_set.cc b/src/core/dense_set.cc
index 6231ae9c1bcd..69f55f22c167 100644
--- a/src/core/dense_set.cc
+++ b/src/core/dense_set.cc
@@ -46,6 +46,17 @@ DenseSet::IteratorBase::IteratorBase(const DenseSet* owner, bool is_end)
   }
 }
 
+void DenseSet::IteratorBase::SetExpiryTime(uint32_t ttl_sec) {
+  if (!HasExpiry()) {
+    auto src = curr_entry_->GetObject();
+    void* new_obj = owner_->ObjectClone(src, false, true);
+    curr_entry_->SetObject(new_obj);
+    curr_entry_->SetTtl(true);
+    owner_->ObjDelete(src, false);
+  }
+  owner_->ObjUpdateExpireTime(curr_entry_->GetObject(), ttl_sec);
+}
+
 void DenseSet::IteratorBase::Advance() {
   bool step_link = false;
   DCHECK(curr_entry_);
@@ -211,7 +222,7 @@ void DenseSet::CloneBatch(unsigned len, CloneItem* items, DenseSet* other) const
       auto& src = items[i];
       if (src.obj) {
         // The majority of the CPU is spent in this block.
-        void* new_obj = other->ObjectClone(src.obj, src.has_ttl);
+        void* new_obj = other->ObjectClone(src.obj, src.has_ttl, false);
         uint64_t hash = Hash(src.obj, 0);
         other->AddUnique(new_obj, src.has_ttl, hash);
         src.obj = nullptr;
diff --git a/src/core/dense_set.h b/src/core/dense_set.h
index a5b7ef1db4b7..8cbb1a3fecd3 100644
--- a/src/core/dense_set.h
+++ b/src/core/dense_set.h
@@ -187,6 +187,8 @@ class DenseSet {
       return curr_entry_->HasTtl() ? owner_->ObjExpireTime(curr_entry_->GetObject()) : UINT32_MAX;
     }
 
+    void SetExpiryTime(uint32_t ttl_sec);
+
     bool HasExpiry() const {
       return curr_entry_->HasTtl();
     }
@@ -265,8 +267,9 @@ 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 void ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) = 0;
   virtual void ObjDelete(void* obj, bool has_ttl) const = 0;
-  virtual void* ObjectClone(const void* obj, bool has_ttl) const = 0;
+  virtual void* ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const = 0;
 
   void CollectExpired();
 
diff --git a/src/core/score_map.cc b/src/core/score_map.cc
index 649c97c2e5ea..276627bb06d7 100644
--- a/src/core/score_map.cc
+++ b/src/core/score_map.cc
@@ -128,12 +128,16 @@ uint32_t ScoreMap::ObjExpireTime(const void* obj) const {
   return UINT32_MAX;
 }
 
+void ScoreMap::ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) {
+  // Should not reach.
+}
+
 void ScoreMap::ObjDelete(void* obj, bool has_ttl) const {
   sds s1 = (sds)obj;
   sdsfree(s1);
 }
 
-void* ScoreMap::ObjectClone(const void* obj, bool has_ttl) const {
+void* ScoreMap::ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const {
   return nullptr;
 }
 
diff --git a/src/core/score_map.h b/src/core/score_map.h
index f859b6f1f1cd..847877297387 100644
--- a/src/core/score_map.h
+++ b/src/core/score_map.h
@@ -122,8 +122,9 @@ 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;
-  void ObjDelete(void* obj, bool has_ttl) const final;
-  void* ObjectClone(const void* obj, bool has_ttl) const final;
+  void ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) override;
+  void ObjDelete(void* obj, bool has_ttl) const override;
+  void* ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const final;
 };
 
 }  // namespace dfly
diff --git a/src/core/sds_utils.cc b/src/core/sds_utils.cc
index ebd32d4fc7e1..49bbdac680e2 100644
--- a/src/core/sds_utils.cc
+++ b/src/core/sds_utils.cc
@@ -4,6 +4,8 @@
 
 #include "core/sds_utils.h"
 
+#include "base/endian.h"
+
 extern "C" {
 #include "redis/sds.h"
 #include "redis/zmalloc.h"
@@ -43,6 +45,12 @@ inline int SdsHdrSize(char type) {
 
 }  // namespace
 
+void SdsUpdateExpireTime(const void* obj, uint32_t time_at, uint32_t offset) {
+  sds str = (sds)obj;
+  char* valptr = str + sdslen(str) + 1;
+  absl::little_endian::Store32(valptr + offset, time_at);
+}
+
 char* AllocSdsWithSpace(uint32_t strlen, uint32_t space) {
   size_t usable;
   char type = SdsReqType(strlen);
diff --git a/src/core/sds_utils.h b/src/core/sds_utils.h
index f604a3af9acb..a12b4cb823e7 100644
--- a/src/core/sds_utils.h
+++ b/src/core/sds_utils.h
@@ -13,4 +13,7 @@ namespace dfly {
 // sds string (keys) with metadata attached to them.
 char* AllocSdsWithSpace(uint32_t strlen, uint32_t space);
 
+// Updates the expire time of the sds object. The offset is the number of bytes
+void SdsUpdateExpireTime(const void* obj, uint32_t time_at, uint32_t offset);
+
 }  // namespace dfly
diff --git a/src/core/string_map.cc b/src/core/string_map.cc
index a2f6b1d7df41..eaad4d6a5a55 100644
--- a/src/core/string_map.cc
+++ b/src/core/string_map.cc
@@ -276,6 +276,10 @@ uint32_t StringMap::ObjExpireTime(const void* obj) const {
   return UINT32_MAX;
 }
 
+void StringMap::ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) {
+  return SdsUpdateExpireTime(obj, time_now() + ttl_sec, 8);
+}
+
 void StringMap::ObjDelete(void* obj, bool has_ttl) const {
   sds s1 = (sds)obj;
   sds value = GetValue(s1);
@@ -283,8 +287,13 @@ void StringMap::ObjDelete(void* obj, bool has_ttl) const {
   sdsfree(s1);
 }
 
-void* StringMap::ObjectClone(const void* obj, bool has_ttl) const {
-  return nullptr;
+void* StringMap::ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const {
+  uint32_t ttl_sec = add_ttl ? 0 : (has_ttl ? ObjExpireTime(obj) : UINT32_MAX);
+  sds str = (sds)obj;
+  auto pair = detail::SdsPair(str, GetValue(str));
+  auto [newkey, sdsval_tag] = CreateEntry(pair->first, pair->second, time_now(), ttl_sec);
+
+  return (void*)newkey;
 }
 
 detail::SdsPair StringMap::iterator::BreakToPair(void* obj) {
diff --git a/src/core/string_map.h b/src/core/string_map.h
index eb513cc7063a..7930dd66d838 100644
--- a/src/core/string_map.h
+++ b/src/core/string_map.h
@@ -100,9 +100,9 @@ class StringMap : public DenseSet {
 
     using IteratorBase::ExpiryTime;
     using IteratorBase::HasExpiry;
+    using IteratorBase::SetExpiryTime;
   };
 
-  // Returns true if field was added
   // otherwise updates its value and returns false.
   bool AddOrUpdate(std::string_view field, std::string_view value, uint32_t ttl_sec = UINT32_MAX);
 
@@ -114,7 +114,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,8 +157,9 @@ 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;
-  void ObjDelete(void* obj, bool has_ttl) const final;
-  void* ObjectClone(const void* obj, bool has_ttl) const final;
+  void ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) override;
+  void ObjDelete(void* obj, bool has_ttl) const override;
+  void* ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const final;
 };
 
 }  // namespace dfly
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 bc81a716b856..9c1e28ea326f 100644
--- a/src/core/string_set.cc
+++ b/src/core/string_set.cc
@@ -27,7 +27,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 +43,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 = MakeSetSds(src, ttl_sec);
+  bool has_ttl = ttl_sec != UINT32_MAX;
 
   if (AddOrFindObj(newsds, has_ttl) != nullptr) {
     sdsfree(newsds);
@@ -129,22 +115,34 @@ uint32_t StringSet::ObjExpireTime(const void* str) const {
   return absl::little_endian::Load32(ttlptr);
 }
 
+void StringSet::ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) {
+  return SdsUpdateExpireTime(obj, time_now() + ttl_sec, 0);
+}
+
 void StringSet::ObjDelete(void* obj, bool has_ttl) const {
   sdsfree((sds)obj);
 }
 
-void* StringSet::ObjectClone(const void* obj, bool has_ttl) const {
+void* StringSet::ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const {
   sds src = (sds)obj;
-  if (has_ttl) {
-    size_t slen = sdslen(src);
-    char* ttlptr = src + slen + 1;
-    uint32_t at = absl::little_endian::Load32(ttlptr);
-    sds newsds = AllocImmutableWithTtl(slen, at);
-    if (slen)
-      memcpy(newsds, src, slen);
+  string_view sv{src, sdslen(src)};
+  uint32_t ttl_sec = add_ttl ? 0 : (has_ttl ? ObjExpireTime(obj) : UINT32_MAX);
+  return (void*)MakeSetSds(sv, ttl_sec);
+}
+
+sds StringSet::MakeSetSds(string_view src, uint32_t ttl_sec) const {
+  DCHECK_GT(ttl_sec, 0u);  // ttl_sec == 0 would mean find and delete immediately
+  if (ttl_sec != UINT32_MAX) {
+    uint32_t at = time_now() + ttl_sec;
+    DCHECK_LT(time_now(), at);
+
+    sds newsds = AllocImmutableWithTtl(src.size(), at);
+    if (!src.empty())
+      memcpy(newsds, src.data(), src.size());
     return newsds;
   }
-  return sdsnewlen(src, sdslen(src));
+
+  return sdsnewlen(src.data(), src.size());
 }
 
 }  // namespace dfly
diff --git a/src/core/string_set.h b/src/core/string_set.h
index 698d77f05d78..544d4b3f7e52 100644
--- a/src/core/string_set.h
+++ b/src/core/string_set.h
@@ -88,6 +88,7 @@ class StringSet : public DenseSet {
 
     using IteratorBase::ExpiryTime;
     using IteratorBase::HasExpiry;
+    using IteratorBase::SetExpiryTime;
   };
 
   iterator begin() {
@@ -111,8 +112,10 @@ class StringSet : public DenseSet {
 
   size_t ObjectAllocSize(const void* s1) const override;
   uint32_t ObjExpireTime(const void* obj) const override;
+  void ObjUpdateExpireTime(const void* obj, uint32_t ttl_sec) override;
   void ObjDelete(void* obj, bool has_ttl) const override;
-  void* ObjectClone(const void* obj, bool has_ttl) const override;
+  void* ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const override;
+  sds MakeSetSds(std::string_view src, uint32_t ttl_sec) const;
 };
 
 }  // end namespace dfly
diff --git a/src/core/string_set_test.cc b/src/core/string_set_test.cc
index 0fb1e064f333..24c314749103 100644
--- a/src/core/string_set_test.cc
+++ b/src/core/string_set_test.cc
@@ -377,6 +377,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));

From 969a24f228825de93b5eae13fb5895ff4d4e8735 Mon Sep 17 00:00:00 2001
From: NegatioN <4037769+NegatioN@users.noreply.github.com>
Date: Tue, 1 Oct 2024 09:48:31 +0200
Subject: [PATCH 2/2] fix: 0 is a valid input for MakeSetSds

---
 src/core/string_set.cc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/core/string_set.cc b/src/core/string_set.cc
index 9c1e28ea326f..75d670a3d090 100644
--- a/src/core/string_set.cc
+++ b/src/core/string_set.cc
@@ -131,10 +131,8 @@ void* StringSet::ObjectClone(const void* obj, bool has_ttl, bool add_ttl) const
 }
 
 sds StringSet::MakeSetSds(string_view src, uint32_t ttl_sec) const {
-  DCHECK_GT(ttl_sec, 0u);  // ttl_sec == 0 would mean find and delete immediately
   if (ttl_sec != UINT32_MAX) {
     uint32_t at = time_now() + ttl_sec;
-    DCHECK_LT(time_now(), at);
 
     sds newsds = AllocImmutableWithTtl(src.size(), at);
     if (!src.empty())