rocksdb/db/pinned_iterators_manager.h

//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under the BSD-style license found in the
//  LICENSE file in the root directory of this source tree. An additional grant
//  of patent rights can be found in the PATENTS file in the same directory.
//
#pragma once
#include <algorithm>
#include <memory>
#include <utility>
#include <vector>

#include "table/internal_iterator.h"

namespace rocksdb {

// PinnedIteratorsManager will be notified whenever we need to pin an Iterator
// and it will be responsible for deleting pinned Iterators when they are
// not needed anymore.
class PinnedIteratorsManager {
 public:
  PinnedIteratorsManager() : pinning_enabled(false) {}
  ~PinnedIteratorsManager() {
    if (pinning_enabled) {
      ReleasePinnedData();
    }
  }

  // Enable Iterators pinning
  void StartPinning() {
    assert(pinning_enabled == false);
    pinning_enabled = true;
  }

  // Is pinning enabled ?
  bool PinningEnabled() { return pinning_enabled; }

  // 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_ptrs_.emplace_back(ptr, release_func);
  }

  // Release pinned Iterators
  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());

    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::vector<std::pair<void*, ReleaseFunction>> pinned_ptrs_;
};

}  // namespace rocksdb