feat(tiering): Decoders

Signed-off-by: Vladislav Oleshko <[email protected]>

Author: dranikpg

src/core/string_or_view.h

@@ -65,11 +65,17 @@ class StringOrView {
       val_ = std::string{std::get<std::string_view>(val_)};
   }
 
+  // Move out of value as string
   std::string Take() && {
     MakeOwned();
     return std::move(std::get<std::string>(val_));
   }
 
+  std::string* BorrowMut() {
+    MakeOwned();
+    return &std::get<std::string>(val_);
+  }
+
   bool empty() const {
     return visit([](const auto& s) { return s.empty(); }, val_);
   }

src/server/CMakeLists.txt

@@ -26,7 +26,7 @@ set_property(SOURCE dfly_main.cc APPEND PROPERTY COMPILE_DEFINITIONS
 
 if (WITH_TIERING)
     SET(TX_LINUX_SRCS tiering/disk_storage.cc tiering/op_manager.cc tiering/small_bins.cc
-      tiering/external_alloc.cc)
+      tiering/external_alloc.cc tiering/decoders.cc)
 
     add_executable(dfly_bench dfly_bench.cc)
     cxx_link(dfly_bench dfly_parser_lib fibers2 absl::random_random redis_lib)

src/server/tiered_storage.cc

@@ -123,8 +123,7 @@ class TieredStorage::ShardOpManager : public tiering::OpManager {
     }
   }
 
-  bool NotifyFetched(EntryId id, string_view value, tiering::DiskSegment segment,
-                     bool modified) override;
+  bool NotifyFetched(EntryId id, tiering::DiskSegment segment, tiering::Decoder* decoder) override;
 
   bool NotifyDelete(tiering::DiskSegment segment) override;
 
@@ -210,12 +209,12 @@ void TieredStorage::ShardOpManager::Defragment(tiering::DiskSegment segment, str
   }
 }
 
-bool TieredStorage::ShardOpManager::NotifyFetched(EntryId id, string_view value,
-                                                  tiering::DiskSegment segment, bool modified) {
+bool TieredStorage::ShardOpManager::NotifyFetched(EntryId id, tiering::DiskSegment segment,
+                                                  tiering::Decoder* decoder) {
   ++stats_.total_fetches;
 
   if (id == EntryId{kFragmentedBin}) {  // Generally we read whole bins only for defrag
-    Defragment(segment, value);
+    // Defragment(segment, value);
     return true;  // delete
   }
 
@@ -224,20 +223,20 @@ bool TieredStorage::ShardOpManager::NotifyFetched(EntryId id, string_view value,
   //    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 = modified;
-  should_upload |=
-      (ts_->UploadBudget() > int64_t(value.length())) && !SliceSnapshot::IsSnaphotInProgress();
+  bool should_upload = decoder->modified;
+  should_upload |= (ts_->UploadBudget() > int64_t(decoder->estimated_mem_usage)) &&
+                   !SliceSnapshot::IsSnaphotInProgress();
 
   if (!should_upload)
     return false;
 
   auto key = get<OpManager::KeyRef>(id);
   auto* pv = Find(key);
   if (pv && pv->IsExternal() && segment == pv->GetExternalSlice()) {
-    if (modified || pv->WasTouched()) {
-      bool is_raw = !modified;
+    if (decoder->modified || pv->WasTouched()) {
       ++stats_.total_uploads;
-      Upload(key.first, value, is_raw, segment.length, pv);
+      decoder->Upload(pv);
+      // Upload(key.first, value, is_raw, segment.length, pv);
       return true;
     }
     pv->SetTouched(true);
@@ -262,11 +261,11 @@ bool TieredStorage::ShardOpManager::NotifyDelete(tiering::DiskSegment segment) {
   if (bin.fragmented) {
     // Trigger read to signal need for defragmentation. NotifyFetched will handle it.
     DVLOG(2) << "Enqueueing bin defragmentation for: " << bin.segment.offset;
-    auto cb = [dummy = 5](auto res) -> bool {
-      (void)dummy;  // a hack to make cb non constexpr that confuses some old) compilers.
-      return false;
-    };
-    Enqueue(kFragmentedBin, bin.segment, std::move(cb));
+    // auto cb = [dummy = 5](auto res) -> bool {
+    //   (void)dummy;  // a hack to make cb non constexpr that confuses some old) compilers.
+    //   return false;
+    // };
+    //  Enqueue(kFragmentedBin, bin.segment, std::move(cb));
   }
 
   return false;
@@ -324,14 +323,11 @@ void TieredStorage::Read(DbIndex dbid, std::string_view key, const PrimeValue& v
                          std::function<void(io::Result<std::string>)> readf) {
   DCHECK(value.IsExternal());
   DCHECK(!value.IsCool());
-  auto cb = [readf = std::move(readf), enc = value.GetStrEncoding()](auto res) mutable {
-    readf(res.transform([enc](tiering::OpManager::FetchedEntry entry) {
-      auto [ptr, raw] = entry;
-      return raw ? enc.Decode(*ptr).Take() : *ptr;  // TODO(vlad): optimize last value copy
-    }));
-    return false;
+  auto cb = [readf = std::move(readf)](io::Result<tiering::StringDecoder*> res) mutable {
+    readf(res.transform([](auto* d) { return string{d->Read()}; }));
   };
-  op_manager_->Enqueue(KeyRef(dbid, key), value.GetExternalSlice(), std::move(cb));
+  op_manager_->Enqueue(KeyRef(dbid, key), value.GetExternalSlice(), tiering::StringDecoder{value},
+                       std::move(cb));
 }
 
 template <typename T>
@@ -341,21 +337,11 @@ TieredStorage::TResult<T> TieredStorage::Modify(DbIndex dbid, std::string_view k
   DCHECK(value.IsExternal());
 
   util::fb2::Future<io::Result<T>> future;
-  auto cb = [future, modf = std::move(modf), enc = value.GetStrEncoding()](auto res) mutable {
-    if (!res.has_value()) {
-      future.Resolve(res.get_unexpected());
-      return false;
-    }
-
-    auto [raw_val, is_raw] = *res;
-    if (is_raw) {
-      raw_val->resize(enc.DecodedSize(*raw_val));
-      enc.Decode(*raw_val, raw_val->data());
-    }
-    future.Resolve(modf(raw_val));
-    return true;
+  auto cb = [future, modf = std::move(modf)](io::Result<tiering::StringDecoder*> res) mutable {
+    future.Resolve(res.transform([&modf](auto* d) { return modf(d->Write()); }));
   };
-  op_manager_->Enqueue(KeyRef(dbid, key), value.GetExternalSlice(), std::move(cb));
+  op_manager_->Enqueue(KeyRef(dbid, key), value.GetExternalSlice(), tiering::StringDecoder{value},
+                       std::move(cb));
   return future;
 }
 

src/server/tiering/decoders.cc

@@ -0,0 +1,37 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "server/tiering/decoders.h"
+
+namespace dfly::tiering {
+
+StringDecoder::StringDecoder(const CompactObj& obj) : StringDecoder{obj.GetStrEncoding()} {
+}
+
+StringDecoder::StringDecoder(CompactObj::StrEncoding encoding) : encoding_{encoding} {
+}
+
+std::unique_ptr<Decoder> StringDecoder::Clone() const {
+  return std::unique_ptr<StringDecoder>{new StringDecoder(encoding_)};
+}
+
+void StringDecoder::Initialize(std::string_view slice) {
+  value_ = encoding_.Decode(slice);
+  estimated_mem_usage = value_.view().length();
+}
+
+void StringDecoder::Upload(CompactObj* obj) {
+  obj->Materialize(value_.view(), false);
+}
+
+std::string_view StringDecoder::Read() const {
+  return value_.view();
+}
+
+std::string* StringDecoder::Write() {
+  modified = true;
+  return value_.BorrowMut();
+}
+
+}  // namespace dfly::tiering

src/server/tiering/decoders.h

@@ -0,0 +1,52 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <string_view>
+
+#include "core/compact_object.h"
+#include "core/string_or_view.h"
+
+namespace dfly::tiering {
+
+// Decodes serialized value and provides it to callbacks.
+// Acts as generic interface to callback driver (OpManager)
+struct Decoder {
+  virtual ~Decoder() = default;
+
+  // Poor man's type-erasure copy
+  virtual std::unique_ptr<Decoder> Clone() const = 0;
+
+  // Initialize decoder from slice
+  virtual void Initialize(std::string_view slice) = 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
+};
+
+struct StringDecoder : public Decoder {
+  explicit StringDecoder(const CompactObj& obj);
+
+  std::unique_ptr<Decoder> Clone() const override;
+  void Initialize(std::string_view slice) override;
+  void Upload(CompactObj* obj) override;
+
+  std::string_view Read() const;
+  std::string* Write();
+
+ private:
+  explicit StringDecoder(CompactObj::StrEncoding encoding);
+
+  CompactObj::StrEncoding encoding_;
+  dfly::StringOrView value_;
+};
+
+}  // namespace dfly::tiering

src/server/tiering/op_manager.cc

@@ -58,10 +58,11 @@ void OpManager::Close() {
   DCHECK(pending_reads_.empty());
 }
 
-void OpManager::Enqueue(EntryId id, DiskSegment segment, ReadCallback cb) {
+void OpManager::EnqueueInternal(EntryId id, DiskSegment segment, const Decoder& decoder,
+                                ReadCallback cb) {
   // Fill pages for prepared read as it has no penalty and potentially covers more small segments
   PrepareRead(segment.ContainingPages())
-      .ForSegment(segment, id)
+      .ForSegment(segment, id, decoder)
       .callbacks.emplace_back(std::move(cb));
 }
 
@@ -147,16 +148,15 @@ void OpManager::ProcessRead(size_t offset, io::Result<std::string_view> page) {
   }
 
   bool deleting_full = false;
-  std::string key_value;
 
   // Notify functions in the loop may append items to info->key_ops during the traversal
   for (size_t i = 0; i < info->key_ops.size(); i++) {
-    bool modified = false;
     auto& ko = info->key_ops[i];
     if (page) {
-      key_value = page->substr(ko.segment.offset - info->segment.offset, ko.segment.length);
+      size_t offset = ko.segment.offset - info->segment.offset;
+      ko.decoder->Initialize(page->substr(offset, ko.segment.length));
       for (auto& cb : ko.callbacks)
-        modified |= cb(std::pair{&key_value, !modified});
+        cb(&*ko.decoder);
     } else {
       for (auto& cb : ko.callbacks)
         cb(page.get_unexpected());
@@ -167,7 +167,7 @@ void OpManager::ProcessRead(size_t offset, io::Result<std::string_view> page) {
     // If the item is not being deleted, report is as fetched to be cached potentially.
     // In case it's cached, we might need to delete it.
     if (page.has_value() && !delete_from_storage)
-      delete_from_storage |= NotifyFetched(Borrowed(ko.id), key_value, ko.segment, modified);
+      delete_from_storage |= NotifyFetched(Borrowed(ko.id), ko.segment, &*ko.decoder);
 
     // If the item is being deleted, check if the full page needs to be deleted.
     if (delete_from_storage)
@@ -181,15 +181,22 @@ void OpManager::ProcessRead(size_t offset, io::Result<std::string_view> page) {
   pending_reads_.erase(offset);
 }
 
-OpManager::EntryOps& OpManager::ReadOp::ForSegment(DiskSegment key_segment, EntryId id) {
+OpManager::EntryOps::EntryOps(OwnedEntryId id, DiskSegment segment, const Decoder& decoder)
+    : id{std::move(id)}, segment{segment}, decoder{decoder.Clone()} {
+}
+
+OpManager::EntryOps& OpManager::ReadOp::ForSegment(DiskSegment key_segment, EntryId id,
+                                                   const Decoder& decoder) {
   DCHECK_GE(key_segment.offset, segment.offset);
   DCHECK_LE(key_segment.length, segment.length);
 
   for (auto& ops : key_ops) {
-    if (ops.segment.offset == key_segment.offset)
+    if (ops.segment.offset == key_segment.offset) {
+      DCHECK(typeid(*ops.decoder) == typeid(decoder));
       return ops;
+    }
   }
-  return key_ops.emplace_back(ToOwned(id), key_segment);
+  return key_ops.emplace_back(ToOwned(id), key_segment, decoder);
 }
 
 OpManager::EntryOps* OpManager::ReadOp::Find(DiskSegment key_segment) {

src/server/tiering/op_manager.h

@@ -10,6 +10,7 @@
 #include <variant>
 
 #include "server/tiering/common.h"
+#include "server/tiering/decoders.h"
 #include "server/tiering/disk_storage.h"
 #include "server/tx_base.h"
 #include "util/fibers/future.h"
@@ -36,19 +37,6 @@ class OpManager {
   using EntryId = std::variant<unsigned, KeyRef>;
   using OwnedEntryId = std::variant<unsigned, std::pair<DbIndex, std::string>>;
 
-  using FetchedEntry = std::pair<std::string* /* ptr */, bool /* raw */>;
-
-  // Callback for post-read completion. Returns whether the value was modified.
-  // We use fu2 function to allow moveable semantics. The arguments are:
-  // bool - true if the string is raw as it was extracted from the prime value.
-  // string* - the string that may potentially be modified by the callbacks that subsribed to this
-  //           read. The callback run in the same order as the order of invocation, guaranteeing
-  //           consistent read after modifications.
-  using ReadCallback =
-      fu2::function_base<true /*owns*/, false /*moveable*/, fu2::capacity_fixed<40, 8>,
-                         false /* non-throwing*/, false /* strong exceptions guarantees*/,
-                         bool(io::Result<FetchedEntry>)>;
-
   explicit OpManager(size_t max_size);
   virtual ~OpManager();
 
@@ -60,7 +48,14 @@ class OpManager {
   // Enqueue callback to be executed once value is read. Trigger read if none is pending yet for
   // this segment. Multiple entries can be obtained from a single segment, but every distinct id
   // will have it's own independent callback loop that can safely modify the underlying value
-  void Enqueue(EntryId id, DiskSegment segment, ReadCallback cb);
+  template <typename D, typename F>
+  void Enqueue(EntryId id, DiskSegment segment, const D& decoder, F&& cb) {
+    static_assert(std::is_base_of_v<Decoder, D>);
+    auto erased_cb = [cb = std::forward<F>(cb)](io::Result<Decoder*> res) mutable {
+      cb(res.transform([](Decoder* ptr) { return static_cast<D*>(ptr); }));
+    };
+    EnqueueInternal(id, segment, decoder, std::move(erased_cb));
+  }
 
   // Delete entry with pending io
   void Delete(EntryId id);
@@ -74,27 +69,34 @@ class OpManager {
   Stats GetStats() const;
 
  protected:
+  using ReadCallback =
+      fu2::function_base<true /*owns*/, false /*moveable*/, fu2::capacity_fixed<40, 8>,
+                         false /* non-throwing*/, false /* strong exceptions guarantees*/,
+                         void(io::Result<Decoder*>)>;
+
+  // Type erased continuation of Enqueue
+  void EnqueueInternal(EntryId id, DiskSegment segment, const Decoder& decoder, ReadCallback cb);
+
   // Notify that a stash succeeded and the entry was stored at the provided segment or failed with
   // given error
   virtual void NotifyStashed(EntryId id, const io::Result<DiskSegment>& segment) = 0;
 
   // Notify that an entry was successfully fetched. Includes whether entry was modified.
   // Returns true if value needs to be deleted from the storage.
-  virtual bool NotifyFetched(EntryId id, std::string_view value, DiskSegment segment,
-                             bool modified) = 0;
+  virtual bool NotifyFetched(EntryId id, DiskSegment segment, Decoder*) = 0;
 
   // Notify delete. Return true if the filled segment needs to be marked as free.
   virtual bool NotifyDelete(DiskSegment segment) = 0;
 
  protected:
   // Describes pending futures for a single entry
   struct EntryOps {
-    EntryOps(OwnedEntryId id, DiskSegment segment) : id(std::move(id)), segment(segment) {
-    }
+    EntryOps(OwnedEntryId id, DiskSegment segment, const Decoder& decoder);
 
     OwnedEntryId id;
     DiskSegment segment;
     absl::InlinedVector<ReadCallback, 1> callbacks;
+    std::unique_ptr<Decoder> decoder;
     bool deleting = false;
   };
 
@@ -104,7 +106,7 @@ class OpManager {
     }
 
     // Get ops for id or create new
-    EntryOps& ForSegment(DiskSegment segment, EntryId id);
+    EntryOps& ForSegment(DiskSegment segment, EntryId id, const Decoder& decoder);
 
     // Find if there are operations for the given segment, return nullptr otherwise
     EntryOps* Find(DiskSegment segment);