#pragma once
#include <algorithm>
#include <memory>
#include <utility>
#include <vector>
#include "table/internal_iterator.h"
namespace
rocksdb {
class
PinnedIteratorsManager :
public
Cleanable {
public
:
PinnedIteratorsManager() : pinning_enabled(
false
) {}
~PinnedIteratorsManager() {
if
(pinning_enabled) {
ReleasePinnedData();
}
}
void
StartPinning() {
assert
(pinning_enabled ==
false
);
pinning_enabled =
true
;
}
bool
PinningEnabled() {
return
pinning_enabled; }
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_ptrs_.emplace_back(ptr, release_func);
}
inline
void
ReleasePinnedData() {
assert
(pinning_enabled ==
true
);
pinning_enabled =
false
;
std::sort(pinned_ptrs_.begin(), pinned_ptrs_.end());
auto
unique_end = std::unique(pinned_ptrs_.begin(), pinned_ptrs_.end());
for
(
auto
i = pinned_ptrs_.begin(); i != unique_end; ++i) {
void
* ptr = i->first;
ReleaseFunction release_func = i->second;
release_func(ptr);
}
pinned_ptrs_.clear();
Cleanable::Reset();
}
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::vector<std::pair<
void
*, ReleaseFunction>> pinned_ptrs_;
};
}