chore(strings): Simplify Set flow (#164)

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

CMakeLists.txt

@@ -17,8 +17,6 @@ option(BUILD_SHARED_LIBS "Build shared libraries" OFF)
 option(DF_USE_SSL "Provide support for SSL connections" ON)
 
 include(third_party)
-message(STATUS "after thirdpary")
-
 include(internal)
 
 include_directories(src)

src/core/dash.h

@@ -708,6 +708,9 @@ auto DashTable<_Key, _Value, Policy>::InsertInternal(U&& key, V&& value, Evictio
     assert(seg_id < segment_.size());
     SegmentType* target = segment_[seg_id];
 
+    // Load heap allocated segment data - to avoid TLB miss when accessing the bucket.
+    __builtin_prefetch(target, 0, 1);
+
     auto [it, res] =
         target->Insert(std::forward<U>(key), std::forward<V>(value), key_hash, EqPred());
 

src/core/dash_internal.h

@@ -1055,6 +1055,10 @@ auto Segment<Key, Value, Policy>::FindIt(U&& key, Hash_t key_hash, Pred&& cf) co
   uint8_t bidx = BucketIndex(key_hash);
   const Bucket& target = bucket_[bidx];
 
+  // It helps a bit (10% on my home machine) and more importantly, it does not hurt
+  // since we are going to access this memory in a bit.
+  __builtin_prefetch(&target);
+
   uint8_t fp_hash = key_hash & kFpMask;
   SlotId sid = target.FindByFp(fp_hash, false, key, cf);
   if (sid != BucketType::kNanSlot) {

src/core/dash_test.cc

@@ -23,7 +23,9 @@ extern "C" {
 #include "redis/zmalloc.h"
 }
 
+#if defined(__clang__)
 #pragma clang diagnostic ignored "-Wunused-const-variable"
+#endif
 
 namespace dfly {
 
@@ -1027,7 +1029,7 @@ static void BM_Insert(benchmark::State& state) {
     }
   }
 }
-BENCHMARK(BM_Insert)->Arg(1000)->Arg(10000)->Arg(100000);
+BENCHMARK(BM_Insert)->Arg(10000)->Arg(100000)->Arg(1000000);
 
 struct NoDestroySdsPolicy : public SdsDashPolicy {
   static void DestroyKey(sds s) {
@@ -1109,7 +1111,7 @@ static void BM_RedisDictInsert(benchmark::State& state) {
     dictRelease(d);
   }
 }
-BENCHMARK(BM_RedisDictInsert)->Arg(1000)->Arg(10000)->Arg(100000);
+BENCHMARK(BM_RedisDictInsert)->Arg(10000)->Arg(100000)->Arg(1000000);
 
 static void BM_RedisStringInsert(benchmark::State& state) {
   unsigned count = state.range(0);

src/redis/sds.c

@@ -622,7 +622,7 @@ int sdsull2str(char *s, unsigned long long v) {
  * sdscatprintf(sdsempty(),"%lld\n", value);
  */
 sds sdsfromlonglong(long long value) {
-    char buf[SDS_LLSTR_SIZE];
+    char buf[SDS_LLSTR_SIZE + 10];
     int len = sdsll2str(buf,value);
 
     return sdsnewlen(buf,len);

src/server/db_slice.cc

@@ -278,18 +278,23 @@ OpResult<pair<PrimeIterator, unsigned>> DbSlice::FindFirst(DbIndex db_index, Arg
   return OpStatus::KEY_NOTFOUND;
 }
 
-auto DbSlice::AddOrFind(DbIndex db_index, string_view key) noexcept(false)
-    -> pair<PrimeIterator, bool> {
+pair<PrimeIterator, bool> DbSlice::AddOrFind(DbIndex db_index, string_view key) noexcept(false) {
+  auto res = AddOrFind2(db_index, key);
+  return make_pair(get<0>(res), get<2>(res));
+}
+
+tuple<PrimeIterator, ExpireIterator, bool> DbSlice::AddOrFind2(DbIndex db_index,
+                                                             std::string_view key) noexcept(false) {
   DCHECK(IsDbValid(db_index));
 
   auto& db = db_arr_[db_index];
 
   // If we have some registered onchange callbacks, we must know in advance whether its Find or Add.
   if (!change_cb_.empty()) {
-    auto it = FindExt(db_index, key).first;
+    auto res = FindExt(db_index, key);
 
-    if (IsValid(it)) {
-      return make_pair(it, true);
+    if (IsValid(res.first)) {
+      return tuple_cat(res, make_tuple(true));
     }
 
     // It's a new entry.
@@ -326,15 +331,16 @@ auto DbSlice::AddOrFind(DbIndex db_index, string_view key) noexcept(false)
     it.SetVersion(NextVersion());
     memory_budget_ = evp.mem_budget();
 
-    return make_pair(it, true);
+    return make_tuple(it, ExpireIterator{}, true);
   }
 
   auto& existing = it;
 
   DCHECK(IsValid(existing));
 
+  ExpireIterator expire_it;
   if (existing->second.HasExpire()) {
-    auto expire_it = db->expire.Find(existing->first);
+    expire_it = db->expire.Find(existing->first);
     CHECK(IsValid(expire_it));
 
     // TODO: to implement the incremental update of expiry values using multi-generation
@@ -357,11 +363,11 @@ auto DbSlice::AddOrFind(DbIndex db_index, string_view key) noexcept(false)
       existing->second.Reset();
       events_.expired_keys++;
 
-      return make_pair(existing, true);
+      return make_tuple(existing, ExpireIterator{}, true);
     }
   }
 
-  return make_pair(existing, false);
+  return make_tuple(existing, expire_it, false);
 }
 
 void DbSlice::ActivateDb(DbIndex db_ind) {
@@ -429,7 +435,7 @@ void DbSlice::FlushDb(DbIndex db_ind) {
 }
 
 // Returns true if a state has changed, false otherwise.
-bool DbSlice::Expire(DbIndex db_ind, PrimeIterator it, uint64_t at) {
+bool DbSlice::UpdateExpire(DbIndex db_ind, PrimeIterator it, uint64_t at) {
   auto& db = *db_arr_[db_ind];
   if (at == 0 && it->second.HasExpire()) {
     CHECK_EQ(1u, db.expire.Erase(it->first));
@@ -485,19 +491,15 @@ pair<PrimeIterator, bool> DbSlice::AddOrFind(DbIndex db_ind, string_view key, Pr
     return res;
 
   auto& db = *db_arr_[db_ind];
-  auto& new_it = res.first;
+  auto& it = res.first;
 
-  size_t value_heap_size = obj.MallocUsed();
-  db.stats.obj_memory_usage += value_heap_size;
-  if (obj.ObjType() == OBJ_STRING)
-    db.stats.strval_memory_usage += value_heap_size;
-
-  new_it->second = std::move(obj);
+  it->second = std::move(obj);
+  PostUpdate(db_ind, it);
 
   if (expire_at_ms) {
-    new_it->second.SetExpire(true);
+    it->second.SetExpire(true);
     uint64_t delta = expire_at_ms - expire_base_[0];
-    CHECK(db.expire.Insert(new_it->first.AsRef(), ExpirePeriod(delta)).second);
+    CHECK(db.expire.Insert(it->first.AsRef(), ExpirePeriod(delta)).second);
   }
 
   return res;

src/server/db_slice.h

@@ -124,7 +124,10 @@ class DbSlice {
 
   // Return .second=true if insertion ocurred, false if we return the existing key.
   // throws: bad_alloc is insertion could not happen due to out of memory.
-  std::pair<PrimeIterator, bool> AddOrFind(DbIndex db_ind, std::string_view key) noexcept(false);
+  std::pair<PrimeIterator, bool> AddOrFind(DbIndex db_index, std::string_view key) noexcept(false);
+
+  std::tuple<PrimeIterator, ExpireIterator, bool> AddOrFind2(DbIndex db_index,
+                                                             std::string_view key) noexcept(false);
 
   // Returns second=true if insertion took place, false otherwise.
   // expire_at_ms equal to 0 - means no expiry.
@@ -134,7 +137,7 @@ class DbSlice {
 
   // Either adds or removes (if at == 0) expiry. Returns true if a change was made.
   // Does not change expiry if at != 0 and expiry already exists.
-  bool Expire(DbIndex db_ind, PrimeIterator main_it, uint64_t at);
+  bool UpdateExpire(DbIndex db_ind, PrimeIterator main_it, uint64_t at);
 
   void SetMCFlag(DbIndex db_ind, PrimeKey key, uint32_t flag);
   uint32_t GetMCFlag(DbIndex db_ind, const PrimeKey& key) const;
@@ -222,8 +225,8 @@ class DbSlice {
   void UnregisterOnChange(uint64_t id);
 
   struct DeleteExpiredStats {
-    uint32_t deleted = 0;    // number of deleted items due to expiry (less than traversed).
-    uint32_t traversed = 0;  // number of traversed items that have ttl bit
+    uint32_t deleted = 0;         // number of deleted items due to expiry (less than traversed).
+    uint32_t traversed = 0;       // number of traversed items that have ttl bit
     size_t survivor_ttl_sum = 0;  // total sum of ttl of survivors (traversed - deleted).
   };
 

src/server/generic_family.cc

@@ -149,7 +149,7 @@ OpStatus Renamer::UpdateDest(Transaction* t, EngineShard* es) {
         dest_it->second = std::move(pv_);
       }
       dest_it->second.SetExpire(has_expire);  // preserve expire flag.
-      db_slice.Expire(db_indx_, dest_it, src_res_.expire_ts);
+      db_slice.UpdateExpire(db_indx_, dest_it, src_res_.expire_ts);
     } else {
       if (src_res_.ref_val.ObjType() == OBJ_STRING) {
         pv_.SetString(str_val_);
@@ -621,7 +621,7 @@ OpStatus GenericFamily::OpExpire(const OpArgs& op_args, string_view key,
   } else if (IsValid(expire_it)) {
     expire_it->second = db_slice.FromAbsoluteTime(now_msec + rel_msec);
   } else {
-    db_slice.Expire(op_args.db_ind, it, rel_msec + now_msec);
+    db_slice.UpdateExpire(op_args.db_ind, it, rel_msec + now_msec);
   }
 
   return OpStatus::OK;
@@ -698,10 +698,10 @@ OpResult<void> GenericFamily::OpRen(const OpArgs& op_args, string_view from_key,
     to_it->second = std::move(from_obj);
     to_it->second.SetExpire(IsValid(to_expire));  // keep the expire flag on 'to'.
 
-    // It is guaranteed that Expire() call does not erase the element because then
-    // from_it would be invalid. Therefore, it Expire does not invalidate any iterators and
-    // we can delete via from_it.
-    db_slice.Expire(op_args.db_ind, to_it, exp_ts);
+    // It is guaranteed that UpdateExpire() call does not erase the element because then
+    // from_it would be invalid. Therefore, UpdateExpire does not invalidate any iterators,
+    // therefore we can delete 'from_it'.
+    db_slice.UpdateExpire(op_args.db_ind, to_it, exp_ts);
     CHECK(db_slice.Del(op_args.db_ind, from_it));
   } else {
     // Here we first delete from_it because AddNew below could invalidate from_it.

src/server/main_service.cc

@@ -39,9 +39,6 @@ extern "C" {
 
 using namespace std;
 
-// TODO: to move to absl flags and keep legacy flags only for glog library.
-// absl flags allow parsing of custom types and allow specifying which flags appear
-// for helpshort.
 ABSL_FLAG(uint32_t, port, 6379, "Redis port");
 ABSL_FLAG(uint32_t, memcache_port, 0, "Memcached port");
 ABSL_FLAG(uint64_t, maxmemory, 0,

src/server/string_family.cc

@@ -143,25 +143,41 @@ OpResult<void> SetCmd::Set(const SetParams& params, std::string_view key, std::s
 
   VLOG(2) << "Set " << key << "(" << db_slice_.shard_id() << ") ";
 
-  auto [it, expire_it] = db_slice_.FindExt(params.db_index, key);
+  if (params.how == SET_IF_EXISTS) {
+    auto [it, expire_it] = db_slice_.FindExt(params.db_index, key);
 
-  if (IsValid(it)) {  // existing
-    return SetExisting(params, it, expire_it, value);
+    if (IsValid(it)) {  // existing
+      return SetExisting(params, it, expire_it, value);
+    }
+    return OpStatus::SKIPPED;
   }
 
   // New entry
-  if (params.how == SET_IF_EXISTS)
-    return OpStatus::SKIPPED;
-
-  PrimeValue tvalue{value};
-  tvalue.SetFlag(params.memcache_flags != 0);
-  uint64_t at_ms = params.expire_after_ms ? params.expire_after_ms + db_slice_.Now() : 0;
+  tuple<PrimeIterator, ExpireIterator, bool> add_res;
   try {
-    it = db_slice_.AddNew(params.db_index, key, std::move(tvalue), at_ms);
+    add_res = db_slice_.AddOrFind2(params.db_index, key);
   } catch (bad_alloc& e) {
     return OpStatus::OUT_OF_MEMORY;
   }
 
+  PrimeIterator it = get<0>(add_res);
+  if (!get<2>(add_res)) {
+    return SetExisting(params, it, get<1>(add_res), value);
+  }
+
+  // adding new value.
+  PrimeValue tvalue{value};
+  tvalue.SetFlag(params.memcache_flags != 0);
+  it->second = std::move(tvalue);
+  db_slice_.PostUpdate(params.db_index, it);
+
+  if (params.expire_after_ms) {
+    db_slice_.UpdateExpire(params.db_index, it, params.expire_after_ms + db_slice_.Now());
+  }
+
+  if (params.memcache_flags)
+    db_slice_.SetMCFlag(params.db_index, it->first.AsRef(), params.memcache_flags);
+
   EngineShard* shard = db_slice_.shard_owner();
 
   if (shard->tiered_storage()) {  // external storage enabled.
@@ -170,9 +186,6 @@ OpResult<void> SetCmd::Set(const SetParams& params, std::string_view key, std::s
     }
   }
 
-  if (params.memcache_flags)
-    db_slice_.SetMCFlag(params.db_index, it->first.AsRef(), params.memcache_flags);
-
   return OpStatus::OK;
 }
 
@@ -194,7 +207,7 @@ OpStatus SetCmd::SetExisting(const SetParams& params, PrimeIterator it, ExpireIt
   if (IsValid(e_it) && at_ms) {
     e_it->second = db_slice_.FromAbsoluteTime(at_ms);
   } else {
-    bool changed = db_slice_.Expire(params.db_index, it, at_ms);
+    bool changed = db_slice_.UpdateExpire(params.db_index, it, at_ms);
     if (changed && at_ms == 0)  // erased.
       return OpStatus::OK;
   }