Remove callbacks and move tiering responsibility to qlist

Author: mkaruza

src/core/qlist.cc

@@ -459,15 +459,12 @@ size_t QList::DefragIfNeeded(PageUsage* page_usage) {
   return reallocated;
 }
 
-void QList::SetTieringParams(const TieringParams& params) {
-  tiering_params_ = make_unique<TieringParams>(params);
-}
-
 QList::QList(int fill, int compress)
     : fill_(fill),
       dict_learning_failed_(0),
       dict_compress_failed_(0),
       dict_bulk_finished_(0),
+      tiering_enabled_(0),
       compress_(compress),
       bookmark_count_(0) {
 }
@@ -480,10 +477,10 @@ QList::QList(QList&& other) noexcept
       dict_learning_failed_(other.dict_learning_failed_),
       dict_compress_failed_(other.dict_compress_failed_),
       dict_bulk_finished_(other.dict_bulk_finished_),
+      tiering_enabled_(other.tiering_enabled_),
       compress_(other.compress_),
       bookmark_count_(other.bookmark_count_),
-      num_offloaded_nodes_(other.num_offloaded_nodes_),
-      tiering_params_(std::move(other.tiering_params_)) {
+      num_offloaded_nodes_(other.num_offloaded_nodes_) {
   other.head_ = nullptr;
   other.len_ = other.count_ = 0;
   other.num_offloaded_nodes_ = 0;
@@ -503,9 +500,9 @@ QList& QList::operator=(QList&& other) noexcept {
     dict_learning_failed_ = other.dict_learning_failed_;
     dict_compress_failed_ = other.dict_compress_failed_;
     dict_bulk_finished_ = other.dict_bulk_finished_;
+    tiering_enabled_ = other.tiering_enabled_;
     compress_ = other.compress_;
     bookmark_count_ = other.bookmark_count_;
-    tiering_params_ = std::move(other.tiering_params_);
     num_offloaded_nodes_ = other.num_offloaded_nodes_;
     other.head_ = nullptr;
     other.len_ = other.count_ = other.num_offloaded_nodes_ = 0;
@@ -521,10 +518,8 @@ void QList::Clear() noexcept {
 
     // If entry is offloaded we should skip freeing its memory.
     bool free_entry = current->offloaded == 0;
-    if (current->offloaded || current->io_pending) {
-      if (tiering_params_ && tiering_params_->delete_cb) {
-        tiering_params_->delete_cb(current);
-      }
+    if (tiering_enabled_ && (current->offloaded || current->io_pending)) {
+      CleanupOffloadedNode(current);
     } else {
       if (current->encoding != QUICKLIST_NODE_ENCODING_RAW) {
         quicklistLZF* lzf = (quicklistLZF*)current->entry;
@@ -919,7 +914,7 @@ void QList::Replace(Iterator it, std::string_view elem) {
 }
 
 void QList::CoolOff(Node* node, uint32_t node_id) {
-  if (tiering_params_) {
+  if (tiering_enabled_) {
     // Dry run for offloading decision.
     // a. Node id is withing the offloadable depth - offload it if not already offloaded.
     // b. Node id is outside the offloadable depth - but we have too many nodes that are not
@@ -930,12 +925,12 @@ void QList::CoolOff(Node* node, uint32_t node_id) {
     //    we won't need to traverse them again for "trivial" access patterns unless they
     //    get accessed again. Another reason for missing offloaded nodes is that node_id can be
     //    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_id >= tiering_node_depth_threshold_ &&
+        node_id + tiering_node_depth_threshold_ < len_) {
       if (!node->offloaded && !node->io_pending) {
         OffloadNode(node);
       }
-    } else if (num_offloaded_nodes_ * 2 + tiering_params_->node_depth_threshold * 2 < len_) {
+    } else if (num_offloaded_nodes_ * 2 + tiering_node_depth_threshold_ * 2 < len_) {
       // We check `num_offloaded_nodes_ * 2` above to avoid frequent traversals.
       // So only when the gap between offloaded and non-offloaded nodes is large enough,
       // we do a traversal to offload more nodes.
@@ -946,8 +941,8 @@ void QList::CoolOff(Node* node, uint32_t node_id) {
       // Traverse from both ends towards the middle as we expect more offloads towards the ends
       // due to usual access patterns of adding items via lpush/rpush.
       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) {
+             (num_offloaded_nodes_ + 2 * tiering_node_depth_threshold_) < len_) {
+        if (traverse_node_id >= tiering_node_depth_threshold_) {
           if (fw->offloaded == 0 && fw->io_pending == 0) {
             OffloadNode(fw);
           }
@@ -1047,18 +1042,17 @@ void QList::CompressByDepth(Node* node) {
 }
 
 void QList::Materialize(Node* node) {
-  if (!tiering_params_ || (!node->offloaded && !node->io_pending))
+  if (!tiering_enabled_ || (!node->offloaded && !node->io_pending))
     return;
 
   // Cancel stash in progress before loading.
-  if (node->io_pending && tiering_params_->delete_cb) {
-    tiering_params_->delete_cb(node);
+  if (node->io_pending) {
+    CleanupOffloadedNode(node);
   }
 
   // Load the offloaded node data back into memory.
-  if (node->offloaded && tiering_params_->onload_cb) {
-    stats.onload_requests++;
-    tiering_params_->onload_cb(node);
+  if (node->offloaded) {
+    ReadOffloadedNode(node);
   }
 
   DCHECK(!node->offloaded);
@@ -1198,10 +1192,8 @@ void QList::DelNode(Node* node) {
     malloc_size_ -= node->sz;
   }
 
-  if (tiering_params_ && (node->offloaded || node->io_pending)) {
-    if (tiering_params_->delete_cb) {
-      tiering_params_->delete_cb(node);
-    }
+  if (tiering_enabled_ && (node->offloaded || node->io_pending)) {
+    CleanupOffloadedNode(node);
   }
 
   /* If we deleted a node within our compress depth, we
@@ -1249,12 +1241,18 @@ bool QList::DelPackedIndex(Node* node, uint8_t* p) {
   return false;
 }
 
-void QList::OffloadNode(Node* node) {
-  DCHECK(tiering_params_ && node->offloaded == 0 && node->io_pending == 0);
+void QList::OffloadNode(Node* node) const {
+  DCHECK(tiering_enabled_ && node->offloaded == 0 && node->io_pending == 0);
   stats.offload_requests++;
-  if (tiering_params_->offload_cb) {
-    tiering_params_->offload_cb(node);
-  }
+  node->io_pending = 1;
+}
+
+void QList::ReadOffloadedNode(QList::Node* node) const {
+  stats.onload_requests++;
+}
+
+void QList::CleanupOffloadedNode(QList::Node* node) const {
+  node->io_pending = 0;
 }
 
 void QList::InitIteratorEntry(Iterator* it) const {

src/core/qlist.h

@@ -126,17 +126,6 @@ class QList {
     num_offloaded_nodes_ += delta;
   }
 
-  struct TieringParams {
-    uint32_t node_depth_threshold = 2;
-    // Called when a node should be offloaded.
-    // True if node meets criteria for offloading and stashing was initiated.
-    std::function<bool(Node*)> offload_cb;
-    // Called when an offloaded node needs its data loaded back into memory
-    std::function<void(Node*)> onload_cb;
-    // Called when an offloaded or io_pending node is being deleted.
-    std::function<void(Node*)> delete_cb;
-  };
-
   /**
    * fill: The number of entries allowed per internal list node can be specified
    * as a fixed maximum size or a maximum number of elements.
@@ -265,14 +254,18 @@ class QList {
   // Returns count of nodes reallocated to help in testing.
   size_t DefragIfNeeded(PageUsage* page_usage);
 
-  void SetTieringParams(const TieringParams& params);
-
   // Sets the malloc_size_ threshold at which ZSTD dictionary training is triggered.
   // 0 disables ZSTD dictionary compression.
   void set_compr_threshold(uint32_t threshold) {
     zstd_threshold_ = threshold;
   }
 
+  // Enable tiered storage and set node depth threshold
+  void EnableTiering(uint32_t threshold) {
+    tiering_enabled_ = 1;
+    tiering_node_depth_threshold_ = threshold;
+  }
+
   struct Stats {
     uint64_t compression_attempts = 0;
 
@@ -346,7 +339,13 @@ class QList {
 
   void DelNode(Node* node);
   bool DelPackedIndex(Node* node, uint8_t* p);
-  void OffloadNode(Node* node);
+
+  // Offload node to tiered storage
+  void OffloadNode(Node* node) const;
+  // Read offloaded node from tiered storage
+  void ReadOffloadedNode(Node* node) const;
+  // Delete offloaded node or cancel offloading of node
+  void CleanupOffloadedNode(Node* node) const;
 
   // Initializes iterator's zi_ to point to the element at offset_.
   // Decompresses the node if needed. Assumes current_ is not null.
@@ -360,13 +359,14 @@ class QList {
   uint16_t dict_learning_failed_ : 1; /* thread-local dict training failed for this list's data */
   uint16_t dict_compress_failed_ : 1; /* compression with thread-local dict failed for this list */
   uint16_t dict_bulk_finished_ : 1;   /* bulk compression done, per-node compression active */
-  uint16_t reserved1_ : 13;
+  uint16_t tiering_enabled_ : 1;      /* tiering storage enabled */
+  uint16_t reserved1_ : 12;
   unsigned compress_ : QL_COMP_BITS; /* depth of end nodes not to compress;0=off */
   unsigned bookmark_count_ : QL_BM_BITS;
   unsigned reserved2_ : 12;
   uint32_t num_offloaded_nodes_ = 0;
-  uint32_t zstd_threshold_ = 0;  // 0 = disabled
-  std::unique_ptr<TieringParams> tiering_params_;
+  uint32_t zstd_threshold_ = 0;                // 0 = disabled
+  uint32_t tiering_node_depth_threshold_ = 0;  // 0 = disabled
 };
 
 }  // namespace dfly

src/core/qlist_test.cc

@@ -413,28 +413,13 @@ TEST_F(QListTest, DefragmentListpackCompressed) {
 
 TEST_F(QListTest, Tiering) {
   QList::stats.offload_requests = 0;
+  QList tiered_ql;
 
-  // Simulate tiering setup where we mark offloaded nodes as pending so that
-  // throttling logic works correctly. We need to free memory manually in the
-  // delete callback.
-  QList::TieringParams params{.node_depth_threshold = 1,
-                              .offload_cb =
-                                  [](QList::Node* node) {
-                                    node->io_pending = 1;
-                                    return true;
-                                  },
-                              .onload_cb = nullptr,
-                              .delete_cb =
-                                  [](QList::Node* node) {
-                                    zfree(node->entry);
-                                    node->entry = nullptr;
-                                    node->io_pending = 0;
-                                  }};
-
-  ql_.SetTieringParams(params);
+  // Enable tiering and set node_depth_threshold = 1
+  tiered_ql.EnableTiering(1);
 
   for (int i = 0; i < 8000; i++) {
-    ql_.Push(absl::StrCat("value", i), QList::TAIL);
+    tiered_ql.Push(absl::StrCat("value", i), QList::TAIL);
   }
 
   EXPECT_EQ(QList::stats.offload_requests, 9);

src/server/list_family.cc

@@ -26,7 +26,6 @@ extern "C" {
 #include "server/error.h"
 #include "server/family_utils.h"
 #include "server/namespaces.h"
-#include "server/tiered_storage.h"
 #include "server/transaction.h"
 
 /**
@@ -101,17 +100,15 @@ class ListWrapper {
     if (ShouldStoreAsListPack(sz + additional_size)) {
       return nullptr;
     }
+
     QList* ql = CompactObj::AllocateMR<QList>(GetFlag(FLAGS_list_max_listpack_size),
                                               GetFlag(FLAGS_list_compress_depth));
 
-    if (GetFlag(FLAGS_list_tiering_threshold) > 0 && EngineShard::tlocal()->tiered_storage()) {
-      QList::TieringParams params{.node_depth_threshold = GetFlag(FLAGS_list_tiering_threshold),
-                                  .offload_cb = nullptr,
-                                  .onload_cb = nullptr,
-                                  .delete_cb = nullptr};
-
-      ql->SetTieringParams(params);
+    const uint32_t tiering_node_depth_threshold = absl::GetFlag(FLAGS_list_tiering_threshold);
+    if (tiering_node_depth_threshold > 0 && EngineShard::tlocal()->tiered_storage()) {
+      ql->EnableTiering(tiering_node_depth_threshold);
     }
+
     if (uint32_t zstd_thresh = GetFlag(FLAGS_list_experimental_zstd_dict_threshold);
         zstd_thresh > 0) {
       ql->set_compr_threshold(zstd_thresh);