Implement FlashCacheComponent's allocate & insert APIs

Summary:
^

BlockCache isn't currently exposed to the outside world, so it always expects to have the value ready to write at allocation time.  Break up inserting into BlockCache according to the interface with an allocate (make space for writing) separate from inserting (make visible in cache).

This isn't too much work since Regions already support refcounting, so we can create an AllocatedHandle and then release the region when the item is inserted (or discarded).

Note: clean up writeEntry to return void since it always returns an ok status.

Reviewed By: AlnisM

Differential Revision: D88516565

fbshipit-source-id: cf83e4e36b6192e22c33924b381083f673945fd6

Author: Rob Lyerly

cachelib/interface/CacheComponent.h

@@ -73,6 +73,9 @@ class CacheComponent {
    * Returns an AllocatedHandle to the old item if it was replaced. The input
    * AllocatedHandle is no longer usable (moved out).
    *
+   * Note: it may be too expensive for some implementations to return the old
+   * item; they'll just return std::nullopt.
+   *
    * @param handle AllocatedHandle returned by allocate()
    * @return an empty optional if the item was inserted, an AllocatedHandle to
    * the replaced item if it was replaced (no longer findable replaced) or an
@@ -110,6 +113,15 @@ class CacheComponent {
   virtual folly::coro::Task<UnitResult> remove(ReadHandle&& handle) = 0;
 
  protected:
+  /**
+   * Mark an item as inserted into cache.
+   * @param handle handle to the cache item
+   * @param inserted whether the item was inserted into cache
+   */
+  FOLLY_ALWAYS_INLINE void setInserted(Handle& handle, bool inserted) {
+    handle.inserted_ = inserted;
+  }
+
   /**
    * Release a handle (make unusable) without adjusting refcounts. Useful when
    * CacheComponent takes an rvalue ref to a handle that needs to be destroyed.

cachelib/interface/Result.h

@@ -26,6 +26,7 @@ struct Error {
     INVALID_CONFIG,
 
     // Allocation errors
+    ALLOCATE_FAILED,
     NO_SPACE,
 
     // Insertion errors

cachelib/interface/components/FlashCacheComponent.cpp

@@ -20,6 +20,89 @@ using namespace facebook::cachelib::navy;
 
 namespace facebook::cachelib::interface {
 
+/**
+ * In-memory buffer for data cached in flash. Reference counting via
+ * RegionDescriptors is used to ensure the backing Region is still resident &
+ * valid while working on this cache item. The buffer's layout is:
+ *
+ *   ---------------------------------------------------------------------------
+ *   | creation & expiry time ||  value  || (empty) |  key  | entry descriptor |
+ *   ---------------------------------------------------------------------------
+ */
+class FlashCacheItem : public CacheItem {
+ public:
+  using EntryDesc = BlockCache::EntryDesc;
+
+  static constexpr size_t kMetadataSize = 2 * sizeof(uint32_t);
+
+  // Allocation constructor
+  // Note: valueSize *must* include space for creation & expiration time
+  FlashCacheItem(const HashedKey& hashedKey,
+                 uint32_t valueSize,
+                 std::tuple<RegionDescriptor, size_t, RelAddress>&& allocation,
+                 const uint32_t creationTime,
+                 const uint32_t ttlSecs)
+      : desc_(std::move(std::get<0>(allocation))),
+        relAddr_(std::move(std::get<2>(allocation))),
+        buffer_(std::get<1>(allocation)) {
+    XDCHECK_GE(valueSize, 2 * sizeof(uint32_t))
+        << "valueSize doesn't include space for metadata";
+    BlockCache::writeEntryDescAndKey(buffer_, hashedKey, valueSize);
+    uint32_t* metadata = reinterpret_cast<uint32_t*>(buffer_.data());
+    metadata[0] = creationTime;
+    metadata[1] = ttlSecs > 0 ? creationTime + ttlSecs : 0;
+  }
+
+  RegionDescriptor& getRegionDescriptor() noexcept { return desc_; }
+  RelAddress& getRelAddress() noexcept { return relAddr_; }
+  Buffer& getBuffer() noexcept { return buffer_; }
+  EntryDesc* getEntryDescriptor() const noexcept {
+    auto entryEnd = const_cast<uint8_t*>(buffer_.data()) + buffer_.size();
+    return reinterpret_cast<EntryDesc*>(entryEnd - sizeof(EntryDesc));
+  }
+
+  // ------------------------------ Interface ------------------------------ //
+
+  uint32_t getCreationTime() const noexcept override {
+    return reinterpret_cast<const uint32_t*>(buffer_.data())[0];
+  }
+  uint32_t getExpiryTime() const noexcept override {
+    return reinterpret_cast<const uint32_t*>(buffer_.data())[1];
+  }
+  // Ref-counting for FlashCache is different than RAMCache - multiple
+  // RAMCacheItems all refer to a single cache item in memory whereas each
+  // FlashCacheItem is separate (even if it references the same cache item in
+  // Region memory).
+  //
+  // The only thing we care about is Region refcounting via RegionDescriptor. We
+  // don't need to increment on create (creating the RegionDescriptor does that)
+  // and we should always release the cache item to release the descriptor.
+  void incrementRefCount() noexcept override {}
+  bool decrementRefCount() noexcept override {
+    // get cache to release the descriptor (need access to RegionManager)
+    return true;
+  }
+  Key getKey() const noexcept override {
+    auto* desc = getEntryDescriptor();
+    return Key{reinterpret_cast<const char*>(desc) - desc->keySize,
+               desc->keySize};
+  }
+  void* getMemory() const noexcept override {
+    return const_cast<uint8_t*>(buffer_.data() + kMetadataSize);
+  }
+  uint32_t getMemorySize() const noexcept override {
+    return getEntryDescriptor()->valueSize - kMetadataSize;
+  }
+  uint32_t getTotalSize() const noexcept override { return buffer_.size(); }
+
+ private:
+  // Used for region lifecycle management & indexing
+  RegionDescriptor desc_;
+  RelAddress relAddr_;
+  // Data is buffered in RAM
+  Buffer buffer_;
+};
+
 /* static */ Result<FlashCacheComponent> FlashCacheComponent::create(
     std::string name, navy::BlockCache::Config&& config) noexcept {
   try {
@@ -34,21 +117,83 @@ const std::string& FlashCacheComponent::getName() const noexcept {
 }
 
 folly::coro::Task<Result<AllocatedHandle>> FlashCacheComponent::allocate(
-    Key /* key */,
-    uint32_t /* size */,
-    uint32_t /* creationTime */,
-    uint32_t /* ttlSecs */) {
-  co_return makeError(Error::Code::UNIMPLEMENTED, "not yet implemented");
+    Key key, uint32_t size, uint32_t creationTime, uint32_t ttlSecs) {
+  if (key.empty()) {
+    co_return makeError(Error::Code::INVALID_ARGUMENTS, "empty key");
+  }
+
+  HashedKey hashedKey(key);
+  uint32_t valueSize = size + FlashCacheItem::kMetadataSize;
+  auto ret = co_await onWorkerThread(
+      [=, this]() { return cache_->allocateForInsert(hashedKey, valueSize); },
+      [this](auto&& result) {
+        if (result.hasValue()) {
+          // Successfully allocated Region memory but coroutine was cancelled -
+          // close the descriptor to release refcount on the Region
+          auto& [desc, slotSize, _] = result.value();
+          cache_->addHole(slotSize);
+          cache_->regionManager_.close(std::move(desc));
+        }
+      });
+  if (ret.hasError()) {
+    co_return makeError(Error::Code::ALLOCATE_FAILED,
+                        "could not allocate space in a region");
+  }
+  // Note: we can't checksum until the user has finished writing their data,
+  // defer until they call insert()
+  auto* item = new FlashCacheItem(hashedKey, valueSize, std::move(ret).value(),
+                                  creationTime, ttlSecs);
+  co_return AllocatedHandle(*this, *item);
+}
+
+folly::coro::Task<UnitResult> FlashCacheComponent::insertImpl(
+    AllocatedHandle&& handle, bool allowReplace) {
+  if (!handle) {
+    co_return makeError(Error::Code::INVALID_ARGUMENTS,
+                        "empty AllocatedHandle");
+  }
+
+  auto& fccItem = reinterpret_cast<FlashCacheItem&>(*handle);
+  auto* desc = fccItem.getEntryDescriptor();
+  auto logicalSizeWritten = desc->keySize + desc->valueSize;
+  auto totalSize = fccItem.getTotalSize();
+  auto keyHash = desc->keyHash;
+
+  // Note: we don't need to run this on a worker thread because everything here
+  // is synchronous - only allocation is async
+
+  if (cache_->checksumData_) {
+    desc->cs = checksum(fccItem.getBuffer().view().slice(0, desc->valueSize));
+  }
+  cache_->regionManager_.write(fccItem.getRelAddress(),
+                               std::move(fccItem.getBuffer()));
+  cache_->logicalWrittenCount_.add(logicalSizeWritten);
+  auto inserted = cache_->updateIndex(keyHash, totalSize,
+                                      fccItem.getRelAddress(), allowReplace);
+  if (inserted) {
+    setInserted(handle, true);
+    auto _ = std::move(handle);
+    co_return folly::unit;
+  } else {
+    // Note: do hole accounting in release()
+    co_return makeError(Error::Code::ALREADY_INSERTED, "key already inserted");
+  }
 }
 
 folly::coro::Task<UnitResult> FlashCacheComponent::insert(
-    AllocatedHandle&& /* handle */) {
-  co_return makeError(Error::Code::UNIMPLEMENTED, "not yet implemented");
+    AllocatedHandle&& handle) {
+  co_return co_await insertImpl(std::move(handle), /* allowReplace */ false);
 }
 
 folly::coro::Task<Result<std::optional<AllocatedHandle>>>
-FlashCacheComponent::insertOrReplace(AllocatedHandle&& /* handle */) {
-  co_return makeError(Error::Code::UNIMPLEMENTED, "not yet implemented");
+FlashCacheComponent::insertOrReplace(AllocatedHandle&& handle) {
+  auto inserted =
+      co_await insertImpl(std::move(handle), /* allowReplace */ true);
+  XDCHECK(inserted) << "insertOrReplace should never fail";
+  // Note: returning the old data requires a flash read, don't expose for now
+  co_return std::move(inserted).then([](auto&&) {
+    return Result<std::optional<AllocatedHandle>>(std::nullopt);
+  });
 }
 
 folly::coro::Task<Result<std::optional<ReadHandle>>> FlashCacheComponent::find(
@@ -75,8 +220,14 @@ FlashCacheComponent::FlashCacheComponent(std::string&& name,
     : name_(std::move(name)),
       cache_(std::make_unique<navy::BlockCache>(std::move(config))) {}
 
-folly::coro::Task<void> FlashCacheComponent::release(CacheItem& /* item */,
-                                                     bool /* inserted */) {
+folly::coro::Task<void> FlashCacheComponent::release(CacheItem& item,
+                                                     bool inserted) {
+  auto& fccItem = reinterpret_cast<FlashCacheItem&>(item);
+  if (!inserted) {
+    cache_->addHole(fccItem.getTotalSize());
+  }
+  cache_->regionManager_.close(std::move(fccItem.getRegionDescriptor()));
+  delete &item;
   co_return;
 }
 

cachelib/interface/components/FlashCacheComponent.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include "cachelib/interface/CacheComponent.h"
+#include "cachelib/interface/utils/CoroFiberAdapter.h"
 #include "cachelib/navy/block_cache/BlockCache.h"
 
 namespace facebook::cachelib::interface {
@@ -66,6 +67,24 @@ class FlashCacheComponent : public CacheComponent {
 
   FlashCacheComponent(std::string&& name, navy::BlockCache::Config&& config);
 
+  // Runs func() on a RegionManager worker fiber. Should not be called from an
+  // existing region manager worker fiber!
+  template <typename FuncT,
+            typename CleanupFuncT = std::function<void()>,
+            typename ReturnT = std::invoke_result_t<FuncT>>
+  folly::coro::Task<ReturnT> onWorkerThread(FuncT&& func,
+                                            CleanupFuncT&& cleanup = {}) {
+    XDCHECK(!cache_->regionManager_.isOnWorker())
+        << "Calling public APIs from a worker thread is unsupported";
+    co_return co_await utils::onWorkerThread(
+        cache_->regionManager_.getNextWorker(),
+        std::forward<FuncT>(func),
+        std::forward<CleanupFuncT>(cleanup));
+  }
+
+  folly::coro::Task<UnitResult> insertImpl(AllocatedHandle&& handle,
+                                           bool allowReplace);
+
   // ------------------------------ Interface ------------------------------ //
 
   folly::coro::Task<void> release(CacheItem& item, bool inserted) override;