Minor PinnedIteratorsManager Refactoring

Summary:
This diff include these simple change
- Rename ReleasePinnedIterators to ReleasePinnedData
- Rename PinIteratorIfNeeded to PinIterator
- Use std::vector directly in PinnedIteratorsManager instead of std::unique_ptr<std::vector>
- Generalize PinnedIteratorsManager by adding PinPtr which can pin any pointer

Test Plan: existing tests

Reviewers: sdong, yiwu, andrewkr

Reviewed By: andrewkr

Subscribers: andrewkr, dhruba

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

db/compaction_iterator.cc

@@ -90,7 +90,7 @@ void CompactionIterator::Next() {
       valid_ = true;
     } else {
       // We consumed all pinned merge operands, release pinned iterators
-      pinned_iters_mgr_.ReleasePinnedIterators();
+      pinned_iters_mgr_.ReleasePinnedData();
       // MergeHelper moves the iterator to the first record after the merged
       // records, so even though we reached the end of the merge output, we do
       // not want to advance the iterator.
@@ -404,7 +404,7 @@ void CompactionIterator::NextFromInput() {
         // batch consumed by the merge operator should not shadow any keys
         // coming after the merges
         has_current_user_key_ = false;
-        pinned_iters_mgr_.ReleasePinnedIterators();
+        pinned_iters_mgr_.ReleasePinnedData();
       }
     } else {
       valid_ = true;

db/db_iter.cc

@@ -133,7 +133,9 @@ class DBIter: public Iterator {
   }
   virtual ~DBIter() {
     // Release pinned data if any
-    pinned_iters_mgr_.ReleasePinnedIterators();
+    if (pinned_iters_mgr_.PinningEnabled()) {
+      pinned_iters_mgr_.ReleasePinnedData();
+    }
     RecordTick(statistics_, NO_ITERATORS, -1);
     local_stats_.BumpGlobalStatistics(statistics_);
     if (!arena_mode_) {
@@ -223,8 +225,8 @@ class DBIter: public Iterator {
 
   // Release blocks pinned by TempPinData()
   void ReleaseTempPinnedData() {
-    if (!pin_thru_lifetime_) {
-      pinned_iters_mgr_.ReleasePinnedIterators();
+    if (!pin_thru_lifetime_ && pinned_iters_mgr_.PinningEnabled()) {
+      pinned_iters_mgr_.ReleasePinnedData();
     }
   }
 

db/pinned_iterators_manager.h

@@ -6,6 +6,7 @@
 #pragma once
 #include <algorithm>
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "table/internal_iterator.h"
@@ -17,50 +18,70 @@ namespace rocksdb {
 // not needed anymore.
 class PinnedIteratorsManager {
  public:
-  PinnedIteratorsManager() : pinning_enabled(false), pinned_iters_(nullptr) {}
-  ~PinnedIteratorsManager() { ReleasePinnedIterators(); }
+  PinnedIteratorsManager() : pinning_enabled(false) {}
+  ~PinnedIteratorsManager() {
+    if (pinning_enabled) {
+      ReleasePinnedData();
+    }
+  }
 
   // Enable Iterators pinning
   void StartPinning() {
-    if (!pinning_enabled) {
+    assert(pinning_enabled == false);
     pinning_enabled = true;
-      if (!pinned_iters_) {
-        pinned_iters_.reset(new std::vector<InternalIterator*>());
-      }
-    }
   }
 
   // Is pinning enabled ?
   bool PinningEnabled() { return pinning_enabled; }
 
-  // Take ownership of iter if pinning is enabled and delete it when
-  // ReleasePinnedIterators() is called
-  void PinIteratorIfNeeded(InternalIterator* iter) {
-    if (!pinning_enabled || !iter) {
+  // Take ownership of iter and delete it when ReleasePinnedData() is called
+  void PinIterator(InternalIterator* iter, bool arena = false) {
+    if (arena) {
+      PinPtr(iter, &PinnedIteratorsManager::ReleaseArenaInternalIterator);
+    } else {
+      PinPtr(iter, &PinnedIteratorsManager::ReleaseInternalIterator);
+    }
+  }
+
+  typedef void (*ReleaseFunction)(void* arg1);
+  void PinPtr(void* ptr, ReleaseFunction release_func) {
+    assert(pinning_enabled);
+    if (ptr == nullptr) {
       return;
     }
-    pinned_iters_->push_back(iter);
+    pinned_ptrs_.emplace_back(ptr, release_func);
   }
 
   // Release pinned Iterators
-  inline void ReleasePinnedIterators() {
-    if (pinning_enabled) {
+  inline void ReleasePinnedData() {
+    assert(pinning_enabled == true);
     pinning_enabled = false;
 
     // Remove duplicate pointers
-      std::sort(pinned_iters_->begin(), pinned_iters_->end());
-      std::unique(pinned_iters_->begin(), pinned_iters_->end());
+    std::sort(pinned_ptrs_.begin(), pinned_ptrs_.end());
+    std::unique(pinned_ptrs_.begin(), pinned_ptrs_.end());
 
-      for (auto& iter : *pinned_iters_) {
-        delete iter;
-      }
-      pinned_iters_->clear();
+    for (size_t i = 0; i < pinned_ptrs_.size(); i++) {
+      assert(i == 0 || pinned_ptrs_[i].first != pinned_ptrs_[i - 1].first);
+
+      void* ptr = pinned_ptrs_[i].first;
+      ReleaseFunction release_func = pinned_ptrs_[i].second;
+      release_func(ptr);
     }
+    pinned_ptrs_.clear();
   }
 
  private:
+  static void ReleaseInternalIterator(void* ptr) {
+    delete reinterpret_cast<InternalIterator*>(ptr);
+  }
+
+  static void ReleaseArenaInternalIterator(void* ptr) {
+    reinterpret_cast<InternalIterator*>(ptr)->~InternalIterator();
+  }
+
   bool pinning_enabled;
-  std::unique_ptr<std::vector<InternalIterator*>> pinned_iters_;
+  std::vector<std::pair<void*, ReleaseFunction>> pinned_ptrs_;
 };
 
 }  // namespace rocksdb

table/block_based_table_reader.cc

@@ -1446,7 +1446,7 @@ Status BlockBasedTable::Get(const ReadOptions& read_options, const Slice& key,
         if (pin_blocks) {
           if (get_context->State() == GetContext::kMerge) {
             // Pin blocks as long as we are merging
-            pinned_iters_mgr->PinIteratorIfNeeded(biter);
+            pinned_iters_mgr->PinIterator(biter);
           } else {
             delete biter;
           }

table/internal_iterator.h

@@ -71,7 +71,7 @@ class InternalIterator : public Cleanable {
   virtual void SetPinnedItersMgr(PinnedIteratorsManager* pinned_iters_mgr) {}
 
   // If true, this means that the Slice returned by key() is valid as long as
-  // PinnedIteratorsManager::ReleasePinnedIterators is not called and the
+  // PinnedIteratorsManager::ReleasePinnedData is not called and the
   // Iterator is not deleted.
   //
   // IsKeyPinned() is guaranteed to always return true if
@@ -81,7 +81,7 @@ class InternalIterator : public Cleanable {
   virtual bool IsKeyPinned() const { return false; }
 
   // If true, this means that the Slice returned by value() is valid as long as
-  // PinnedIteratorsManager::ReleasePinnedIterators is not called and the
+  // PinnedIteratorsManager::ReleasePinnedData is not called and the
   // Iterator is not deleted.
   virtual bool IsValuePinned() const { return false; }
 

table/two_level_iterator.cc

@@ -202,7 +202,7 @@ void TwoLevelIterator::SetSecondLevelIterator(InternalIterator* iter) {
 
   InternalIterator* old_iter = second_level_iter_.Set(iter);
   if (pinned_iters_mgr_ && pinned_iters_mgr_->PinningEnabled()) {
-    pinned_iters_mgr_->PinIteratorIfNeeded(old_iter);
+    pinned_iters_mgr_->PinIterator(old_iter);
   } else {
     delete old_iter;
   }