fix(tiering): Recompute size before upload (#6031)

Author: dranikpg

src/core/compact_object.cc

@@ -1255,8 +1255,7 @@ std::pair<size_t, size_t> CompactObj::GetExternalSlice() const {
 void CompactObj::Materialize(std::string_view blob, bool is_raw) {
   CHECK(IsExternal()) << int(taglen_);
   DCHECK_EQ(u_.ext_ptr.representation, static_cast<uint8_t>(ExternalRep::STRING));
-
-  DCHECK_GT(blob.size(), kInlineLen);
+  DCHECK_GT(blob.size(), kInlineLen);  // There are no mutable commands that shrink strings
 
   if (is_raw) {
     if (kUseSmallStrings && SmallString::CanAllocate(blob.size())) {
@@ -1267,6 +1266,7 @@ void CompactObj::Materialize(std::string_view blob, bool is_raw) {
       u_.r_obj.SetString(blob, tl.local_mr);
     }
   } else {
+    mask_bits_.encoding = NONE_ENC;  // reset encoding
     EncodeString(blob, false);
   }
 }

src/core/compact_object_test.cc

@@ -620,6 +620,12 @@ TEST_F(CompactObjectTest, StrEncodingAndMaterialize) {
       obj.SetExternal(0, 0, CompactObj::ExternalRep::STRING);  // dummy values
       obj.Materialize(raw_str, true);
       EXPECT_EQ(test_str, obj.ToString());
+
+      // Restore from external again, but not as a raw value
+      obj.SetExternal(0, 0, CompactObj::ExternalRep::STRING);
+      auto test_str2 = test_str + "updated";
+      obj.Materialize(test_str2, false);
+      EXPECT_EQ(obj.ToString(), test_str2);
     }
   }
 }

src/server/tiered_storage.cc

@@ -230,13 +230,15 @@ bool TieredStorage::ShardOpManager::NotifyFetched(EntryId id, tiering::DiskSegme
     return true;  // delete
   }
 
+  tiering::Decoder::UploadMetrics metrics = decoder->GetMetrics();
+
   // 1. When modified is true we MUST upload the value back to memory.
   // 2. On the other hand, if read is caused by snapshotting we do not want to fetch it.
   //    Currently, our heuristic is not very smart, because we stop uploading any reads during
   //    the snapshotting.
   // TODO: to revisit this when we rewrite it with more efficient snapshotting algorithm.
-  bool should_upload = decoder->modified;
-  should_upload |= (ts_->UploadBudget() > int64_t(decoder->estimated_mem_usage)) &&
+  bool should_upload = metrics.modified;
+  should_upload |= (ts_->UploadBudget() > int64_t(metrics.estimated_mem_usage)) &&
                    !SliceSnapshot::IsSnaphotInProgress();
 
   if (!should_upload)
@@ -245,7 +247,7 @@ bool TieredStorage::ShardOpManager::NotifyFetched(EntryId id, tiering::DiskSegme
   auto key = get<OpManager::KeyRef>(id);
   auto* pv = Find(key);
   if (pv && pv->IsExternal() && segment == pv->GetExternalSlice()) {
-    if (decoder->modified || pv->WasTouched()) {
+    if (metrics.modified || pv->WasTouched()) {
       ++stats_.total_uploads;
       decoder->Upload(pv);
       RecordDeleted(*pv, segment.length, GetDbTableStats(key.first));

src/server/tiered_storage_test.cc

@@ -67,6 +67,8 @@ class TieredStorageTest : public BaseFamilyTest {
 TEST_F(TieredStorageTest, SimpleGetSet) {
   absl::FlagSaver saver;
   SetFlag(&FLAGS_tiered_offload_threshold, 0.0f);  // disable offloading
+  UpdateFromFlags();
+
   const int kMin = 256;
   const int kMax = tiering::kPageSize + 10;
 
@@ -122,16 +124,58 @@ TEST_F(TieredStorageTest, MGET) {
     EXPECT_EQ(elements[i], values[i]);
 }
 
-TEST_F(TieredStorageTest, SimpleAppend) {
-  // TODO: use pipelines to issue APPEND/GET/APPEND sequence,
-  // currently it's covered only for op_manager_test
-  for (size_t sleep : {0, 100, 500, 1000}) {
-    Run({"SET", "k0", BuildString(3000)});
-    if (sleep)
-      util::ThisFiber::SleepFor(sleep * 1us);
-    EXPECT_THAT(Run({"APPEND", "k0", "B"}), IntArg(3001));
-    ASSERT_EQ(Run({"GET", "k0"}), BuildString(3000) + 'B') << sleep;
+// Issue many APPEND commands to an offloaded value that are executed at once (with CLIENT PAUSE).
+// They should all finish within the same io completion loop.
+TEST_F(TieredStorageTest, AppendStorm) {
+  const size_t kAppends = 20;
+
+  absl::FlagSaver saver;
+  absl::SetFlag(&FLAGS_tiered_offload_threshold, 1.0);
+  absl::SetFlag(&FLAGS_tiered_upload_threshold, 0.0);
+  absl::SetFlag(&FLAGS_tiered_experimental_cooling, false);
+  UpdateFromFlags();
+
+  // Offload single value
+  string base_value(4096, 'a');
+  Run({"SET", "key", base_value});
+  ExpectConditionWithinTimeout([this] { return GetMetrics().tiered_stats.total_stashes == 1; });
+
+  // Accumulate APPENDs
+  Run({"CLIENT", "pause", "1000"});
+  vector<Fiber> fibs;
+  for (size_t i = 0; i < kAppends; i++) {
+    fibs.emplace_back(pp_->at(0)->LaunchFiber([this, i] {
+      Run(absl::StrCat(i), {"APPEND", "key", string(96, 'b')});
+    }));
   }
+
+  // Throw in a SETRANGE
+  fibs.emplace_back(pp_->at(0)->LaunchFiber([this] {
+    Run("range", {"SETRANGE", "key", "0", string(96, 'x')});
+  }));
+
+  // Throw in a GETRANGE to a range that keeps constant
+  string get_range;
+  fibs.emplace_back(pp_->at(0)->LaunchFiber([this, &get_range] {
+    get_range = Run("get", {"GETRANGE", "key", "96", "191"}).GetString();
+  }));
+
+  // Unlock and wait
+  Run({"CLIENT", "unpause"});
+  for (auto& f : fibs)
+    f.JoinIfNeeded();
+
+  // Check partial result is right
+  EXPECT_EQ(get_range, string(96, 'a'));
+
+  // Get value and verify it
+  auto value = Run({"GET", "key"});
+  EXPECT_EQ(value, string(96, 'x') + string(4000, 'a') + string(kAppends * 96, 'b'));
+
+  // Check value was read no more than once for APPENDs and once for GET
+  auto metrics = GetMetrics();
+  EXPECT_LE(metrics.tiered_stats.total_fetches, 2u);
+  EXPECT_LE(metrics.tiered_stats.total_uploads, 2u);
 }
 
 TEST_F(TieredStorageTest, Ranges) {
@@ -207,6 +251,7 @@ TEST_F(TieredStorageTest, Defrag) {
 TEST_F(TieredStorageTest, BackgroundOffloading) {
   absl::FlagSaver saver;
   SetFlag(&FLAGS_tiered_offload_threshold, 1.0f);      // offload all values
+  SetFlag(&FLAGS_tiered_upload_threshold, 0.0f);       // upload all values
   SetFlag(&FLAGS_tiered_experimental_cooling, false);  // The setup works without cooling buffers
   UpdateFromFlags();
 
@@ -247,7 +292,6 @@ TEST_F(TieredStorageTest, BackgroundOffloading) {
   // should be re-stashed again.
   EXPECT_EQ(metrics.tiered_stats.total_stashes, kNum + metrics.tiered_stats.total_uploads)
       << resp.GetString();
-  EXPECT_EQ(metrics.tiered_stats.total_fetches, kNum * 2);
   EXPECT_EQ(metrics.tiered_stats.allocated_bytes, kNum * 4096);
 }
 

src/server/tiering/decoders.cc

@@ -18,6 +18,11 @@ void BareDecoder::Upload(CompactObj* obj) {
   ABSL_UNREACHABLE();
 }
 
+Decoder::UploadMetrics BareDecoder::GetMetrics() const {
+  ABSL_UNREACHABLE();
+  return UploadMetrics{};
+}
+
 StringDecoder::StringDecoder(const CompactObj& obj) : StringDecoder{obj.GetStrEncoding()} {
 }
 
@@ -31,22 +36,28 @@ std::unique_ptr<Decoder> StringDecoder::Clone() const {
 void StringDecoder::Initialize(std::string_view slice) {
   slice_ = slice;
   value_ = encoding_.Decode(slice);
-  estimated_mem_usage = slice.length();  // will be encoded back
 }
 
 void StringDecoder::Upload(CompactObj* obj) {
-  if (modified)
+  if (modified_)
     obj->Materialize(value_.view(), false);
   else
     obj->Materialize(slice_, true);
 }
 
+Decoder::UploadMetrics StringDecoder::GetMetrics() const {
+  return UploadMetrics{
+      .modified = modified_,
+      .estimated_mem_usage = value_.view().size(),
+  };
+}
+
 std::string_view StringDecoder::Read() const {
   return value_.view();
 }
 
 std::string* StringDecoder::Write() {
-  modified = true;
+  modified_ = true;
   return value_.GetMutable();
 }
 

src/server/tiering/decoders.h

@@ -17,6 +17,11 @@ namespace dfly::tiering {
 // Decodes serialized value and provides it to callbacks.
 // Acts as generic interface to callback driver (OpManager)
 struct Decoder {
+  struct UploadMetrics {
+    bool modified;               // whether the value as modified
+    size_t estimated_mem_usage;  // Estimated memory usage if uploaded
+  };
+
   virtual ~Decoder() = default;
 
   // Poor man's type-erasure copy
@@ -25,17 +30,18 @@ struct Decoder {
   // Initialize decoder from slice
   virtual void Initialize(std::string_view slice) = 0;
 
+  // Compute upload metrics to determine if its worth
+  virtual UploadMetrics GetMetrics() const = 0;
+
   // Store value in compact object
   virtual void Upload(CompactObj* obj) = 0;
-
-  bool modified = false;           // Must be set if modified (not equal to original slice)
-  size_t estimated_mem_usage = 0;  // Estimated usage if uploaded
 };
 
 // Basic "bare" decoder that just stores the provided slice
 struct BareDecoder : public Decoder {
   std::unique_ptr<Decoder> Clone() const override;
   void Initialize(std::string_view slice) override;
+  UploadMetrics GetMetrics() const override;
   void Upload(CompactObj* obj) override;
 
   std::string_view slice;
@@ -46,6 +52,7 @@ struct StringDecoder : public Decoder {
 
   std::unique_ptr<Decoder> Clone() const override;
   void Initialize(std::string_view slice) override;
+  UploadMetrics GetMetrics() const override;
   void Upload(CompactObj* obj) override;
 
   std::string_view Read() const;
@@ -54,6 +61,7 @@ struct StringDecoder : public Decoder {
  private:
   explicit StringDecoder(CompactObj::StrEncoding encoding);
 
+  bool modified_;
   std::string_view slice_;
   CompactObj::StrEncoding encoding_;
   dfly::StringOrView value_;