diff --git a/db/compaction_iterator.cc b/db/compaction_iterator.cc index aa160bee7..81b93a2f0 100644 --- a/db/compaction_iterator.cc +++ b/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; diff --git a/db/db_iter.cc b/db/db_iter.cc index 5de3b406b..2a12534f2 100644 --- a/db/db_iter.cc +++ b/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(); } } diff --git a/db/pinned_iterators_manager.h b/db/pinned_iterators_manager.h index 38cb89ce0..f8d64762c 100644 --- a/db/pinned_iterators_manager.h +++ b/db/pinned_iterators_manager.h @@ -6,6 +6,7 @@ #pragma once #include #include +#include #include #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) { - pinning_enabled = true; - if (!pinned_iters_) { - pinned_iters_.reset(new std::vector()); - } - } + assert(pinning_enabled == false); + pinning_enabled = true; } // 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) { - pinning_enabled = false; + inline void ReleasePinnedData() { + assert(pinning_enabled == true); + pinning_enabled = false; + + // Remove duplicate pointers + std::sort(pinned_ptrs_.begin(), pinned_ptrs_.end()); + std::unique(pinned_ptrs_.begin(), pinned_ptrs_.end()); - // Remove duplicate pointers - std::sort(pinned_iters_->begin(), pinned_iters_->end()); - std::unique(pinned_iters_->begin(), pinned_iters_->end()); + for (size_t i = 0; i < pinned_ptrs_.size(); i++) { + assert(i == 0 || pinned_ptrs_[i].first != pinned_ptrs_[i - 1].first); - for (auto& iter : *pinned_iters_) { - delete iter; - } - pinned_iters_->clear(); + 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(ptr); + } + + static void ReleaseArenaInternalIterator(void* ptr) { + reinterpret_cast(ptr)->~InternalIterator(); + } + bool pinning_enabled; - std::unique_ptr> pinned_iters_; + std::vector> pinned_ptrs_; }; } // namespace rocksdb diff --git a/table/block_based_table_reader.cc b/table/block_based_table_reader.cc index 5a1e8ebd3..635a1c877 100644 --- a/table/block_based_table_reader.cc +++ b/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; } diff --git a/table/internal_iterator.h b/table/internal_iterator.h index f1f1e0bff..2850a6773 100644 --- a/table/internal_iterator.h +++ b/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; } diff --git a/table/two_level_iterator.cc b/table/two_level_iterator.cc index 81dc8792b..54523fad7 100644 --- a/table/two_level_iterator.cc +++ b/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; }