fixes and small tests

Author: dranikpg

src/core/compact_object.cc

@@ -1224,11 +1224,13 @@ CompactObj::ExternalRep CompactObj::GetExternalRep() const {
   return static_cast<CompactObj::ExternalRep>(u_.ext_ptr.representation);
 }
 
-void CompactObj::SetCool(size_t offset, uint32_t sz, detail::TieredColdRecord* record) {
+void CompactObj::SetCool(size_t offset, uint32_t sz, ExternalRep rep,
+                         detail::TieredColdRecord* record) {
   // We copy the mask of the "cooled" referenced object because it contains the encoding info.
   SetMeta(EXTERNAL_TAG, record->value.mask_);
 
   u_.ext_ptr.is_cool = 1;
+  u_.ext_ptr.representation = static_cast<uint8_t>(rep);
   u_.ext_ptr.page_offset = offset % 4096;
   u_.ext_ptr.serialized_size = sz;
   u_.ext_ptr.cool_record = record;

src/core/compact_object.h

@@ -364,7 +364,8 @@ class CompactObj {
   }
 
   // Assigns a cooling record to the object together with its external slice.
-  void SetCool(size_t offset, uint32_t serialized_size, detail::TieredColdRecord* record);
+  void SetCool(size_t offset, uint32_t serialized_size, ExternalRep rep,
+               detail::TieredColdRecord* record);
 
   struct CoolItem {
     uint16_t page_offset;
@@ -377,7 +378,7 @@ class CompactObj {
   CoolItem GetCool() const;
 
   // Prequisite: IsCool() is true.
-  // Keeps cool record only as external value.
+  // Keeps cool record only as external value and discard in-memory part.
   void KeepExternal(size_t offset, size_t sz);
 
   std::pair<size_t, size_t> GetExternalSlice() const;

src/server/common.cc

@@ -162,7 +162,7 @@ bool ParseDouble(string_view src, double* value) {
 #define ADD(x) (x) += o.x
 
 TieredStats& TieredStats::operator+=(const TieredStats& o) {
-  static_assert(sizeof(TieredStats) == 160);
+  static_assert(sizeof(TieredStats) == 168);
 
   ADD(total_stashes);
   ADD(total_fetches);
@@ -182,6 +182,8 @@ TieredStats& TieredStats::operator+=(const TieredStats& o) {
   ADD(small_bins_cnt);
   ADD(small_bins_entries_cnt);
   ADD(small_bins_filling_bytes);
+  ADD(small_bins_filling_entries_cnt);
+
   ADD(total_stash_overflows);
   ADD(cold_storage_bytes);
   ADD(total_offloading_steps);

src/server/common.h

@@ -91,6 +91,7 @@ struct TieredStats {
   uint64_t small_bins_cnt = 0;
   uint64_t small_bins_entries_cnt = 0;
   size_t small_bins_filling_bytes = 0;
+  size_t small_bins_filling_entries_cnt = 0;
   size_t cold_storage_bytes = 0;
 
   uint64_t clients_throttled = 0;        // current number of throttled clients

src/server/hset_family.cc

@@ -492,9 +492,8 @@ OpResult<uint32_t> OpSet(const OpArgs& op_args, string_view key, CmdArgList valu
 
   op_args.shard->search_indices()->AddDoc(key, op_args.db_cntx, pv);
 
-  // TODO: Enable in final PR under flag
-  // if (auto* ts = op_args.shard->tiered_storage(); ts)
-  //  ts->TryStash(op_args.db_cntx.db_index, key, &pv);
+  if (auto* ts = op_args.shard->tiered_storage(); ts)
+    ts->TryStash(op_args.db_cntx.db_index, key, &pv);
 
   return created;
 }

src/server/tiered_storage.cc

@@ -52,7 +52,7 @@ ABSL_FLAG(float, tiered_offload_threshold, 0.5,
 ABSL_FLAG(float, tiered_upload_threshold, 0.1,
           "Ratio of free memory (free/max memory) below which uploading stops");
 
-ABSL_FLAG(bool, tiered_experimental_hash_offload, false, "Experimental hash datatype offloading");
+ABSL_FLAG(bool, tiered_experimental_hash_support, false, "Experimental hash datatype offloading");
 
 namespace dfly {
 
@@ -81,9 +81,12 @@ tiering::DiskSegment FromCoolItem(const PrimeValue::CoolItem& item) {
   return {item.record->page_index * tiering::kPageSize + item.page_offset, item.serialized_size};
 }
 
+// Determine required byte size and encoding type based on value.
+// TODO(vlad): Maybe split into different accessors?
 // Do NOT enforce rules depending on dynamic runtime values as this is called
 // when scheduling stash and just before succeeeding and is expected to return the same results
-optional<std::pair<size_t, CompactObj::ExternalRep>> EstimateSerializedSize(const PrimeValue& pv) {
+optional<pair<size_t /*size*/, CompactObj::ExternalRep>> EstimateSerializedSize(
+    const PrimeValue& pv) {
   switch (pv.ObjType()) {
     case OBJ_STRING:
       return std::make_pair(pv.GetRawString().view().size(), CompactObj::ExternalRep::STRING);
@@ -111,11 +114,11 @@ size_t Serialize(CompactObj::ExternalRep rep, const PrimeValue& pv, io::MutableB
     case CompactObj::ExternalRep::SERIALIZED_MAP: {
       DCHECK_EQ(pv.Encoding(), kEncodingListPack);
 
+      // TODO(vlad): Optimize copy for serialization
       detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
       vector<pair<string, string>> entries(lw.begin(), lw.end());
-      vector<pair<string_view, string_view>> entries_sv(entries.begin(), entries.end());
       return tiering::SerializedMap::Serialize(
-          entries_sv, {reinterpret_cast<char*>(buffer.data()), buffer.length()});
+          entries, {reinterpret_cast<char*>(buffer.data()), buffer.length()});
     }
   };
   return 0;
@@ -218,13 +221,13 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
       stats->tiered_used_bytes += segment.length;
       stats_.total_stashes++;
 
+      CompactObj::ExternalRep rep = EstimateSerializedSize(*pv)->second;
       if (ts_->config_.experimental_cooling) {
         RetireColdEntries(pv->MallocUsed());
-        ts_->CoolDown(key.first, key.second, segment, pv);
+        ts_->CoolDown(key.first, key.second, segment, rep, pv);
       } else {
         stats->AddTypeMemoryUsage(pv->ObjType(), -pv->MallocUsed());
-        auto estimation = EstimateSerializedSize(*pv);
-        pv->SetExternal(segment.offset, segment.length, estimation->second);
+        pv->SetExternal(segment.offset, segment.length, rep);
       }
     } else {
       LOG(DFATAL) << "Should not reach here";
@@ -459,7 +462,7 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
   error_code ec;
 
   value->SetStashPending(true);
-  if (true /*OccupiesWholePages(*estimated)*/) {  // large enough for own page
+  if (OccupiesWholePages(estimated->first)) {  // large enough for own page
     id = KeyRef(dbid, key);
     if (auto prepared = op_manager_->PrepareStash(estimated->first); prepared) {
       auto [offset, buf] = *prepared;
@@ -560,6 +563,7 @@ TieredStats TieredStorage::GetStats() const {
     stats.small_bins_cnt = bins_stats.stashed_bins_cnt;
     stats.small_bins_entries_cnt = bins_stats.stashed_entries_cnt;
     stats.small_bins_filling_bytes = bins_stats.current_bin_bytes;
+    stats.small_bins_filling_entries_cnt = bins_stats.current_entries_cnt;
   }
 
   {  // Own stats
@@ -584,14 +588,14 @@ void TieredStorage::UpdateFromFlags() {
       .write_depth_limit = absl::GetFlag(FLAGS_tiered_storage_write_depth),
       .offload_threshold = absl::GetFlag(FLAGS_tiered_offload_threshold),
       .upload_threshold = absl::GetFlag(FLAGS_tiered_upload_threshold),
-      .experimental_hash_offload = absl::GetFlag(FLAGS_tiered_experimental_hash_offload),
+      .experimental_hash_offload = absl::GetFlag(FLAGS_tiered_experimental_hash_support),
   };
 }
 
 std::vector<std::string> TieredStorage::GetMutableFlagNames() {
   return base::GetFlagNames(FLAGS_tiered_min_value_size, FLAGS_tiered_experimental_cooling,
                             FLAGS_tiered_storage_write_depth, FLAGS_tiered_offload_threshold,
-                            FLAGS_tiered_upload_threshold, FLAGS_tiered_experimental_hash_offload);
+                            FLAGS_tiered_upload_threshold, FLAGS_tiered_experimental_hash_support);
 }
 
 bool TieredStorage::ShouldOffload() const {
@@ -702,7 +706,8 @@ bool TieredStorage::ShouldStash(const PrimeValue& pv) const {
 }
 
 void TieredStorage::CoolDown(DbIndex db_ind, std::string_view str,
-                             const tiering::DiskSegment& segment, PrimeValue* pv) {
+                             const tiering::DiskSegment& segment, CompactObj::ExternalRep rep,
+                             PrimeValue* pv) {
   detail::TieredColdRecord* record = CompactObj::AllocateMR<detail::TieredColdRecord>();
   cool_queue_.push_front(*record);
   stats_.cool_memory_used += (sizeof(detail::TieredColdRecord) + pv->MallocUsed());
@@ -712,7 +717,7 @@ void TieredStorage::CoolDown(DbIndex db_ind, std::string_view str,
   record->page_index = segment.offset / tiering::kPageSize;
   record->value = std::move(*pv);
 
-  pv->SetCool(segment.offset, segment.length, record);
+  pv->SetCool(segment.offset, segment.length, rep, record);
 }
 
 PrimeValue TieredStorage::Warmup(DbIndex dbid, PrimeValue::CoolItem item) {
@@ -721,10 +726,6 @@ PrimeValue TieredStorage::Warmup(DbIndex dbid, PrimeValue::CoolItem item) {
   // We remove it from both cool storage and the offline storage.
   PrimeValue hot = DeleteCool(item.record);
   op_manager_->DeleteOffloaded(dbid, segment);
-
-  // Bring it back to the PrimeTable.
-  DCHECK(hot.ObjType() == OBJ_STRING);
-
   return hot;
 }
 

src/server/tiered_storage.h

@@ -116,7 +116,7 @@ class TieredStorage {
 
   // Moves pv contents to the cool storage and updates pv to point to it.
   void CoolDown(DbIndex db_ind, std::string_view str, const tiering::DiskSegment& segment,
-                PrimeValue* pv);
+                CompactObj::ExternalRep rep, PrimeValue* pv);
 
   PrimeValue DeleteCool(detail::TieredColdRecord* record);
   detail::TieredColdRecord* PopCool();

src/server/tiered_storage_test.cc

@@ -28,6 +28,7 @@ ABSL_DECLARE_FLAG(float, tiered_offload_threshold);
 ABSL_DECLARE_FLAG(float, tiered_upload_threshold);
 ABSL_DECLARE_FLAG(unsigned, tiered_storage_write_depth);
 ABSL_DECLARE_FLAG(bool, tiered_experimental_cooling);
+ABSL_DECLARE_FLAG(bool, tiered_experimental_hash_support);
 
 namespace dfly {
 
@@ -454,4 +455,41 @@ TEST_F(PureDiskTSTest, Dump) {
   EXPECT_EQ(resp, "OK");
 }
 
+TEST_P(LatentCoolingTSTest, SimpleHash) {
+  absl::FlagSaver saver;
+  absl::SetFlag(&FLAGS_tiered_experimental_hash_support, true);
+  UpdateFromFlags();
+
+  const size_t kNUM = 100;
+
+  auto build_command = [](string_view key) {
+    vector<string> cmd = {"HSET", string{key}};
+    for (char c = 'a'; c <= 'z'; c++) {
+      cmd.push_back(string{1, c});
+      cmd.push_back(string{31, 'x'} + c);
+    }
+    return cmd;
+  };
+
+  // Create some hashes
+  for (size_t i = 0; i < kNUM; i++) {
+    Run(build_command(absl::StrCat("k", i)));
+  }
+
+  // Wait for all to be stashed or in end up in bins
+  ExpectConditionWithinTimeout([=] {
+    auto metrics = GetMetrics();
+    return metrics.tiered_stats.total_stashes +
+               metrics.tiered_stats.small_bins_filling_entries_cnt ==
+           kNUM;
+  });
+
+  // Verify correctness
+  for (size_t i = 0; i < kNUM; i++) {
+    auto resp = Run({"HGET", absl::StrCat("k", i), string{1, 'f'}});
+    auto v = string{31, 'x'} + 'f';
+    EXPECT_EQ(resp, v);
+  }
+}
+
 }  // namespace dfly

src/server/tiering/serialized_map.h

@@ -50,7 +50,7 @@ struct SerializedMap {
   size_t DataBytes() const;
 
   // Input for serialization
-  using Input = const absl::Span<const std::pair<std::string_view, std::string_view>>;
+  using Input = const absl::Span<const std::pair<std::string, std::string>>;
 
   // Buffer size required for serialization
   static size_t SerializeSize(Input);

src/server/tiering/serialized_map_test.cc

@@ -12,11 +12,11 @@ using namespace std;
 class SerializedMapTest : public ::testing::Test {};
 
 TEST_F(SerializedMapTest, TestBasic) {
-  const vector<std::pair<string_view, string_view>> kBase = {{"first key", "first value"},
-                                                             {"second key", "second value"},
-                                                             {"third key", "third value"},
-                                                             {"fourth key", "fourth value"},
-                                                             {"fifth key", "fifth value"}};
+  const vector<std::pair<string, string>> kBase = {{"first key", "first value"},
+                                                   {"second key", "second value"},
+                                                   {"third key", "third value"},
+                                                   {"fourth key", "fourth value"},
+                                                   {"fifth key", "fifth value"}};
 
   // Serialize kBase to buffer
   std::string buffer;