EvictingCacheMap.h

Go to the documentation of this file.

/*
 * Copyright 2014-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <algorithm>
#include <exception>
#include <functional>

#include <boost/intrusive/list.hpp>
#include <boost/intrusive/unordered_set.hpp>
#include <boost/iterator/iterator_adaptor.hpp>
#include <boost/utility.hpp>

#include <folly/lang/Exception.h>

namespace folly {

template <
    class TKey,
    class TValue,
    class THash = std::hash<TKey>,
    class TKeyEqual = std::equal_to<TKey>>
class EvictingCacheMap {
 private:
  // typedefs for brevity
  struct Node;
  struct KeyHasher;
  struct KeyValueEqual;
  typedef boost::intrusive::link_mode<boost::intrusive::safe_link> link_mode;
  typedef boost::intrusive::unordered_set<
      Node,
      boost::intrusive::hash<KeyHasher>,
      boost::intrusive::equal<KeyValueEqual>>
      NodeMap;
  typedef boost::intrusive::list<Node> NodeList;
  typedef std::pair<const TKey, TValue> TPair;

 public:
  typedef std::function<void(TKey, TValue&&)> PruneHookCall;

  // iterator base : returns TPair on dereference
  template <typename Value, typename TIterator>
  class iterator_base : public boost::iterator_adaptor<
                            iterator_base<Value, TIterator>,
                            TIterator,
                            Value,
                            boost::bidirectional_traversal_tag> {
   public:
    iterator_base() {}
    explicit iterator_base(TIterator it)
        : iterator_base::iterator_adaptor_(it) {}
    Value& dereference() const {
      return this->base_reference()->pr;
    }
  };

  // iterators
  typedef iterator_base<TPair, typename NodeList::iterator> iterator;
  typedef iterator_base<const TPair, typename NodeList::const_iterator>
      const_iterator;
  typedef iterator_base<TPair, typename NodeList::reverse_iterator>
      reverse_iterator;
  typedef iterator_base<const TPair, typename NodeList::const_reverse_iterator>
      const_reverse_iterator;

  // the default map typedefs
  using key_type = TKey;
  using mapped_type = TValue;
  using hasher = THash;

  explicit EvictingCacheMap(
      std::size_t maxSize,
      std::size_t clearSize = 1,
      const THash& keyHash = THash(),
      const TKeyEqual& keyEqual = TKeyEqual())
      : nIndexBuckets_(std::max(maxSize / 2, std::size_t(kMinNumIndexBuckets))),
        indexBuckets_(new typename NodeMap::bucket_type[nIndexBuckets_]),
        indexTraits_(indexBuckets_.get(), nIndexBuckets_),
        keyHash_(keyHash),
        keyEqual_(keyEqual),
        index_(indexTraits_, keyHash_, keyEqual_),
        maxSize_(maxSize),
        clearSize_(clearSize) {}

  EvictingCacheMap(const EvictingCacheMap&) = delete;
  EvictingCacheMap& operator=(const EvictingCacheMap&) = delete;
  EvictingCacheMap(EvictingCacheMap&&) = default;
  EvictingCacheMap& operator=(EvictingCacheMap&&) = default;

  ~EvictingCacheMap() {
    setPruneHook(nullptr);
    // ignore any potential exceptions from pruneHook_
    pruneWithFailSafeOption(size(), nullptr, true);
  }

  void setMaxSize(size_t maxSize, PruneHookCall pruneHook = nullptr) {
    if (maxSize != 0 && maxSize < size()) {
      // Prune the excess elements with our new constraints.
      prune(std::max(size() - maxSize, clearSize_), pruneHook);
    }
    maxSize_ = maxSize;
  }

  size_t getMaxSize() const {
    return maxSize_;
  }

  void setClearSize(size_t clearSize) {
    clearSize_ = clearSize;
  }

  bool exists(const TKey& key) const {
    return findInIndex(key) != index_.end();
  }

  TValue& get(const TKey& key) {
    auto it = find(key);
    if (it == end()) {
      throw_exception<std::out_of_range>("Key does not exist");
    }
    return it->second;
  }

  iterator find(const TKey& key) {
    auto it = findInIndex(key);
    if (it == index_.end()) {
      return end();
    }
    lru_.erase(lru_.iterator_to(*it));
    lru_.push_front(*it);
    return iterator(lru_.iterator_to(*it));
  }

  const TValue& getWithoutPromotion(const TKey& key) const {
    auto it = findWithoutPromotion(key);
    if (it == end()) {
      throw_exception<std::out_of_range>("Key does not exist");
    }
    return it->second;
  }

  TValue& getWithoutPromotion(const TKey& key) {
    auto const& cThis = *this;
    return const_cast<TValue&>(cThis.getWithoutPromotion(key));
  }

  const_iterator findWithoutPromotion(const TKey& key) const {
    auto it = findInIndex(key);
    return (it == index_.end()) ? end() : const_iterator(lru_.iterator_to(*it));
  }

  iterator findWithoutPromotion(const TKey& key) {
    auto it = findInIndex(key);
    return (it == index_.end()) ? end() : iterator(lru_.iterator_to(*it));
  }

  bool erase(const TKey& key) {
    auto it = findInIndex(key);
    if (it == index_.end()) {
      return false;
    }
    auto node = &(*it);
    std::unique_ptr<Node> nptr(node);
    lru_.erase(lru_.iterator_to(*node));
    index_.erase(it);
    return true;
  }

  void set(
      const TKey& key,
      TValue value,
      bool promote = true,
      PruneHookCall pruneHook = nullptr) {
    auto it = findInIndex(key);
    if (it != index_.end()) {
      it->pr.second = std::move(value);
      if (promote) {
        lru_.erase(lru_.iterator_to(*it));
        lru_.push_front(*it);
      }
    } else {
      auto node = new Node(key, std::move(value));
      index_.insert(*node);
      lru_.push_front(*node);

      // no evictions if maxSize_ is 0 i.e. unlimited capacity
      if (maxSize_ > 0 && size() > maxSize_) {
        prune(clearSize_, pruneHook);
      }
    }
  }

  std::size_t size() const {
    return index_.size();
  }

  bool empty() const {
    return index_.empty();
  }

  void clear(PruneHookCall pruneHook = nullptr) {
    prune(size(), pruneHook);
  }

  void setPruneHook(PruneHookCall pruneHook) {
    pruneHook_ = pruneHook;
  }

  void prune(std::size_t pruneSize, PruneHookCall pruneHook = nullptr) {
    // do not swallow exceptions for prunes not triggered from destructor
    pruneWithFailSafeOption(pruneSize, pruneHook, false);
  }

  // Iterators and such
  iterator begin() {
    return iterator(lru_.begin());
  }
  iterator end() {
    return iterator(lru_.end());
  }
  const_iterator begin() const {
    return const_iterator(lru_.begin());
  }
  const_iterator end() const {
    return const_iterator(lru_.end());
  }

  const_iterator cbegin() const {
    return const_iterator(lru_.cbegin());
  }
  const_iterator cend() const {
    return const_iterator(lru_.cend());
  }

  reverse_iterator rbegin() {
    return reverse_iterator(lru_.rbegin());
  }
  reverse_iterator rend() {
    return reverse_iterator(lru_.rend());
  }

  const_reverse_iterator rbegin() const {
    return const_reverse_iterator(lru_.rbegin());
  }
  const_reverse_iterator rend() const {
    return const_reverse_iterator(lru_.rend());
  }

  const_reverse_iterator crbegin() const {
    return const_reverse_iterator(lru_.crbegin());
  }
  const_reverse_iterator crend() const {
    return const_reverse_iterator(lru_.crend());
  }

 private:
  struct Node : public boost::intrusive::unordered_set_base_hook<link_mode>,
                public boost::intrusive::list_base_hook<link_mode> {
    Node(const TKey& key, TValue&& value)
        : pr(std::make_pair(key, std::move(value))) {}
    TPair pr;
  };

  struct KeyHasher {
    KeyHasher(const THash& keyHash) : hash(keyHash) {}
    std::size_t operator()(const Node& node) const {
      return hash(node.pr.first);
    }
    std::size_t operator()(const TKey& key) const {
      return hash(key);
    }
    THash hash;
  };

  struct KeyValueEqual {
    KeyValueEqual(const TKeyEqual& keyEqual) : equal(keyEqual) {}
    bool operator()(const TKey& lhs, const Node& rhs) const {
      return equal(lhs, rhs.pr.first);
    }
    bool operator()(const Node& lhs, const TKey& rhs) const {
      return equal(lhs.pr.first, rhs);
    }
    bool operator()(const Node& lhs, const Node& rhs) const {
      return equal(lhs.pr.first, rhs.pr.first);
    }
    TKeyEqual equal;
  };

  typename NodeMap::iterator findInIndex(const TKey& key) {
    return index_.find(key, KeyHasher(keyHash_), KeyValueEqual(keyEqual_));
  }

  typename NodeMap::const_iterator findInIndex(const TKey& key) const {
    return index_.find(key, KeyHasher(keyHash_), KeyValueEqual(keyEqual_));
  }

  void pruneWithFailSafeOption(
      std::size_t pruneSize,
      PruneHookCall pruneHook,
      bool failSafe) {
    auto& ph = (nullptr == pruneHook) ? pruneHook_ : pruneHook;

    for (std::size_t i = 0; i < pruneSize && !lru_.empty(); i++) {
      auto* node = &(*lru_.rbegin());
      std::unique_ptr<Node> nptr(node);

      lru_.erase(lru_.iterator_to(*node));
      index_.erase(index_.iterator_to(*node));
      if (ph) {
        try {
          ph(node->pr.first, std::move(node->pr.second));
        } catch (...) {
          if (!failSafe) {
            throw;
          }
        }
      }
    }
  }

  static const std::size_t kMinNumIndexBuckets = 100;
  PruneHookCall pruneHook_;
  std::size_t nIndexBuckets_;
  std::unique_ptr<typename NodeMap::bucket_type[]> indexBuckets_;
  typename NodeMap::bucket_traits indexTraits_;
  THash keyHash_;
  TKeyEqual keyEqual_;
  NodeMap index_;
  NodeList lru_;
  std::size_t maxSize_;
  std::size_t clearSize_;
};

} // namespace folly