feat(tiering): Support for partial offloading of LIST nodes

Signed-off-by: mkaruza <mario@dragonflydb.io>

Author: mkaruza

src/core/compact_object.cc

@@ -821,6 +821,8 @@ CompactObjType CompactObj::ObjType() const {
         return OBJ_STRING;
       case ExternalRep::SERIALIZED_MAP:
         return OBJ_HASH;
+      case ExternalRep::LIST_NODE:
+        return OBJ_LIST;
     };
   }
 

src/core/compact_object.h

@@ -172,8 +172,9 @@ class CompactObj {
 
   // Different representations of external values
   enum class ExternalRep : uint8_t {
-    STRING,         // OBJ_STRING, Basic representation with various string encodings
-    SERIALIZED_MAP  // OBJ_HASH, Serialized map
+    STRING,          // OBJ_STRING, Basic representation with various string encodings
+    SERIALIZED_MAP,  // OBJ_HASH, Serialized map
+    LIST_NODE        // OBJ_LIST, QList::Node
   };
 
   explicit CompactObj(bool is_key)
@@ -452,7 +453,7 @@ class CompactObj {
   }
 
   struct SdsTtlString {
-    char* sds_ptr;    // SDS string (length via sdslen)
+    char* sds_ptr;  // SDS string (length via sdslen)
 
     uint64_t exp_ms;  // absolute expiry time in ms
     std::string_view view() const;

src/core/qlist.cc

@@ -18,6 +18,7 @@ extern "C" {
 
 #include "base/logging.h"
 #include "core/page_usage/page_usage_stats.h"
+#include "core/tiering_types.h"
 
 using namespace std;
 
@@ -49,7 +50,7 @@ namespace dfly {
 namespace {
 
 static_assert(sizeof(QList) == 48);
-static_assert(sizeof(QList::Node) == 40);
+static_assert(sizeof(QList::Node) == 48);
 
 enum IterDir : uint8_t { FWD = 1, REV = 0 };
 
@@ -167,6 +168,8 @@ QList::Node* CreateRAW(int container, uint8_t* entry, size_t sz) {
   node->recompress = 0;
   node->dont_compress = 0;
   node->offloaded = 0;
+  node->io_pending = 0;
+  node->fragment = nullptr;
 
   return node;
 }
@@ -515,6 +518,14 @@ void QList::Clear() noexcept {
 
   while (len_) {
     Node* next = current->next;
+
+    // Clean up offloaded/pending nodes before freeing.
+    if (current->offloaded || current->io_pending) {
+      if (tiering_params_ && tiering_params_->delete_cb) {
+        tiering_params_->delete_cb(current);
+      }
+    }
+
     if (current->encoding != QUICKLIST_NODE_ENCODING_RAW) {
       quicklistLZF* lzf = (quicklistLZF*)current->entry;
       stats.compressed_bytes -= lzf->sz;
@@ -905,7 +916,7 @@ void QList::CoolOff(Node* node, uint32_t node_id) {
     //    off due to merges (can be improved in future).
     if (node_id >= tiering_params_->node_depth_threshold &&
         node_id + tiering_params_->node_depth_threshold < len_) {
-      if (!node->offloaded) {
+      if (!node->offloaded && !node->io_pending) {
         OffloadNode(node);
       }
     } else if (num_offloaded_nodes_ * 2 + tiering_params_->node_depth_threshold * 2 < len_) {
@@ -921,12 +932,12 @@ void QList::CoolOff(Node* node, uint32_t node_id) {
       while (traverse_node_id <= len_ / 2 &&
              (num_offloaded_nodes_ + 2 * tiering_params_->node_depth_threshold) < len_) {
         if (traverse_node_id >= tiering_params_->node_depth_threshold) {
-          if (fw->offloaded == 0) {
+          if (fw->offloaded == 0 && fw->io_pending == 0) {
             OffloadNode(fw);
           }
 
           // Avoid offloading the same node twice when fw and rev meet in the middle.
-          if (rev != fw && rev->offloaded == 0) {
+          if (rev != fw && rev->offloaded == 0 && rev->io_pending == 0) {
             OffloadNode(rev);
           }
         }
@@ -994,12 +1005,28 @@ void QList::CompressByDepth(Node* node) {
 void QList::AccessForReads(bool recompress, Node* node) {
   DCHECK(node);
   stats.total_node_reads++;
+
+  if (node->io_pending) {
+    // A stash is in flight — cancel it. The delete_cb handles cancellation.
+    DCHECK(tiering_params_);
+    if (tiering_params_->delete_cb) {
+      tiering_params_->delete_cb(node);
+    }
+    DCHECK(!node->io_pending);
+  }
+
   if (node->offloaded) {
     DCHECK(tiering_params_);
     stats.onload_requests++;
+    if (tiering_params_->onload_cb) {
+      tiering_params_->onload_cb(node);
+    }
+    // After onload_cb returns, the node entry must be restored.
+    DCHECK(!node->offloaded);
+    DCHECK(node->entry != nullptr);
     num_offloaded_nodes_--;
-    node->offloaded = 0;
   }
+
   if (len_ > 2 && node != head_ && node->next != nullptr) {
     stats.interior_node_reads++;
   }
@@ -1121,8 +1148,16 @@ void QList::DelNode(Node* node) {
   len_--;
   count_ -= node->count;
   malloc_size_ -= node->sz;
-  if (node->offloaded) {
-    num_offloaded_nodes_--;
+
+  if (node->offloaded || node->io_pending) {
+    if (node->offloaded)
+      num_offloaded_nodes_--;
+    else if (node->io_pending)
+      num_offloaded_nodes_--;  // was pre-counted in OffloadNode
+    // Clean up disk segment and deregister fragment.
+    if (tiering_params_ && tiering_params_->delete_cb) {
+      tiering_params_->delete_cb(node);
+    }
   }
 
   /* If we deleted a node within our compress depth, we
@@ -1158,9 +1193,16 @@ bool QList::DelPackedIndex(Node* node, uint8_t* p) {
 
 void QList::OffloadNode(Node* node) {
   DCHECK(tiering_params_ && node->offloaded == 0);
+  if (node->io_pending || node->IsCompressed())
+    return;
+
   num_offloaded_nodes_++;
   stats.offload_requests++;
-  node->offloaded = 1;
+  if (tiering_params_->offload_cb) {
+    tiering_params_->offload_cb(node);
+  } else {
+    node->offloaded = 1;
+  }
 }
 
 void QList::InitIteratorEntry(Iterator* it) const {
@@ -1387,6 +1429,16 @@ uint8_t* QList::TryExtractListpack() {
   return res;
 }
 
+void QList::Node::SetExternal(tiering::Fragment* f) {
+  // Update malloc size for parent QList node
+  DCHECK(f->GetOpaque());
+  reinterpret_cast<QList*>(f->GetOpaque())->AdjustMallocSize(-sz);
+  zfree(entry);
+  entry = nullptr;
+  fragment = f;
+  offloaded = 1;
+}
+
 bool QList::Iterator::Next() {
   if (!current_)
     return false;

src/core/qlist.h

@@ -7,6 +7,7 @@
 #include <absl/functional/function_ref.h>
 
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <string>
 
@@ -26,6 +27,10 @@
 
 namespace dfly {
 
+namespace tiering {
+class Fragment;
+}
+
 class PageUsage;
 
 // Heuristic: for values smaller than 2 KiB we prefer the compact listpack
@@ -64,15 +69,26 @@ class QList {
     uint16_t attempted_compress : 1; /* node can't compress; too small */
     uint16_t dont_compress : 1;      /* prevent compression of entry that will be used later */
     uint16_t offloaded : 1;          /* node is offloaded to colder storage */
-    uint16_t reserved1 : 8;          /* reserved for future use */
-
-    uint16_t reserved2; /* more bits to steal for future usage */
-    uint32_t reserved3; /* more bits to steal for future usage */
+    uint16_t io_pending : 1;         /* stash I/O is in flight for this node */
+    uint16_t reserved1 : 7;          /* reserved for future use */
 
     bool IsCompressed() const {
       return encoding != QUICKLIST_NODE_ENCODING_RAW;
     }
 
+    tiering::Fragment* fragment = nullptr;
+
+    tiering::Fragment* GetFragment() const {
+      return fragment;
+    }
+
+    void RemoveExternal() {
+      fragment = nullptr;
+      offloaded = 0;
+    }
+
+    void SetExternal(tiering::Fragment* f);
+
     size_t GetLZF(void** data) const;
   };
 
@@ -103,9 +119,18 @@ class QList {
   using IterateFunc = absl::FunctionRef<bool(Entry)>;
   enum InsertOpt : uint8_t { BEFORE, AFTER };
 
+  void AdjustMallocSize(size_t delta) {
+    malloc_size_ += delta;
+  }
+
   struct TieringParams {
-    // TODO: hook functions and params that allow qlist offloading nodes to colder storage.
     uint32_t node_depth_threshold = 2;
+    // Called when a node should be offloaded to disk.
+    std::function<void(Node*)> offload_cb;
+    // Called when an offloaded node needs its data loaded back from disk.
+    std::function<void(Node*)> onload_cb;
+    // Called when an offloaded or io_pending node is being deleted.
+    std::function<void(Node*)> delete_cb;
   };
 
   /**

src/core/tiering_types.cc

@@ -5,32 +5,51 @@
 #include "core/tiering_types.h"
 
 #include "core/compact_object.h"
+#include "core/overloaded.h"
 #include "redis/redis_aux.h"
 
 namespace dfly::tiering {
 
 bool Fragment::IsExternal() const {
-  return std::visit([](CompactValue* pv) { return pv->IsExternal(); }, val_);
+  return std::visit(Overloaded{[](CompactValue* pv) { return pv->IsExternal(); },
+                               [](QList::Node* node) { return node->offloaded != 0; }},
+                    val_);
 }
 
 void Fragment::RemoveExternal() {
-  std::visit([](CompactValue* pv) { pv->RemoveExternal(); }, val_);
+  std::visit(Overloaded{[](CompactValue* pv) { pv->RemoveExternal(); },
+                        [](QList::Node* node) { node->RemoveExternal(); }},
+             val_);
 }
 
 void Fragment::SetExternal() {
-  std::visit([this](CompactValue* pv) { pv->SetExternal(this); }, val_);
+  std::visit(Overloaded{[this](CompactValue* pv) { pv->SetExternal(this); },
+                        [this](QList::Node* node) { node->SetExternal(this); }},
+             val_);
 }
 
 bool Fragment::HasStashPending() const {
-  return std::visit([](CompactValue* pv) { return pv->HasStashPending(); }, val_);
+  return std::visit(Overloaded{[](CompactValue* pv) { return pv->HasStashPending(); },
+                               [](QList::Node* node) { return node->io_pending != 0; }},
+                    val_);
 }
 
 void Fragment::SetStashPending(bool b) {
-  std::visit([b](CompactValue* pv) { pv->SetStashPending(b); }, val_);
+  std::visit(Overloaded{[b](CompactValue* pv) { pv->SetStashPending(b); },
+                        [b](QList::Node* node) { node->io_pending = b ? 1 : 0; }},
+             val_);
 }
 
 CompactObjType Fragment::ObjType() const {
-  return std::visit([](CompactValue* pv) { return pv->ObjType(); }, val_);
+  return std::visit(Overloaded{[](CompactValue* pv) { return pv->ObjType(); },
+                               [](QList::Node*) -> CompactObjType { return OBJ_LIST; }},
+                    val_);
+}
+
+size_t Fragment::MallocUsed() const {
+  return std::visit(Overloaded{[](CompactValue* pv) { return pv->MallocUsed(); },
+                               [](QList::Node* node) { return node->sz; }},
+                    val_);
 }
 
 auto Fragment::GetDescr(const CompactValue* pv) -> SerializationDescr {
@@ -52,8 +71,18 @@ auto Fragment::GetDescr(const CompactValue* pv) -> SerializationDescr {
   };
 }
 
+auto Fragment::GetDescr(const QList::Node* node) -> SerializationDescr {
+  if (!node->entry || node->sz == 0)
+    return {};
+
+  std::string_view entry_view{reinterpret_cast<const char*>(node->entry), node->sz};
+  return {std::array<std::string_view, 2>{entry_view, {}}, CompactObj::ExternalRep::LIST_NODE};
+}
+
 auto Fragment::GetSerializationDescr() const -> SerializationDescr {
-  return std::visit([](CompactValue* pv) { return GetDescr(pv); }, val_);
+  return std::visit(Overloaded{[](CompactValue* pv) { return GetDescr(pv); },
+                               [](QList::Node* node) { return GetDescr(node); }},
+                    val_);
 }
 
 std::pair<size_t, size_t> Fragment::GetExternalSlice() const {

src/core/tiering_types.h

@@ -7,6 +7,7 @@
 #include <boost/intrusive/list_hook.hpp>
 
 #include "core/compact_object.h"
+#include "core/qlist.h"
 
 namespace dfly::tiering {
 
@@ -32,14 +33,17 @@ class Fragment {
     CompactObj::ExternalRep rep = CompactObj::ExternalRep::STRING;
   };
 
-  using FragmentType = std::variant<CompactValue*>;
+  using FragmentType = std::variant<CompactValue*, QList::Node*>;
 
   Fragment(CompactValue& pv) : val_(&pv) {  // NOLINT
   }
 
   Fragment(CompactValue* pv) : val_(pv) {  // NOLINT
   }
 
+  Fragment(QList::Node* node) : val_(node) {  // NOLINT
+  }
+
   bool IsExternal() const;
   void RemoveExternal();
   void SetExternal();
@@ -48,6 +52,7 @@ class Fragment {
   void SetStashPending(bool b);
 
   CompactObjType ObjType() const;
+  size_t MallocUsed() const;
 
   // Determine required byte size and encoding type based on value.
   SerializationDescr GetSerializationDescr() const;
@@ -91,6 +96,14 @@ class Fragment {
     val_ = pv;
   }
 
+  void SetOpaque(void* opaque) {
+    opaque_ = opaque;
+  }
+
+  void* GetOpaque() {
+    return opaque_;
+  }
+
   void SetId(size_t id) {
     id_ = id;
   }
@@ -101,6 +114,7 @@ class Fragment {
 
  private:
   static SerializationDescr GetDescr(const CompactValue* pv);
+  static SerializationDescr GetDescr(const QList::Node* node);
 
   size_t id_ = 0;
 
@@ -113,8 +127,9 @@ class Fragment {
   uint32_t serialized_size_ = 0;
   size_t offset_ = 0;
 
-  // First byte of the value if Huffman encoded
-  uint8_t first_byte_ = 0;
+  uint8_t first_byte_ = 0;  // First byte of the value if Huffman encoded
+
+  void* opaque_;  // Auxilary structure that can be attached to fragment value object
 
   FragmentType val_;
 };