sapling/eden/utils/LeaseCache.h

/*
 *  Copyright (c) 2016, 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 <folly/EvictingCacheMap.h>
#include <folly/futures/Future.h>
#include <folly/futures/SharedPromise.h>

namespace facebook {
namespace eden {

template <typename KEY, typename VAL, typename HASH=std::hash<KEY>>
class LeaseCache {
 public:
  using ValuePtr = std::shared_ptr<VAL>;
  using FutureType = folly::Future<ValuePtr>;
  using SharedPromiseType = std::shared_ptr<folly::SharedPromise<ValuePtr>>;
  using FetchFunc = std::function<FutureType(const KEY &key)>;

 private:
  std::mutex lock_;
  folly::EvictingCacheMap<KEY, SharedPromiseType, HASH> cache_;
  FetchFunc fetcher_;

 public:
  LeaseCache(size_t maxSize, FetchFunc fetcher, size_t clearSize = 1) :
    cache_(maxSize, clearSize),
    fetcher_(fetcher) {
  }

  void set(const KEY &key, ValuePtr val) {
    std::lock_guard<std::mutex> g(lock_);
    auto entry = std::make_shared<typename SharedPromiseType::element_type>();
    entry->setValue(val);
    cache_.set(key, entry);
  }

  void erase(const KEY&key) {
    std::lock_guard<std::mutex> g(lock_);
    cache_.erase(key);
  }

  void setMaxSize(size_t size) {
    cache_.setMaxSize(size);
  }

  FutureType get(const KEY &key) {
    SharedPromiseType entry;

    {
      std::lock_guard<std::mutex> g(lock_);

      auto it = cache_.find(key);
      if (it != cache_.end()) {
        entry = it->second;
        return entry->getFuture();
      }

      entry = std::make_shared<typename SharedPromiseType::element_type>();
      cache_.set(key, entry);
    }

    auto future = entry->getFuture();

    fetcher_(key).then([this, entry](folly::Try<ValuePtr>&& t) {
      entry->setTry(std::move(t));
    });

    return future;
  }

  bool exists(const KEY& key) { return cache_.exists(key); }
};

}
}