@nxtedition/rocksdb 13.1.4 → 13.2.0

package/deps/rocksdb/rocksdb/db/multi_cf_iterator_impl.h

@@ -21,27 +21,21 @@ struct MultiCfIteratorInfo {
   int order;
 };
 
+template <typename ResetFunc, typename PopulateFunc>
 class MultiCfIteratorImpl {
  public:
   MultiCfIteratorImpl(
-      const Comparator* comparator,
-      const std::vector<ColumnFamilyHandle*>& column_families,
-      const std::vector<Iterator*>& child_iterators,
-      std::function<void()> reset_func,
-      std::function<void(const autovector<MultiCfIteratorInfo>&)> populate_func)
-      : comparator_(comparator),
-        heap_(MultiCfMinHeap(
-            MultiCfHeapItemComparator<std::greater<int>>(comparator_))),
+      const ReadOptions& read_options, const Comparator* comparator,
+      std::vector<std::pair<ColumnFamilyHandle*, std::unique_ptr<Iterator>>>&&
+          cfh_iter_pairs,
+      ResetFunc reset_func, PopulateFunc populate_func)
+      : allow_unprepared_value_(read_options.allow_unprepared_value),
+        comparator_(comparator),
+        cfh_iter_pairs_(std::move(cfh_iter_pairs)),
         reset_func_(std::move(reset_func)),
-        populate_func_(std::move(populate_func)) {
-    assert(column_families.size() > 0 &&
-           column_families.size() == child_iterators.size());
-    cfh_iter_pairs_.reserve(column_families.size());
-    for (size_t i = 0; i < column_families.size(); ++i) {
-      cfh_iter_pairs_.emplace_back(
-          column_families[i], std::unique_ptr<Iterator>(child_iterators[i]));
-    }
-  }
+        populate_func_(std::move(populate_func)),
+        heap_(MultiCfMinHeap(
+            MultiCfHeapItemComparator<std::greater<int>>(comparator_))) {}
   ~MultiCfIteratorImpl() { status_.PermitUncheckedError(); }
 
   // No copy allowed
@@ -101,10 +95,28 @@ class MultiCfIteratorImpl {
     AdvanceIterator(max_heap, [](Iterator* iter) { iter->Prev(); });
   }
 
+  bool PrepareValue() {
+    assert(Valid());
+
+    if (!allow_unprepared_value_) {
+      return true;
+    }
+
+    if (std::holds_alternative<MultiCfMaxHeap>(heap_)) {
+      return PopulateIterator(std::get<MultiCfMaxHeap>(heap_));
+    }
+
+    return PopulateIterator(std::get<MultiCfMinHeap>(heap_));
+  }
+
  private:
+  Status status_;
+  bool allow_unprepared_value_;
+  const Comparator* comparator_;
   std::vector<std::pair<ColumnFamilyHandle*, std::unique_ptr<Iterator>>>
       cfh_iter_pairs_;
-  Status status_;
+  ResetFunc reset_func_;
+  PopulateFunc populate_func_;
 
   template <typename CompareOp>
   class MultiCfHeapItemComparator {
@@ -125,7 +137,7 @@ class MultiCfIteratorImpl {
    private:
     const Comparator* comparator_;
   };
-  const Comparator* comparator_;
+
   using MultiCfMinHeap =
       BinaryHeap<MultiCfIteratorInfo,
                  MultiCfHeapItemComparator<std::greater<int>>>;
@@ -136,9 +148,6 @@ class MultiCfIteratorImpl {
 
   MultiCfIterHeap heap_;
 
-  std::function<void()> reset_func_;
-  std::function<void(autovector<MultiCfIteratorInfo>)> populate_func_;
-
   Iterator* current() const {
     if (std::holds_alternative<MultiCfMaxHeap>(heap_)) {
       auto& max_heap = std::get<MultiCfMaxHeap>(heap_);
@@ -163,11 +172,11 @@ class MultiCfIteratorImpl {
   }
 
   void InitMinHeap() {
-    heap_.emplace<MultiCfMinHeap>(
+    heap_.template emplace<MultiCfMinHeap>(
         MultiCfHeapItemComparator<std::greater<int>>(comparator_));
   }
   void InitMaxHeap() {
-    heap_.emplace<MultiCfMaxHeap>(
+    heap_.template emplace<MultiCfMaxHeap>(
         MultiCfHeapItemComparator<std::less<int>>(comparator_));
   }
 
@@ -186,13 +195,14 @@ class MultiCfIteratorImpl {
         if (!status_.ok()) {
           // Non-OK status from the iterator. Bail out early
           heap.clear();
-          break;
+          return;
         }
       }
       ++i;
     }
-    if (!heap.empty()) {
-      PopulateIterator(heap);
+    if (!allow_unprepared_value_ && !heap.empty()) {
+      [[maybe_unused]] const bool result = PopulateIterator(heap);
+      assert(result || (!Valid() && !status_.ok()));
     }
   }
 
@@ -212,32 +222,41 @@ class MultiCfIteratorImpl {
     // 2. Make sure all others have iterated past the top iterator key slice
     // 3. Advance the top iterator, and add it back to the heap if valid
     auto top = heap.top();
+    assert(top.iterator);
+    assert(top.iterator->Valid());
+    assert(top.iterator->status().ok());
+
     heap.pop();
-    if (!heap.empty()) {
+
+    while (!heap.empty()) {
       auto current = heap.top();
       assert(current.iterator);
-      while (current.iterator->Valid() &&
-             comparator_->Compare(top.iterator->key(),
-                                  current.iterator->key()) == 0) {
+      assert(current.iterator->Valid());
+      assert(current.iterator->status().ok());
+
+      if (comparator_->Compare(current.iterator->key(), top.iterator->key()) !=
+          0) {
+        break;
+      }
+
+      advance_func(current.iterator);
+
+      if (current.iterator->Valid()) {
         assert(current.iterator->status().ok());
-        advance_func(current.iterator);
-        if (current.iterator->Valid()) {
-          heap.replace_top(heap.top());
+        heap.replace_top(current);
+      } else {
+        considerStatus(current.iterator->status());
+        if (!status_.ok()) {
+          heap.clear();
+          return;
         } else {
-          considerStatus(current.iterator->status());
-          if (!status_.ok()) {
-            heap.clear();
-            return;
-          } else {
-            heap.pop();
-          }
-        }
-        if (!heap.empty()) {
-          current = heap.top();
+          heap.pop();
         }
       }
     }
+
     advance_func(top.iterator);
+
     if (top.iterator->Valid()) {
       assert(top.iterator->status().ok());
       heap.push(top);
@@ -249,13 +268,14 @@ class MultiCfIteratorImpl {
       }
     }
 
-    if (!heap.empty()) {
-      PopulateIterator(heap);
+    if (!allow_unprepared_value_ && !heap.empty()) {
+      [[maybe_unused]] const bool result = PopulateIterator(heap);
+      assert(result || (!Valid() && !status_.ok()));
     }
   }
 
   template <typename BinaryHeap>
-  void PopulateIterator(BinaryHeap& heap) {
+  bool PopulateIterator(BinaryHeap& heap) {
     // 1. Keep the top iterator (by popping it from the heap) and add it to list
     //    to populate
     // 2. For all non-top iterators having the same key as top iter popped
@@ -265,31 +285,69 @@ class MultiCfIteratorImpl {
     //    populate the value/columns and attribute_groups from the list
     //    collected in step 1 and 2 and add all the iters back to the heap
     assert(!heap.empty());
+
+    auto prepare_value = [this, &heap](Iterator* iterator) {
+      assert(iterator);
+      assert(iterator->Valid());
+      assert(iterator->status().ok());
+
+      if (!iterator->PrepareValue()) {
+        assert(!iterator->Valid());
+        assert(!iterator->status().ok());
+
+        considerStatus(iterator->status());
+        heap.clear();
+
+        assert(!Valid());
+        assert(!status_.ok());
+
+        return false;
+      }
+
+      return true;
+    };
+
     auto top = heap.top();
-    heap.pop();
+    assert(top.iterator);
+    assert(top.iterator->Valid());
+    assert(top.iterator->status().ok());
+
+    if (!prepare_value(top.iterator)) {
+      return false;
+    }
+
     autovector<MultiCfIteratorInfo> to_populate;
+
     to_populate.push_back(top);
-    if (!heap.empty()) {
+    heap.pop();
+
+    while (!heap.empty()) {
       auto current = heap.top();
       assert(current.iterator);
-      while (current.iterator->Valid() &&
-             comparator_->Compare(top.iterator->key(),
-                                  current.iterator->key()) == 0) {
-        assert(current.iterator->status().ok());
-        to_populate.push_back(current);
-        heap.pop();
-        if (!heap.empty()) {
-          current = heap.top();
-        } else {
-          break;
-        }
+      assert(current.iterator->Valid());
+      assert(current.iterator->status().ok());
+
+      if (comparator_->Compare(current.iterator->key(), top.iterator->key()) !=
+          0) {
+        break;
+      }
+
+      if (!prepare_value(current.iterator)) {
+        return false;
       }
+
+      to_populate.push_back(current);
+      heap.pop();
     }
+
     // Add the items back to the heap
     for (auto& item : to_populate) {
       heap.push(item);
     }
+
     populate_func_(to_populate);
+
+    return true;
   }
 };