1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 //
6 #pragma once
7 #include <algorithm>
8 #include <memory>
9 #include <utility>
10 #include <vector>
11 
12 #include "table/internal_iterator.h"
13 
14 namespace ROCKSDB_NAMESPACE {
15 
16 // PinnedIteratorsManager will be notified whenever we need to pin an Iterator
17 // and it will be responsible for deleting pinned Iterators when they are
18 // not needed anymore.
19 class PinnedIteratorsManager : public Cleanable {
20  public:
PinnedIteratorsManager()21   PinnedIteratorsManager() : pinning_enabled(false) {}
~PinnedIteratorsManager()22   ~PinnedIteratorsManager() {
23     if (pinning_enabled) {
24       ReleasePinnedData();
25     }
26   }
27 
28   // Enable Iterators pinning
StartPinning()29   void StartPinning() {
30     assert(pinning_enabled == false);
31     pinning_enabled = true;
32   }
33 
34   // Is pinning enabled ?
PinningEnabled()35   bool PinningEnabled() { return pinning_enabled; }
36 
37   // Take ownership of iter and delete it when ReleasePinnedData() is called
38   void PinIterator(InternalIterator* iter, bool arena = false) {
39     if (arena) {
40       PinPtr(iter, &PinnedIteratorsManager::ReleaseArenaInternalIterator);
41     } else {
42       PinPtr(iter, &PinnedIteratorsManager::ReleaseInternalIterator);
43     }
44   }
45 
46   typedef void (*ReleaseFunction)(void* arg1);
PinPtr(void * ptr,ReleaseFunction release_func)47   void PinPtr(void* ptr, ReleaseFunction release_func) {
48     assert(pinning_enabled);
49     if (ptr == nullptr) {
50       return;
51     }
52     pinned_ptrs_.emplace_back(ptr, release_func);
53   }
54 
55   // Release pinned Iterators
ReleasePinnedData()56   inline void ReleasePinnedData() {
57     assert(pinning_enabled == true);
58     pinning_enabled = false;
59 
60     // Remove duplicate pointers
61     std::sort(pinned_ptrs_.begin(), pinned_ptrs_.end());
62     auto unique_end = std::unique(pinned_ptrs_.begin(), pinned_ptrs_.end());
63 
64     for (auto i = pinned_ptrs_.begin(); i != unique_end; ++i) {
65       void* ptr = i->first;
66       ReleaseFunction release_func = i->second;
67       release_func(ptr);
68     }
69     pinned_ptrs_.clear();
70     // Also do cleanups from the base Cleanable
71     Cleanable::Reset();
72   }
73 
74  private:
ReleaseInternalIterator(void * ptr)75   static void ReleaseInternalIterator(void* ptr) {
76     delete reinterpret_cast<InternalIterator*>(ptr);
77   }
78 
ReleaseArenaInternalIterator(void * ptr)79   static void ReleaseArenaInternalIterator(void* ptr) {
80     reinterpret_cast<InternalIterator*>(ptr)->~InternalIterator();
81   }
82 
83   bool pinning_enabled;
84   std::vector<std::pair<void*, ReleaseFunction>> pinned_ptrs_;
85 };
86 
87 }  // namespace ROCKSDB_NAMESPACE
88