Allocate LevelFileIteratorState and LevelFileNumIterator from DB iterator's arena

Summary: Try to allocate LevelFileIteratorState and LevelFileNumIterator from DB iterator's arena, instead of calling malloc and free.

Test Plan: valgrind check

Reviewers: rven, yhchiang, anthony, kradhakrishnan, igor

Reviewed By: igor

Subscribers: leveldb, dhruba

Differential Revision: https://reviews.facebook.net/D40929

db/version_set.cc

@@ -695,12 +695,14 @@ void Version::AddIterators(const ReadOptions& read_options,
     return;
   }
 
+  auto* arena = merge_iter_builder->GetArena();
+
   // Merge all level zero files together since they may overlap
   for (size_t i = 0; i < storage_info_.LevelFilesBrief(0).num_files; i++) {
     const auto& file = storage_info_.LevelFilesBrief(0).files[i];
     merge_iter_builder->AddIterator(cfd_->table_cache()->NewIterator(
         read_options, soptions, cfd_->internal_comparator(), file.fd, nullptr,
-        false, merge_iter_builder->GetArena()));
+        false, arena));
   }
 
   // For levels > 0, we can use a concatenating iterator that sequentially
@@ -708,14 +710,16 @@ void Version::AddIterators(const ReadOptions& read_options,
   // lazily.
   for (int level = 1; level < storage_info_.num_non_empty_levels(); level++) {
     if (storage_info_.LevelFilesBrief(level).num_files != 0) {
-      merge_iter_builder->AddIterator(NewTwoLevelIterator(
-          new LevelFileIteratorState(
-              cfd_->table_cache(), read_options, soptions,
-              cfd_->internal_comparator(), false /* for_compaction */,
-              cfd_->ioptions()->prefix_extractor != nullptr),
-          new LevelFileNumIterator(cfd_->internal_comparator(),
-                                   &storage_info_.LevelFilesBrief(level)),
-          merge_iter_builder->GetArena()));
+      auto* mem = arena->AllocateAligned(sizeof(LevelFileIteratorState));
+      auto* state = new (mem) LevelFileIteratorState(
+          cfd_->table_cache(), read_options, soptions,
+          cfd_->internal_comparator(), false /* for_compaction */,
+          cfd_->ioptions()->prefix_extractor != nullptr);
+      mem = arena->AllocateAligned(sizeof(LevelFileNumIterator));
+      auto* first_level_iter = new (mem) LevelFileNumIterator(
+          cfd_->internal_comparator(), &storage_info_.LevelFilesBrief(level));
+      merge_iter_builder->AddIterator(
+          NewTwoLevelIterator(state, first_level_iter, arena, false));
     }
   }
 }

table/two_level_iterator.cc

@@ -22,11 +22,17 @@ namespace {
 class TwoLevelIterator: public Iterator {
  public:
   explicit TwoLevelIterator(TwoLevelIteratorState* state,
-      Iterator* first_level_iter);
+                            Iterator* first_level_iter,
+                            bool need_free_iter_and_state);
 
   virtual ~TwoLevelIterator() {
-    first_level_iter_.DeleteIter(false);
+    first_level_iter_.DeleteIter(!need_free_iter_and_state_);
     second_level_iter_.DeleteIter(false);
+    if (need_free_iter_and_state_) {
+      delete state_;
+    } else {
+      state_->~TwoLevelIteratorState();
+    }
   }
 
   virtual void Seek(const Slice& target) override;
@@ -65,9 +71,10 @@ class TwoLevelIterator: public Iterator {
   void SetSecondLevelIterator(Iterator* iter);
   void InitDataBlock();
 
-  std::unique_ptr<TwoLevelIteratorState> state_;
+  TwoLevelIteratorState* state_;
   IteratorWrapper first_level_iter_;
   IteratorWrapper second_level_iter_;  // May be nullptr
+  bool need_free_iter_and_state_;
   Status status_;
   // If second_level_iter is non-nullptr, then "data_block_handle_" holds the
   // "index_value" passed to block_function_ to create the second_level_iter.
@@ -75,8 +82,11 @@ class TwoLevelIterator: public Iterator {
 };
 
 TwoLevelIterator::TwoLevelIterator(TwoLevelIteratorState* state,
-    Iterator* first_level_iter)
-  : state_(state), first_level_iter_(first_level_iter) {}
+                                   Iterator* first_level_iter,
+                                   bool need_free_iter_and_state)
+    : state_(state),
+      first_level_iter_(first_level_iter),
+      need_free_iter_and_state_(need_free_iter_and_state) {}
 
 void TwoLevelIterator::Seek(const Slice& target) {
   if (state_->check_prefix_may_match &&
@@ -186,12 +196,15 @@ void TwoLevelIterator::InitDataBlock() {
 }  // namespace
 
 Iterator* NewTwoLevelIterator(TwoLevelIteratorState* state,
-                              Iterator* first_level_iter, Arena* arena) {
+                              Iterator* first_level_iter, Arena* arena,
+                              bool need_free_iter_and_state) {
   if (arena == nullptr) {
-    return new TwoLevelIterator(state, first_level_iter);
+    return new TwoLevelIterator(state, first_level_iter,
+                                need_free_iter_and_state);
   } else {
     auto mem = arena->AllocateAligned(sizeof(TwoLevelIterator));
-    return new (mem) TwoLevelIterator(state, first_level_iter);
+    return new (mem)
+        TwoLevelIterator(state, first_level_iter, need_free_iter_and_state);
   }
 }
 

table/two_level_iterator.h

@@ -43,8 +43,11 @@ struct TwoLevelIteratorState {
 // arena: If not null, the arena is used to allocate the Iterator.
 //        When destroying the iterator, the destructor will destroy
 //        all the states but those allocated in arena.
+// need_free_iter_and_state: free `state` and `first_level_iter` if
+//                           true. Otherwise, just call destructor.
 extern Iterator* NewTwoLevelIterator(TwoLevelIteratorState* state,
                                      Iterator* first_level_iter,
-                                     Arena* arena = nullptr);
+                                     Arena* arena = nullptr,
+                                     bool need_free_iter_and_state = true);
 
 }  // namespace rocksdb