chore(hset_family): Add mutable functions to listpack wrapper (#5901)

Author: dranikpg

src/core/detail/listpack_wrap.cc

@@ -34,11 +34,75 @@ void ListpackWrap::Iterator::Read() {
   next_ptr_ = lpNext(lp_, next_ptr_);
 }
 
+ListpackWrap::~ListpackWrap() {
+  DCHECK(!dirty_);
+}
+
+uint8_t* ListpackWrap::GetPointer() {
+  dirty_ = false;
+  return lp_;
+}
+
 ListpackWrap::Iterator ListpackWrap::Find(std::string_view key) const {
+  if (size() == 0)
+    return end();
+
   uint8_t* ptr = lpFind(lp_, lpFirst(lp_), (unsigned char*)key.data(), key.size(), 1);
   return Iterator{lp_, ptr, intbuf_};
 }
 
+bool ListpackWrap::Delete(std::string_view key) {
+  if (size() == 0)
+    return false;
+
+  uint8_t* ptr = lpFind(lp_, lpFirst(lp_), (unsigned char*)key.data(), key.size(), 1);
+  if (ptr == nullptr)
+    return false;
+
+  lp_ = lpDeleteRangeWithEntry(lp_, &ptr, 2);
+  dirty_ = true;
+  return true;
+}
+
+bool ListpackWrap::Insert(std::string_view key, std::string_view value, bool skip_exists) {
+  uint8_t* vptr;
+  uint8_t* fptr = lpFirst(lp_);
+  uint8_t* fsrc = key.empty() ? lp_ : (uint8_t*)key.data();
+  // if we vsrc is NULL then lpReplace will delete the element, which is not what we want.
+  // therefore, for an empty val we set it to some other valid address so that lpReplace
+  // will do the right thing and encode empty string instead of deleting the element.
+  uint8_t* vsrc = value.empty() ? lp_ : (uint8_t*)value.data();
+
+  bool updated = false;
+  if (fptr) {
+    fptr = lpFind(lp_, fptr, fsrc, key.size(), 1);
+    if (fptr) {
+      if (skip_exists)
+        return false;
+
+      // Grab pointer to the value (fptr points to the field)
+      vptr = lpNext(lp_, fptr);
+
+      // Replace value
+      lp_ = lpReplace(lp_, &vptr, vsrc, value.size());
+      DCHECK_EQ(0u, lpLength(lp_) % 2);
+
+      dirty_ = true;
+      updated = true;
+    }
+  }
+
+  if (!updated) {
+    // Push new field/value pair onto the tail of the listpack.
+    // TODO: we should at least allocate once for both elements
+    lp_ = lpAppend(lp_, fsrc, key.size());
+    lp_ = lpAppend(lp_, vsrc, value.size());
+    dirty_ = true;
+  }
+
+  return !updated;
+}
+
 size_t ListpackWrap::size() const {
   return lpLength(lp_) / 2;
 }

src/core/detail/listpack_wrap.h

@@ -15,6 +15,8 @@ struct ListpackWrap {
   using IntBuf = uint8_t[2][24];
 
  public:
+  ~ListpackWrap();
+
   struct Iterator {
     using iterator_category = std::forward_iterator_tag;
     using difference_type = std::ptrdiff_t;
@@ -46,14 +48,19 @@ struct ListpackWrap {
   explicit ListpackWrap(uint8_t* lp) : lp_{lp} {
   }
 
+  uint8_t* GetPointer();                      // Get new updated pointer
   Iterator Find(std::string_view key) const;  // Linear search
+  bool Delete(std::string_view key);
+  bool Insert(std::string_view key, std::string_view value, bool skip_exists);
+
   Iterator begin() const;
   Iterator end() const;
   size_t size() const;  // number of entries
 
  private:
-  uint8_t* lp_;
-  mutable IntBuf intbuf_;
+  uint8_t* lp_;            // the listpack itself
+  mutable IntBuf intbuf_;  // buffer for integers decoded to strings
+  bool dirty_ = false;     // whether lp_ was updated, but never retrieved with GetPointer
 };
 
 }  // namespace dfly::detail

src/server/container_utils.cc

@@ -272,17 +272,6 @@ StringMap* GetStringMap(const PrimeValue& pv, const DbContext& db_context) {
   return res;
 }
 
-optional<string_view> LpFind(uint8_t* lp, string_view key, uint8_t int_buf[]) {
-  uint8_t* fptr = lpFirst(lp);
-  DCHECK(fptr);
-
-  fptr = lpFind(lp, fptr, (unsigned char*)key.data(), key.size(), 1);
-  if (!fptr)
-    return std::nullopt;
-  uint8_t* vptr = lpNext(lp, fptr);
-  return LpGetView(vptr, int_buf);
-}
-
 string_view LpGetView(uint8_t* lp_it, uint8_t int_buf[]) {
   int64_t ele_len = 0;
   uint8_t* elem = lpGet(lp_it, &ele_len, int_buf);

src/server/container_utils.h

@@ -77,9 +77,6 @@ StringMap* GetStringMap(const PrimeValue& pv, const DbContext& db_context);
 // Get string_view from listpack poiner. Intbuf to store integer values as strings.
 std::string_view LpGetView(uint8_t* lp_it, uint8_t int_buf[]);
 
-// Find value by key and return stringview to it, otherwise nullopt.
-std::optional<std::string_view> LpFind(uint8_t* lp, std::string_view key, uint8_t int_buf[]);
-
 using BlockingResultCb =
     std::function<void(Transaction*, EngineShard*, std::string_view /* key */)>;
 

src/server/hset_family.cc

@@ -51,63 +51,6 @@ bool IsGoodForListpack(CmdArgList args, const uint8_t* lp) {
 }
 
 using container_utils::GetStringMap;
-using container_utils::LpFind;
-
-pair<uint8_t*, bool> LpDelete(uint8_t* lp, string_view field) {
-  uint8_t* fptr = lpFirst(lp);
-  DCHECK(fptr);
-  fptr = lpFind(lp, fptr, (unsigned char*)field.data(), field.size(), 1);
-  if (fptr == NULL) {
-    return make_pair(lp, false);
-  }
-
-  /* Delete both of the key and the value. */
-  lp = lpDeleteRangeWithEntry(lp, &fptr, 2);
-  return make_pair(lp, true);
-}
-
-// returns a new pointer to lp. Returns true if field was inserted or false it it already existed.
-// skip_exists controls what happens if the field already existed. If skip_exists = true,
-// then val does not override the value and listpack is not changed. Otherwise, the corresponding
-// value is overridden by val.
-pair<uint8_t*, bool> LpInsert(uint8_t* lp, string_view field, string_view val, bool skip_exists) {
-  uint8_t* vptr;
-
-  uint8_t* fptr = lpFirst(lp);
-  uint8_t* fsrc = field.empty() ? lp : (uint8_t*)field.data();
-
-  // if we vsrc is NULL then lpReplace will delete the element, which is not what we want.
-  // therefore, for an empty val we set it to some other valid address so that lpReplace
-  // will do the right thing and encode empty string instead of deleting the element.
-  uint8_t* vsrc = val.empty() ? lp : (uint8_t*)val.data();
-
-  bool updated = false;
-
-  if (fptr) {
-    fptr = lpFind(lp, fptr, fsrc, field.size(), 1);
-    if (fptr) {
-      if (skip_exists) {
-        return make_pair(lp, false);
-      }
-      /* Grab pointer to the value (fptr points to the field) */
-      vptr = lpNext(lp, fptr);
-      updated = true;
-
-      /* Replace value */
-      lp = lpReplace(lp, &vptr, vsrc, val.size());
-      DCHECK_EQ(0u, lpLength(lp) % 2);
-    }
-  }
-
-  if (!updated) {
-    /* Push new field/value pair onto the tail of the listpack */
-    // TODO: we should at least allocate once for both elements.
-    lp = lpAppend(lp, fsrc, field.size());
-    lp = lpAppend(lp, vsrc, val.size());
-  }
-
-  return make_pair(lp, !updated);
-}
 
 struct HMapWrap {
  private:
@@ -263,12 +206,13 @@ OpStatus OpIncrBy(const OpArgs& op_args, string_view key, string_view field, Inc
   unsigned enc = pv.Encoding();
 
   if (enc == kEncodingListPack) {
-    uint8_t intbuf[LP_INTBUF_SIZE];
-    uint8_t* lp = (uint8_t*)pv.RObjPtr();
+    detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
     optional<string_view> res;
 
-    if (!add_res.is_new)
-      res = LpFind(lp, field, intbuf);
+    if (!add_res.is_new) {
+      if (auto it = lw.Find(field); it != lw.end())
+        res = (*it).second;
+    }
 
     OpStatus status = IncrementValue(res, param);
     if (status != OpStatus::OK) {
@@ -279,14 +223,14 @@ OpStatus OpIncrBy(const OpArgs& op_args, string_view key, string_view field, Inc
       double new_val = get<double>(*param);
       char buf[128];
       char* str = RedisReplyBuilder::FormatDouble(new_val, buf, sizeof(buf));
-      lp = LpInsert(lp, field, str, false).first;
+      lw.Insert(field, str, false);
     } else {  // integer increment
       int64_t new_val = get<int64_t>(*param);
       absl::AlphaNum an(new_val);
-      lp = LpInsert(lp, field, an.Piece(), false).first;
+      lw.Insert(field, an.Piece(), false);
     }
 
-    pv.SetRObjPtr(lp);
+    pv.SetRObjPtr(lw.GetPointer());
   } else {
     DCHECK_EQ(enc, kEncodingStrMap2);
     StringMap* sm = GetStringMap(pv, op_args.db_cntx);
@@ -387,19 +331,13 @@ OpResult<uint32_t> OpDel(const OpArgs& op_args, string_view key, CmdArgList valu
   unsigned enc = pv.Encoding();
 
   if (enc == kEncodingListPack) {
-    uint8_t* lp = (uint8_t*)pv.RObjPtr();
-    for (auto s : values) {
-      auto res = LpDelete(lp, s);
-      if (res.second) {
+    detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
+    for (string_view s : values) {
+      if (lw.Delete(s))
         ++deleted;
-        lp = res.first;
-        if (lpLength(lp) == 0) {
-          key_remove = true;
-          break;
-        }
-      }
     }
-    pv.SetRObjPtr(lp);
+    pv.SetRObjPtr(lw.GetPointer());
+    key_remove = lw.size() == 0;
   } else {
     DCHECK_EQ(enc, kEncodingStrMap2);
     StringMap* sm = GetStringMap(pv, op_args.db_cntx);
@@ -505,17 +443,16 @@ OpResult<uint32_t> OpSet(const OpArgs& op_args, string_view key, CmdArgList valu
   unsigned created = 0;
 
   if (lp) {
-    bool inserted;
     size_t malloc_reserved = zmalloc_size(lp);
     size_t min_sz = EstimateListpackMinBytes(values);
     if (min_sz > malloc_reserved) {
       lp = (uint8_t*)zrealloc(lp, min_sz);
     }
+    detail::ListpackWrap lw{lp};
     for (size_t i = 0; i < values.size(); i += 2) {
-      tie(lp, inserted) = LpInsert(lp, ArgS(values, i), ArgS(values, i + 1), op_sp.skip_if_exists);
-      created += inserted;
+      created += lw.Insert(values[i], values[i + 1], op_sp.skip_if_exists);
     }
-    pv.SetRObjPtr(lp);
+    pv.SetRObjPtr(lw.GetPointer());
   } else {
     DCHECK_EQ(kEncodingStrMap2, pv.Encoding());  // Dictionary
     StringMap* sm = GetStringMap(pv, op_args.db_cntx);
@@ -1102,10 +1039,8 @@ int32_t HSetFamily::FieldExpireTime(const DbContext& db_context, const PrimeValu
   DCHECK_EQ(OBJ_HASH, pv.ObjType());
 
   if (pv.Encoding() == kEncodingListPack) {
-    uint8_t intbuf[LP_INTBUF_SIZE];
-    uint8_t* lp = (uint8_t*)pv.RObjPtr();
-    optional<string_view> res = LpFind(lp, field, intbuf);
-    return res ? -1 : -3;
+    detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
+    return lw.Find(field) == lw.end() ? -3 : -1;
   } else {
     StringMap* string_map = (StringMap*)pv.RObjPtr();
     string_map->set_time(MemberTimeSeconds(db_context.time_now_ms));

src/server/search/doc_accessors.cc

@@ -151,14 +151,13 @@ std::optional<BaseAccessor::StringList> BaseAccessor::GetTags(std::string_view a
 
 std::optional<BaseAccessor::StringList> ListPackAccessor::GetStrings(
     string_view active_field) const {
-  auto strsv = container_utils::LpFind(lp_, active_field, intbuf_[0].data());
-  return strsv.has_value() ? StringList{*strsv} : StringList{};
+  auto it = lw_.Find(active_field);
+  return it != lw_.end() ? StringList{(*it).second} : StringList{};
 }
 
 SearchDocData ListPackAccessor::Serialize(const search::Schema& schema) const {
   SearchDocData out{};
-  detail::ListpackWrap lw{lp_};
-  for (const auto [key, value] : lw) {
+  for (const auto [key, value] : lw_) {
     if (auto field_value = ExtractSortableValue(schema, key, value); field_value) {
       out[key] = std::move(field_value).value();
     }
@@ -399,7 +398,7 @@ unique_ptr<BaseAccessor> GetAccessor(const DbContext& db_cntx, const PrimeValue&
   }
 
   if (pv.Encoding() == kEncodingListPack) {
-    auto ptr = reinterpret_cast<ListPackAccessor::LpPtr>(pv.RObjPtr());
+    auto ptr = reinterpret_cast<uint8_t*>(pv.RObjPtr());
     return make_unique<ListPackAccessor>(ptr);
   } else {
     auto* sm = container_utils::GetStringMap(pv, db_cntx);

src/server/search/doc_accessors.h

@@ -10,6 +10,7 @@
 #include <string>
 #include <utility>
 
+#include "core/detail/listpack_wrap.h"
 #include "core/json/json_object.h"
 #include "core/search/search.h"
 #include "core/search/vector_utils.h"
@@ -40,17 +41,14 @@ struct BaseAccessor : public search::DocumentAccessor {
 
 // Accessor for hashes stored with listpack
 struct ListPackAccessor : public BaseAccessor {
-  using LpPtr = uint8_t*;
-
-  explicit ListPackAccessor(LpPtr ptr) : lp_{ptr} {
+  explicit ListPackAccessor(uint8_t* ptr /* listpack ptr */) : lw_{ptr} {
   }
 
   std::optional<StringList> GetStrings(std::string_view field) const override;
   SearchDocData Serialize(const search::Schema& schema) const override;
 
  private:
-  mutable std::array<uint8_t, 33> intbuf_[2];
-  LpPtr lp_;
+  detail::ListpackWrap lw_;
 };
 
 // Accessor for hashes stored with StringMap