Add CacheAllocator::findMany() for batched DRAM lookups

cachelib/allocator/CMakeLists.txt

@@ -141,6 +141,7 @@ if (BUILD_TESTS)
   # timeout from ctest: add_source_test (tests/RebalanceStrategyTest.cpp)
   # failure from ctest: add_source_test (tests/AllocatorTypeTest.cpp)
   add_source_test (tests/ChainedHashTest.cpp)
+  add_source_test (tests/FindManyTest.cpp)
   # timeout from ctest: add_source_test (tests/AllocatorResizeTypeTest.cpp)
   # timeout from ctest: add_source_test (tests/AllocatorHitStatsTypeTest.cpp)
   add_source_test (tests/AllocatorMemoryTiersTest.cpp)

cachelib/allocator/CacheAllocator.h

@@ -590,6 +590,22 @@ class CacheAllocator : public CacheBase {
   //                  key does not exist.
   ReadHandle find(Key key);
 
+  // Looks up several keys at once: handles[i] gets what find(keys[i]) would
+  // return, but the DRAM hash table and item cache misses of different keys
+  // overlap. As with find(), keys missing or expired in DRAM are then looked
+  // up in NVM, and those handles may not be ready yet. Stats and the clock are
+  // updated once per batch. The keys must stay valid during the call, so they
+  // must not point into items held only by the handles being overwritten.
+  //
+  // @param keys     the keys to look up
+  // @param handles  output, the same size as keys; overwritten
+  //
+  // @throw std::invalid_argument if the sizes differ
+  // @throw std::overflow_error if the maximum item refcount is exceeded by
+  //        creating a handle; the handles must then be discarded
+  void findMany(folly::Range<const Key*> keys,
+                folly::Range<ReadHandle*> handles);
+
   // Warning: this API is synchronous today with HybridCache. This means as
   //          opposed to find(), we will block on an item being read from
   //          flash until it is loaded into DRAM-cache. In find(), if an item
@@ -5075,6 +5091,63 @@ CacheAllocator<CacheTrait>::find(typename Item::Key key) {
   return findImpl(key, AccessMode::kRead);
 }
 
+template <typename CacheTrait>
+void CacheAllocator<CacheTrait>::findMany(folly::Range<const Key*> keys,
+                                          folly::Range<ReadHandle*> handles) {
+  if (keys.size() != handles.size()) {
+    throw std::invalid_argument(fmt::format("findMany: {} keys but {} handles",
+                                            keys.size(), handles.size()));
+  }
+
+  stats_.numCacheGets.add(keys.size());
+  accessContainer_->findMany(keys, handles);
+
+  // The per-key logic of findInternalWithExpiration(), but the counters and
+  // the clock are updated once per batch.
+  const uint32_t now = util::getCurrentTimeSec();
+  uint64_t misses = 0;
+  uint64_t expiries = 0;
+  for (size_t i = 0; i < keys.size(); ++i) {
+    auto& handle = handles[i];
+    auto* const item = handle.getInternal();
+    if (UNLIKELY(!item)) {
+      ++misses;
+      recordEvent(AllocatorApiEvent::FIND, keys[i],
+                  AllocatorApiResult::NOT_FOUND);
+      if (nvmCache_) {
+        handle = nvmCache_->find(HashedKey{keys[i]});
+      }
+      continue;
+    }
+    if (UNLIKELY(item->isExpired(now))) {
+      ++misses;
+      ++expiries;
+      recordEvent(AllocatorApiEvent::FIND, keys[i], AllocatorApiResult::EXPIRED,
+                  handle);
+      handle = {};
+      handle.markExpired();
+      if (nvmCache_) {
+        handle = nvmCache_->find(HashedKey{keys[i]});
+      }
+      continue;
+    }
+    recordEvent(AllocatorApiEvent::FIND, keys[i], AllocatorApiResult::FOUND,
+                handle);
+    // Another thread may have accessed the item after the clock was read.
+    const auto lastAccess = item->getLastAccessTime();
+    if (lastAccess > 0) {
+      handle.setTTASecs(now > lastAccess ? now - lastAccess : 0);
+    }
+    markUseful(handle, AccessMode::kRead);
+  }
+  if (misses) {
+    stats_.numCacheGetMiss.add(misses);
+  }
+  if (expiries) {
+    stats_.numCacheGetExpiries.add(expiries);
+  }
+}
+
 template <typename CacheTrait>
 void CacheAllocator<CacheTrait>::markUseful(const ReadHandle& handle,
                                             AccessMode mode) {

cachelib/allocator/ChainedHashTable.h

@@ -18,7 +18,10 @@
 
 #include <fmt/core.h>
 #include <folly/Optional.h>
+#include <folly/synchronization/SanitizeThread.h>
 
+#include <algorithm>
+#include <array>
 #include <cstdint>
 #include <functional>
 #include <map>
@@ -151,6 +154,30 @@ class ChainedHashTable {
     // gets the bucket for the key by using the corresponding hash function.
     BucketId getBucket(Key k) const noexcept;
 
+    // Prefetches the bucket slot. Batched lookups call this for every key
+    // first so that the slot cache misses of the whole batch overlap.
+    void prefetchBucket(BucketId bucket) const noexcept {
+      __builtin_prefetch(hashTable_.get() + bucket);
+    }
+
+    // Prefetches the first node of the bucket's chain. The slot is read without
+    // the bucket lock, so the pointer may be stale and is only prefetched.
+    // Only for 4-byte pointers, as a wider read could tear; the table is
+    // allocated or mapped with at least 4-byte alignment.
+    void prefetchBucketHead([[maybe_unused]] BucketId bucket) const noexcept {
+      if constexpr (sizeof(CompressedPtrType) == sizeof(uint32_t)) {
+        uint32_t raw;
+        {
+          folly::annotate_ignore_thread_sanitizer_guard g(__FILE__, __LINE__);
+          raw = __atomic_load_n(
+              reinterpret_cast<const uint32_t*>(hashTable_.get() + bucket),
+              __ATOMIC_RELAXED);
+        }
+        compressor_.prefetch(CompressedPtrType{
+            static_cast<typename CompressedPtrType::SerializedPtrType>(raw)});
+      }
+    }
+
     // Call 'func' on each element in the given bucket.
     //
     // @param bucket  the bucket id to fetch.
@@ -516,6 +543,18 @@ class ChainedHashTable {
     //        creating this item handle.
     Handle find(Key key) const;
 
+    // Batched find(): handles[i] = find(keys[i]), with the bucket slots and
+    // chain heads prefetched ahead of the lookups.
+    //
+    // @param keys     keys to look up
+    // @param handles  output, same size as keys
+    //
+    // @throw std::overflow_error if the maximum item refcount is exceeded by
+    //        creating a handle
+    template <typename OutputHandle>
+    void findMany(folly::Range<const Key*> keys,
+                  folly::Range<OutputHandle*> handles) const;
+
     // for saving the state of the hash table
     //
     // precondition:  serialization must happen without any reader or writer
@@ -1318,6 +1357,46 @@ typename T::Handle ChainedHashTable::Container<T, HookPtr, LockT>::find(
   return handleMaker_(ht_.findInBucket(key, bucket));
 }
 
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+template <typename OutputHandle>
+void ChainedHashTable::Container<T, HookPtr, LockT>::findMany(
+    folly::Range<const Key*> keys, folly::Range<OutputHandle*> handles) const {
+  XDCHECK_EQ(keys.size(), handles.size());
+  // Bucket slots are prefetched one group of keys ahead of the chain heads,
+  // which can only be decoded once their slot has arrived. Groups of 8
+  // measured fastest.
+  constexpr size_t kGroup = 8;
+  std::array<BucketId, 2 * kGroup> buckets; // the current and the next group
+  const size_t n = keys.size();
+  auto prefetchSlots = [&](size_t first) {
+    for (size_t i = first; i < std::min(n, first + kGroup); ++i) {
+      buckets[i % buckets.size()] = ht_.getBucket(keys[i]);
+      ht_.prefetchBucket(buckets[i % buckets.size()]);
+    }
+  };
+  prefetchSlots(0);
+  for (size_t group = 0; group < n; group += kGroup) {
+    prefetchSlots(group + kGroup);
+    const size_t end = std::min(n, group + kGroup);
+    for (size_t i = group; i < end; ++i) {
+      ht_.prefetchBucketHead(buckets[i % buckets.size()]);
+    }
+    for (size_t i = group; i < end; ++i) {
+      const BucketId bucket = buckets[i % buckets.size()];
+      Handle handle;
+      {
+        auto l = locks_.lockShared(bucket);
+        handle = handleMaker_(ht_.findInBucket(keys[i], bucket));
+      }
+      // The old handle may hold the last reference to a removed item, whose
+      // release runs callbacks that must not run under the bucket lock.
+      handles[i] = std::move(handle);
+    }
+  }
+}
+
 template <typename T,
           typename ChainedHashTable::Hook<T> T::* HookPtr,
           typename LockT>

cachelib/allocator/memory/CompressedPtr.h

@@ -369,6 +369,12 @@ class PtrCompressor {
             compressed, false /* isMultiTiered */));
   }
 
+  // Prefetches the memory of a compressed pointer that may be stale.
+  void prefetch(const CompressedPtrType& compressed) const noexcept {
+    allocator_.template prefetch<CompressedPtrType>(compressed,
+                                                    false /* isMultiTiered */);
+  }
+
   bool operator==(const PtrCompressor& rhs) const noexcept {
     return &allocator_ == &rhs.allocator_;
   }

cachelib/allocator/memory/MemoryAllocator.h

@@ -621,6 +621,14 @@ class MemoryAllocator {
     return slabAllocator_.unCompress<CompressedPtrType>(cPtr, isMultiTiered);
   }
 
+  // Prefetches the memory of a possibly stale compressed pointer, see
+  // SlabAllocator::prefetch().
+  template <typename CompressedPtrType>
+  void prefetch(const CompressedPtrType& cPtr,
+                bool isMultiTiered) const noexcept {
+    slabAllocator_.prefetch<CompressedPtrType>(cPtr, isMultiTiered);
+  }
+
   // a special implementation of pointer compression for benchmarking purposes.
   CompressedPtr4B CACHELIB_INLINE compressAlt(const void* ptr) const {
     return slabAllocator_.compressAlt(ptr);

cachelib/allocator/memory/SlabAllocator.h

@@ -309,6 +309,29 @@ class SlabAllocator {
     return slab->memoryAtOffset(offset);
   }
 
+  // Prefetches the memory of a compressed pointer that may be stale: its slab
+  // may have been released or given to another allocation class since. Unlike
+  // unCompress(), never checks the pointer; out-of-slab offsets are skipped.
+  template <typename CompressedPtrType>
+  void prefetch(const CompressedPtrType& ptr,
+                bool isMultiTiered) const noexcept {
+    if (ptr.isNull()) {
+      return;
+    }
+    const SlabIdx slabIndex = ptr.getSlabIdx(isMultiTiered);
+    uint32_t allocSize;
+    {
+      // allocSize may be 0 or belong to another class; the offset check below
+      // keeps the address inside the slab.
+      folly::annotate_ignore_thread_sanitizer_guard g(__FILE__, __LINE__);
+      allocSize = getSlabHeader(slabIndex)->getAllocSize();
+    }
+    const uint64_t offset = uint64_t{allocSize} * ptr.getAllocIdx();
+    if (offset < Slab::kSize) {
+      __builtin_prefetch(slabMemoryStart_[slabIndex].memoryAtOffset(offset));
+    }
+  }
+
   // a special implementation of pointer compression for benchmarking purposes.
   CompressedPtr4B compressAlt(const void* ptr) const;
   void* unCompressAlt(const CompressedPtr4B ptr) const;

cachelib/allocator/memory/tests/MemoryAllocatorTest.cpp

@@ -494,6 +494,31 @@ TEST_F(MemoryAllocatorTest, GrowShrinkPool) {
   ASSERT_TRUE(m.growPool(p1, Slab::kSize));
 }
 
+// prefetch() takes pointers that may be stale: their slab may since have been
+// released or given to an allocation class of another size. Unlike
+// unCompress(), it must not check them; in debug builds unCompress() and
+// memoryAtOffset() reject these pointers.
+TEST_F(MemoryAllocatorTest, PrefetchStale) {
+  const size_t totalSize = 5 * Slab::kSize;
+  void* memory = allocate(totalSize);
+  const std::set<uint32_t> sizes{1024};
+  MemoryAllocator m(getDefaultMemoryAllocatorConfig(sizes), memory, totalSize);
+  const auto pid = m.addPool("pool", 2 * Slab::kSize, sizes);
+  void* alloc = m.allocate(pid, 1024);
+  ASSERT_NE(nullptr, alloc);
+
+  // One slab holds 1024-byte allocations and the other is not carved yet.
+  const uint32_t pastSlabEnd = Slab::kSize / 1024;
+  for (const uint32_t slabIdx : {0u, 1u}) {
+    for (const uint32_t allocIdx : {1u, pastSlabEnd}) {
+      const CompressedPtr4B ptr{static_cast<int64_t>(
+          compress(CompressedPtr4B{}, slabIdx, allocIdx, false))};
+      m.prefetch(ptr, false /* isMultiTiered */);
+    }
+  }
+  m.free(alloc);
+}
+
 TEST_F(MemoryAllocatorTest, isAllocFreed) {
   const size_t numSlabs = 100;
   const size_t size = numSlabs * Slab::kSize;

cachelib/allocator/nvmcache/tests/NvmCacheTests.cpp

@@ -20,6 +20,7 @@
 
 #include <array>
 #include <climits>
+#include <cstring>
 #include <set>
 #include <thread>
 
@@ -137,6 +138,32 @@ TEST_F(NvmCacheTest, BasicGet) {
   ASSERT_TRUE(this->checkKeyExists(key, false /* ramOnly */));
 }
 
+TEST_F(NvmCacheTest, FindMany) {
+  auto& cache = this->cache();
+  for (const std::string key : {"dram", "nvm"}) {
+    auto it = cache.allocate(this->poolId(), key, 100);
+    ASSERT_NE(nullptr, it);
+    std::memset(it->getMemory(), key[0], 100);
+    cache.insertOrReplace(it);
+  }
+  ASSERT_TRUE(this->pushToNvmCacheFromRamForTesting("nvm"));
+  this->removeFromRamForTesting("nvm");
+  ASSERT_FALSE(this->checkKeyExists("nvm", true /* ramOnly */));
+
+  std::vector<LruAllocator::Key> keys{"dram", "nvm", "missing"};
+  std::vector<ReadHandle> handles(keys.size());
+  cache.findMany(folly::range(keys), folly::range(handles));
+  for (const auto& handle : handles) {
+    handle.wait();
+  }
+  ASSERT_NE(nullptr, handles[0]);
+  EXPECT_EQ('d', *reinterpret_cast<const char*>(handles[0]->getMemory()));
+  ASSERT_NE(nullptr, handles[1]);
+  EXPECT_EQ(LruAllocator::Key{"nvm"}, handles[1]->getKey());
+  EXPECT_EQ('n', *reinterpret_cast<const char*>(handles[1]->getMemory()));
+  EXPECT_EQ(nullptr, handles[2]);
+}
+
 TEST_F(NvmCacheTest, CouldExistFast) {
   // Enable fast negative lookup
   this->makeCache();