fixes

Author: dranikpg

src/core/compact_object.cc

@@ -879,10 +879,13 @@ CompactObjType CompactObj::ObjType() const {
     return OBJ_STRING;
 
   if (taglen_ == EXTERNAL_TAG) {
+    VLOG(0) << "My type is external";
     switch (static_cast<ExternalRep>(u_.ext_ptr.representation)) {
       case ExternalRep::STRING:
+        VLOG(0) << "My type is a string";
         return OBJ_STRING;
       case ExternalRep::SERIALIZED_MAP:
+        VLOG(0) << "Mype is a hash map";
         return OBJ_HASH;
     };
   }
@@ -1244,6 +1247,10 @@ auto CompactObj::GetCool() const -> CoolItem {
   return res;
 }
 
+void CompactObj::KeepExternal(size_t offset, size_t sz) {
+  SetExternal(offset, sz, GetExternalRep());
+}
+
 std::pair<size_t, size_t> CompactObj::GetExternalSlice() const {
   DCHECK_EQ(EXTERNAL_TAG, taglen_);
   auto& ext = u_.ext_ptr;

src/core/compact_object.h

@@ -376,6 +376,10 @@ class CompactObj {
   // Returns the external data of the object incuding its ColdRecord.
   CoolItem GetCool() const;
 
+  // Prequisite: IsCool() is true.
+  // Keeps cool record only as external value.
+  void KeepExternal(size_t offset, size_t sz);
+
   std::pair<size_t, size_t> GetExternalSlice() const;
 
   // Injects either the the raw string (extracted with GetRawString()) or the usual string

src/server/hset_family.cc

@@ -492,6 +492,9 @@ 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);
 
+  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

@@ -79,32 +79,34 @@ tiering::DiskSegment FromCoolItem(const PrimeValue::CoolItem& item) {
   return {item.record->page_index * tiering::kPageSize + item.page_offset, item.serialized_size};
 }
 
-optional<size_t> EstimateSerializedSize(const PrimeValue& pv) {
+// 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) {
   switch (pv.ObjType()) {
     case OBJ_STRING:
-      return pv.GetRawString().view().size();
+      return std::make_pair(pv.GetRawString().view().size(), CompactObj::ExternalRep::STRING);
     case OBJ_HASH:
       if (pv.Encoding() == kEncodingListPack) {
         auto* lp = static_cast<uint8_t*>(pv.RObjPtr());
         size_t bytes = lpBytes(lp);
         bytes += lpLength(lp) * 2 * 4;
-        return bytes;
+        return std::make_pair(bytes, CompactObj::ExternalRep::SERIALIZED_MAP);
       }
       return {};
     default:
       return {};
   };
 }
 
-size_t Serialize(const PrimeValue& pv, io::MutableBytes buffer) {
-  DCHECK_LE(EstimateSerializedSize(pv), buffer.size());
-  switch (pv.ObjType()) {
-    case OBJ_STRING: {
+size_t Serialize(CompactObj::ExternalRep rep, const PrimeValue& pv, io::MutableBytes buffer) {
+  DCHECK_LE(EstimateSerializedSize(pv)->first, buffer.size());
+  switch (rep) {
+    case CompactObj::ExternalRep::STRING: {
       auto sv = pv.GetRawString();
       memcpy(buffer.data(), sv.view().data(), sv.view().size());
       return sv.view().size();
     }
-    case OBJ_HASH: {
+    case CompactObj::ExternalRep::SERIALIZED_MAP: {
       DCHECK_EQ(pv.Encoding(), kEncodingListPack);
 
       detail::ListpackWrap lw{static_cast<uint8_t*>(pv.RObjPtr())};
@@ -113,15 +115,15 @@ size_t Serialize(const PrimeValue& pv, io::MutableBytes buffer) {
       return tiering::SerializedMap::Serialize(
           entries_sv, {reinterpret_cast<char*>(buffer.data()), buffer.length()});
     }
-    default:
-      DCHECK(false);
-      return 0;
-  }
+  };
+  return 0;
 }
 
 string SerializeString(const PrimeValue& pv) {
-  string s(*EstimateSerializedSize(pv), 0);
-  size_t written = Serialize(pv, {reinterpret_cast<uint8_t*>(s.data()), s.size()});
+  auto estimate = EstimateSerializedSize(pv);
+  string s(estimate->first, 0);
+  size_t written =
+      Serialize(estimate->second, pv, {reinterpret_cast<uint8_t*>(s.data()), s.size()});
   s.resize(written);
   return s;
 }
@@ -186,12 +188,20 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
   void RetireColdEntries(size_t additional_memory);
 
   // Set value to be an in-memory type again. Update memory stats.
-  void Upload(DbIndex dbid, string_view value, bool is_raw, size_t serialized_len, PrimeValue* pv) {
+  void Upload(DbIndex dbid, string_view value, PrimeValue* pv) {
     DCHECK(!value.empty());
     DCHECK_EQ(uint8_t(pv->GetExternalRep()), uint8_t(CompactObj::ExternalRep::STRING));
 
-    pv->Materialize(value, is_raw);
-    RecordDeleted(*pv, serialized_len, GetDbTableStats(dbid));
+    switch (pv->GetExternalRep()) {
+      case CompactObj::ExternalRep::STRING:
+        pv->Materialize(value, true);
+        break;
+      case CompactObj::ExternalRep::SERIALIZED_MAP:
+        pv->InitRobj(OBJ_HASH, kEncodingListPack, nullptr);
+        break;
+    };
+
+    RecordDeleted(*pv, value.size(), GetDbTableStats(dbid));
   }
 
   // Find entry by key in db_slice and store external segment in place of original value.
@@ -211,7 +221,8 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
         ts_->CoolDown(key.first, key.second, segment, pv);
       } else {
         stats->AddTypeMemoryUsage(pv->ObjType(), -pv->MallocUsed());
-        pv->SetExternal(segment.offset, segment.length, CompactObj::ExternalRep::STRING);
+        auto estimation = EstimateSerializedSize(*pv);
+        pv->SetExternal(segment.offset, segment.length, estimation->second);
       }
     } else {
       LOG(DFATAL) << "Should not reach here";
@@ -268,7 +279,7 @@ void TieredStorage::ShardOpManager::Defragment(tiering::DiskSegment segment, str
     } else {
       // Cut out relevant part of value and restore it to memory
       string_view value = page.substr(item_segment.offset - segment.offset, item_segment.length);
-      Upload(dbid, value, true, item_segment.length, &pv);
+      Upload(dbid, value, &pv);
     }
   }
 }
@@ -439,18 +450,18 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
     return {};
   }
 
-  optional<size_t> estimated = EstimateSerializedSize(*value);
+  auto estimated = EstimateSerializedSize(*value);
   DCHECK(estimated);
 
   tiering::OpManager::EntryId id;
   error_code ec;
 
   value->SetStashPending(true);
-  if (OccupiesWholePages(*estimated)) {  // large enough for own page
+  if (true /*OccupiesWholePages(*estimated)*/) {  // large enough for own page
     id = KeyRef(dbid, key);
-    if (auto prepared = op_manager_->PrepareStash(*estimated); prepared) {
+    if (auto prepared = op_manager_->PrepareStash(estimated->first); prepared) {
       auto [offset, buf] = *prepared;
-      size_t written = Serialize(*value, buf.bytes);
+      size_t written = Serialize(estimated->second, *value, buf.bytes);
       tiering::DiskSegment segment{offset, written};
       op_manager_->Stash(id, segment, buf);
     } else {
@@ -459,7 +470,7 @@ std::optional<util::fb2::Future<bool>> TieredStorage::TryStash(DbIndex dbid, str
   } else if (auto bin = bins_->Stash(dbid, key, SerializeString(*value)); bin) {
     id = bin->first;
     // TODO(vlad): Write bin to prepared buffer instead of allocating one
-    if (auto prepared = op_manager_->PrepareStash(*estimated); prepared) {
+    if (auto prepared = op_manager_->PrepareStash(estimated->first); prepared) {
       auto [offset, buf] = *prepared;
       memcpy(buf.bytes.data(), bin->second.data(), bin->second.size());
       tiering::DiskSegment segment{offset, bin->second.size()};
@@ -654,10 +665,10 @@ size_t TieredStorage::ReclaimMemory(size_t goal) {
                            ->prime.FindFirst(record->key_hash, predicate);
     CHECK(IsValid(it));
     PrimeValue& pv = it->second;
-    tiering::DiskSegment segment = FromCoolItem(pv.GetCool());
 
     // Now the item is only in storage.
-    pv.SetExternal(segment.offset, segment.length, CompactObj::ExternalRep::STRING);
+    tiering::DiskSegment segment = FromCoolItem(pv.GetCool());
+    pv.KeepExternal(segment.offset, segment.length);
 
     auto* stats = op_manager_->GetDbTableStats(record->db_index);
     stats->AddTypeMemoryUsage(record->value.ObjType(), -record->value.MallocUsed());
@@ -673,13 +684,14 @@ bool TieredStorage::ShouldStash(const PrimeValue& pv) const {
     return false;
 
   // Estimate value size
-  optional<size_t> size = EstimateSerializedSize(pv);
-  if (!size)
+  auto estimation = EstimateSerializedSize(pv);
+  if (!estimation)
     return false;
 
   const auto& disk_stats = op_manager_->GetStats().disk_stats;
-  return *size >= config_.min_value_size &&
-         disk_stats.allocated_bytes + tiering::kPageSize + *size < disk_stats.max_file_size;
+  return estimation->first >= config_.min_value_size &&
+         disk_stats.allocated_bytes + tiering::kPageSize + estimation->first <
+             disk_stats.max_file_size;
 }
 
 void TieredStorage::CoolDown(DbIndex db_ind, std::string_view str,

src/server/tiered_storage_test.cc

@@ -454,4 +454,8 @@ TEST_F(PureDiskTSTest, Dump) {
   EXPECT_EQ(resp, "OK");
 }
 
+TEST_F(PureDiskTSTest, Hash1) {
+  Run({"hmset", "k1", "a", "a", "b", "b", "c", "c"});
+}
+
 }  // namespace dfly