feat(server): Extend string_map api, fix AddOrSet method.

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

src/core/dense_set.cc

@@ -313,7 +313,7 @@ void DenseSet::Grow() {
   }
 }
 
-bool DenseSet::AddInternal(void* ptr, bool has_ttl) {
+void* DenseSet::AddOrFind(void* ptr, bool has_ttl) {
   uint64_t hc = Hash(ptr, 0);
 
   if (entries_.empty()) {
@@ -325,13 +325,14 @@ bool DenseSet::AddInternal(void* ptr, bool has_ttl) {
     ++size_;
     ++num_used_buckets_;
 
-    return true;
+    return nullptr;
   }
 
   // if the value is already in the set exit early
   uint32_t bucket_id = BucketId(hc);
-  if (Find(ptr, bucket_id, 0).second != nullptr) {
-    return false;
+  DensePtr* dptr = Find(ptr, bucket_id, 0).second;
+  if (dptr != nullptr) {
+    return dptr->GetObject();
   }
 
   DCHECK_LT(bucket_id, entries_.size());
@@ -347,7 +348,7 @@ bool DenseSet::AddInternal(void* ptr, bool has_ttl) {
       }
       ++num_used_buckets_;
       ++size_;
-      return true;
+      return nullptr;
     }
 
     if (size_ < entries_.size()) {
@@ -392,7 +393,7 @@ bool DenseSet::AddInternal(void* ptr, bool has_ttl) {
   DCHECK(!entries_[bucket_id].IsDisplaced());
 
   ++size_;
-  return true;
+  return nullptr;
 }
 
 auto DenseSet::Find(const void* ptr, uint32_t bid, uint32_t cookie) -> pair<DensePtr*, DensePtr*> {

src/core/dense_set.h

@@ -245,7 +245,9 @@ class DenseSet {
     return false;
   }
 
-  bool AddInternal(void* obj, bool has_ttl);
+  // Returns previous object if the object with such key already exists,
+  // Returns null if obj was added.
+  void* AddOrFind(void* obj, bool has_ttl);
 
   void* FindInternal(const void* obj, uint32_t cookie) const {
     DensePtr* ptr = const_cast<DenseSet*>(this)->Find(obj, BucketId(obj, cookie), cookie).second;
@@ -259,6 +261,14 @@ class DenseSet {
   // DenseSet. All data allocated by a derived class should be freed before calling this
   void ClearInternal();
 
+  void IncreaseMallocUsed(size_t delta) {
+    obj_malloc_used_ += delta;
+  }
+
+  void DecreaseMallocUsed(size_t delta) {
+    obj_malloc_used_ -= delta;
+  }
+
  private:
   DenseSet(const DenseSet&) = delete;
   DenseSet& operator=(DenseSet&) = delete;

src/core/string_map.cc

@@ -19,9 +19,9 @@ namespace dfly {
 
 namespace {
 
-sds GetValue(sds key) {
-  char* from = key + sdslen(key) + 1;
-  return (char*)absl::little_endian::Load64(from);
+inline sds GetValue(sds key) {
+  char* valptr = key + sdslen(key) + 1;
+  return (char*)absl::little_endian::Load64(valptr);
 }
 
 }  // namespace
@@ -41,11 +41,18 @@ bool StringMap::AddOrSet(string_view field, string_view value, uint32_t ttl_sec)
 
   sds val = sdsnewlen(value.data(), value.size());
   absl::little_endian::Store64(newkey + field.size() + 1, uint64_t(val));
-  bool has_ttl = false;
 
-  if (!AddInternal(newkey, has_ttl)) {
-    ObjDelete(newkey, has_ttl);
-    LOG(FATAL) << "TBD:ORSET";
+  bool has_ttl = false;
+  sds prev_entry = (sds)AddOrFind(newkey, has_ttl);
+  if (prev_entry) {
+    sdsfree(newkey);
+    char* valptr = prev_entry + sdslen(prev_entry) + 1;
+    sds prev_val = (sds)absl::little_endian::Load64(valptr);
+    DecreaseMallocUsed(zmalloc_usable_size(sdsAllocPtr(prev_val)));
+    sdsfree(prev_val);
+
+    absl::little_endian::Store64(valptr, uint64_t(val));
+    IncreaseMallocUsed(zmalloc_usable_size(sdsAllocPtr(val)));
 
     return false;
   }
@@ -53,7 +60,8 @@ bool StringMap::AddOrSet(string_view field, string_view value, uint32_t ttl_sec)
   return true;
 }
 
-bool StringMap::Erase(string_view field) {
+bool StringMap::Erase(string_view key) {
+  LOG(FATAL) << "TBD";
   return false;
 }
 
@@ -71,9 +79,7 @@ sds StringMap::Find(std::string_view key) {
   if (!str)
     return nullptr;
 
-  char* valptr = str + sdslen(str) + 1;
-  sds res = (sds)absl::little_endian::Load64(valptr);
-  return res;
+  return GetValue(str);
 }
 
 uint64_t StringMap::Hash(const void* obj, uint32_t cookie) const {
@@ -128,4 +134,9 @@ void StringMap::ObjDelete(void* obj, bool has_ttl) const {
   sdsfree(s1);
 }
 
+detail::SdsPair StringMap::iterator::BreakToPair(void* obj) {
+  sds f = (sds)obj;
+  return detail::SdsPair(f, GetValue(f));
+}
+
 }  // namespace dfly

src/core/string_map.h

@@ -12,6 +12,27 @@ extern "C" {
 
 namespace dfly {
 
+namespace detail {
+
+class SdsPair {
+ public:
+  SdsPair(sds k, sds v) : first(k), second(v) {
+  }
+
+  SdsPair* operator->() {
+    return this;
+  }
+
+  const SdsPair* operator->() const {
+    return this;
+  }
+
+  const sds first;
+  const sds second;
+};
+
+};  // namespace detail
+
 class StringMap : public DenseSet {
  public:
   StringMap(std::pmr::memory_resource* res = std::pmr::get_default_resource()) : DenseSet(res) {
@@ -20,14 +41,41 @@ class StringMap : public DenseSet {
   ~StringMap();
 
   class iterator : private DenseSet::IteratorBase {
+    static detail::SdsPair BreakToPair(void* obj);
+
    public:
     iterator() : IteratorBase() {
     }
 
     iterator(DenseSet* owner, bool is_end) : IteratorBase(owner, is_end) {
     }
+
+    detail::SdsPair operator->() const {
+      void* ptr = curr_entry_->GetObject();
+      return BreakToPair(ptr);
+    }
+
+    detail::SdsPair operator*() const {
+      void* ptr = curr_entry_->GetObject();
+      return BreakToPair(ptr);
+    }
+
+    iterator& operator++() {
+      Advance();
+      return *this;
+    }
+
+    bool operator==(const iterator& b) const {
+      return curr_list_ == b.curr_list_;
+    }
+
+    bool operator!=(const iterator& b) const {
+      return !(*this == b);
+    }
   };
 
+  // Returns true if field was added
+  // otherwise updates its value and returns false.
   bool AddOrSet(std::string_view field, std::string_view value, uint32_t ttl_sec = UINT32_MAX);
 
   bool Erase(std::string_view s1);
@@ -37,6 +85,14 @@ class StringMap : public DenseSet {
 
   void Clear();
 
+  iterator begin() {
+    return iterator{this, false};
+  }
+
+  iterator end() {
+    return iterator{this, true};
+  }
+
  private:
   uint64_t Hash(const void* obj, uint32_t cookie) const final;
   bool ObjEqual(const void* left, const void* right, uint32_t right_cookie) const final;

src/core/string_map_test.cc

@@ -73,6 +73,23 @@ TEST_F(StringMapTest, Basic) {
   EXPECT_TRUE(sm_->AddOrSet("foo", "bar"));
   EXPECT_TRUE(sm_->Contains("foo"));
   EXPECT_STREQ("bar", sm_->Find("foo"));
+
+  auto it = sm_->begin();
+  EXPECT_STREQ("foo", it->first);
+  EXPECT_STREQ("bar", it->second);
+  ++it;
+  EXPECT_TRUE(it == sm_->end());
+
+  for (const auto& k_v : *sm_) {
+    EXPECT_STREQ("foo", k_v.first);
+    EXPECT_STREQ("bar", k_v.second);
+  }
+
+  size_t sz = sm_->ObjMallocUsed();
+  EXPECT_FALSE(sm_->AddOrSet("foo", "baraaaaaaaaaaaa2"));
+  EXPECT_GT(sm_->ObjMallocUsed(), sz);
+  it = sm_->begin();
+  EXPECT_STREQ("baraaaaaaaaaaaa2", it->second);
 }
 
 }  // namespace dfly

src/core/string_set.cc

@@ -39,7 +39,7 @@ StringSet::~StringSet() {
 }
 
 bool StringSet::AddSds(sds s1) {
-  return AddInternal(s1, false);
+  return AddOrFind(s1, false) == nullptr;
 }
 
 bool StringSet::Add(string_view src, uint32_t ttl_sec) {
@@ -60,7 +60,7 @@ bool StringSet::Add(string_view src, uint32_t ttl_sec) {
     has_ttl = true;
   }
 
-  if (!AddInternal(newsds, has_ttl)) {
+  if (AddOrFind(newsds, has_ttl) != nullptr) {
     sdsfree(newsds);
     return false;
   }

src/core/string_set.h

@@ -23,10 +23,11 @@ class StringSet : public DenseSet {
 
   ~StringSet();
 
+  // Returns true if elem was added.
   bool Add(std::string_view s1, uint32_t ttl_sec = UINT32_MAX);
 
-  // Used currently by rdb_load.
-  bool AddSds(sds s1);
+  // Used currently by rdb_load. Returns true if elem was added.
+  bool AddSds(sds elem);
 
   bool Erase(std::string_view s1);