diff --git a/cachelib/allocator/CMakeLists.txt b/cachelib/allocator/CMakeLists.txt index f3533ede4..63b4a423e 100644 --- a/cachelib/allocator/CMakeLists.txt +++ b/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) diff --git a/cachelib/allocator/CacheAllocator.h b/cachelib/allocator/CacheAllocator.h index b485a5fba..4dc535c07 100644 --- a/cachelib/allocator/CacheAllocator.h +++ b/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 keys, + folly::Range 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::find(typename Item::Key key) { return findImpl(key, AccessMode::kRead); } +template +void CacheAllocator::findMany(folly::Range keys, + folly::Range 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 void CacheAllocator::markUseful(const ReadHandle& handle, AccessMode mode) { diff --git a/cachelib/allocator/ChainedHashTable.h b/cachelib/allocator/ChainedHashTable.h index 9c64f0e62..18d8bb2f8 100644 --- a/cachelib/allocator/ChainedHashTable.h +++ b/cachelib/allocator/ChainedHashTable.h @@ -18,7 +18,10 @@ #include #include +#include +#include +#include #include #include #include @@ -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(hashTable_.get() + bucket), + __ATOMIC_RELAXED); + } + compressor_.prefetch(CompressedPtrType{ + static_cast(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 + void findMany(folly::Range keys, + folly::Range 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::find( return handleMaker_(ht_.findInBucket(key, bucket)); } +template T::* HookPtr, + typename LockT> +template +void ChainedHashTable::Container::findMany( + folly::Range keys, folly::Range 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 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 T::* HookPtr, typename LockT> diff --git a/cachelib/allocator/memory/CompressedPtr.h b/cachelib/allocator/memory/CompressedPtr.h index 93d4502b8..f992764b1 100644 --- a/cachelib/allocator/memory/CompressedPtr.h +++ b/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(compressed, + false /* isMultiTiered */); + } + bool operator==(const PtrCompressor& rhs) const noexcept { return &allocator_ == &rhs.allocator_; } diff --git a/cachelib/allocator/memory/MemoryAllocator.h b/cachelib/allocator/memory/MemoryAllocator.h index 3715e5e57..c7c81839d 100644 --- a/cachelib/allocator/memory/MemoryAllocator.h +++ b/cachelib/allocator/memory/MemoryAllocator.h @@ -621,6 +621,14 @@ class MemoryAllocator { return slabAllocator_.unCompress(cPtr, isMultiTiered); } + // Prefetches the memory of a possibly stale compressed pointer, see + // SlabAllocator::prefetch(). + template + void prefetch(const CompressedPtrType& cPtr, + bool isMultiTiered) const noexcept { + slabAllocator_.prefetch(cPtr, isMultiTiered); + } + // a special implementation of pointer compression for benchmarking purposes. CompressedPtr4B CACHELIB_INLINE compressAlt(const void* ptr) const { return slabAllocator_.compressAlt(ptr); diff --git a/cachelib/allocator/memory/SlabAllocator.h b/cachelib/allocator/memory/SlabAllocator.h index 59c84dd0d..d9fb303eb 100644 --- a/cachelib/allocator/memory/SlabAllocator.h +++ b/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 + 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; diff --git a/cachelib/allocator/memory/tests/MemoryAllocatorTest.cpp b/cachelib/allocator/memory/tests/MemoryAllocatorTest.cpp index 6bea8e697..3340ac8c7 100644 --- a/cachelib/allocator/memory/tests/MemoryAllocatorTest.cpp +++ b/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 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( + 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; diff --git a/cachelib/allocator/nvmcache/tests/NvmCacheTests.cpp b/cachelib/allocator/nvmcache/tests/NvmCacheTests.cpp index 3a9ae5f93..a02bc3eb4 100644 --- a/cachelib/allocator/nvmcache/tests/NvmCacheTests.cpp +++ b/cachelib/allocator/nvmcache/tests/NvmCacheTests.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -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 keys{"dram", "nvm", "missing"}; + std::vector 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(handles[0]->getMemory())); + ASSERT_NE(nullptr, handles[1]); + EXPECT_EQ(LruAllocator::Key{"nvm"}, handles[1]->getKey()); + EXPECT_EQ('n', *reinterpret_cast(handles[1]->getMemory())); + EXPECT_EQ(nullptr, handles[2]); +} + TEST_F(NvmCacheTest, CouldExistFast) { // Enable fast negative lookup this->makeCache(); diff --git a/cachelib/allocator/tests/FindManyTest.cpp b/cachelib/allocator/tests/FindManyTest.cpp new file mode 100644 index 000000000..bc73dff9c --- /dev/null +++ b/cachelib/allocator/tests/FindManyTest.cpp @@ -0,0 +1,248 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * 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. + */ + +#include +#include + +#include +#include +#include +#include +#include + +#include "cachelib/allocator/CacheAllocator.h" +#include "cachelib/common/Time.h" + +namespace facebook::cachelib::tests { +namespace { + +using AllocatorT = LruAllocator; +using ReadHandle = AllocatorT::ReadHandle; +using Key = AllocatorT::Key; + +constexpr uint32_t kValueSize = 100; + +std::string keyOf(size_t i) { return fmt::format("key_{}", i); } + +class FindManyTest : public testing::Test { + protected: + void SetUp() override { + AllocatorT::Config config; + config.setCacheSize(64 * Slab::kSize); + // A small table keeps several items per bucket chain. + config.setAccessConfig({10 /* bucketsPower */, 5 /* locksPower */}); + // The reaper would remove the item that Expired looks up. + config.reaperInterval = std::chrono::seconds(0); + cache_ = std::make_unique(config); + pool_ = + cache_->addPool("default", cache_->getCacheMemoryStats().ramCacheSize); + } + + void insert(const std::string& key, + uint32_t ttlSecs = 0, + uint32_t creationTime = 0) { + auto handle = + cache_->allocate(pool_, key, kValueSize, ttlSecs, creationTime); + ASSERT_NE(nullptr, handle); + cache_->insertOrReplace(handle); + } + + std::unique_ptr cache_; + PoolId pool_{}; +}; + +TEST_F(FindManyTest, MatchesFind) { + constexpr size_t kItems = 2000; + for (size_t i = 0; i < kItems; ++i) { + insert(keyOf(i)); + } + + // Half hits, half misses, with duplicates; batch sizes around multiples of + // the internal prefetch group of 8 keys. + for (size_t batch : {0, 1, 7, 8, 9, 15, 16, 17, 33, 100, 500}) { + std::vector storage; + for (size_t i = 0; i < batch; ++i) { + storage.push_back(keyOf(folly::Random::rand32(2 * kItems))); + } + std::vector keys(storage.begin(), storage.end()); + std::vector handles(batch); + cache_->findMany(folly::range(keys), folly::range(handles)); + + for (size_t i = 0; i < batch; ++i) { + auto expected = cache_->find(keys[i]); + ASSERT_EQ(expected.get(), handles[i].get()) + << "batch " << batch << ", key " << storage[i]; + if (handles[i]) { + EXPECT_EQ(keys[i], handles[i]->getKey()); + } + } + } +} + +TEST_F(FindManyTest, HoldsReferences) { + insert("a"); + std::vector keys{"a", "a"}; + std::vector handles(2); + cache_->findMany(folly::range(keys), folly::range(handles)); + ASSERT_NE(nullptr, handles[0]); + EXPECT_EQ(2, handles[0]->getRefCount()); + handles.clear(); + EXPECT_EQ(1, cache_->find("a")->getRefCount()); +} + +TEST_F(FindManyTest, Expired) { + insert("live"); + insert("expired", 1 /* ttlSecs */, util::getCurrentTimeSec() - 10); + const auto before = cache_->getGlobalCacheStats(); + + std::vector keys{"expired", "live"}; + std::vector handles(keys.size()); + cache_->findMany(folly::range(keys), folly::range(handles)); + EXPECT_EQ(nullptr, handles[0]); + EXPECT_TRUE(handles[0].wasExpired()); + ASSERT_NE(nullptr, handles[1]); + EXPECT_FALSE(handles[1].wasExpired()); + + const auto after = cache_->getGlobalCacheStats(); + EXPECT_EQ(1, after.numCacheGetExpiries - before.numCacheGetExpiries); + EXPECT_EQ(1, after.numCacheGetMiss - before.numCacheGetMiss); +} + +TEST_F(FindManyTest, StatsMatchFind) { + insert("a"); + insert("b"); + insert("expired", 1 /* ttlSecs */, util::getCurrentTimeSec() - 10); + const std::vector keys{"a", "x", "expired", "b", "y", "a", "expired"}; + const auto deltas = [](const GlobalCacheStats& from, + const GlobalCacheStats& to) { + return std::make_tuple(to.numCacheGets - from.numCacheGets, + to.numCacheGetMiss - from.numCacheGetMiss, + to.numCacheGetExpiries - from.numCacheGetExpiries); + }; + + const auto beforeFind = cache_->getGlobalCacheStats(); + for (const auto& key : keys) { + cache_->find(key); + } + const auto afterFind = cache_->getGlobalCacheStats(); + + std::vector handles(keys.size()); + cache_->findMany(folly::range(keys), folly::range(handles)); + const auto afterFindMany = cache_->getGlobalCacheStats(); + + EXPECT_EQ(deltas(beforeFind, afterFind), deltas(afterFind, afterFindMany)); + EXPECT_EQ(std::make_tuple(uint64_t{7}, uint64_t{4}, uint64_t{2}), + deltas(afterFind, afterFindMany)); +} + +// Overwriting a handle may release the last reference to a removed item, and +// its remove callback may use the cache. +TEST_F(FindManyTest, OverwrittenHandleRunsCallbackUnlocked) { + AllocatorT* cache = nullptr; + AllocatorT::Config config; + config.setCacheSize(64 * Slab::kSize); + config.setAccessConfig({10 /* bucketsPower */, 0 /* locksPower */}); + config.setRemoveCallback([&](const AllocatorT::RemoveCbData& data) { + if (data.item.getKey() == "removed") { + cache->remove("other"); + } + }); + AllocatorT allocator(config); + cache = &allocator; + const auto pool = + cache->addPool("default", cache->getCacheMemoryStats().ramCacheSize); + for (const std::string key : {"removed", "other", "key"}) { + auto handle = cache->allocate(pool, key, kValueSize); + ASSERT_NE(nullptr, handle); + cache->insertOrReplace(handle); + } + + std::vector handles; + handles.push_back(cache->find("removed")); + cache->remove("removed"); + std::vector keys{"key"}; + cache->findMany(folly::range(keys), folly::range(handles)); + ASSERT_NE(nullptr, handles[0]); + EXPECT_EQ(nullptr, cache->find("other")); +} + +TEST_F(FindManyTest, InvalidArguments) { + insert("a"); + std::vector keys{"a", "b"}; + + std::vector wrongSize(1); + EXPECT_THROW(cache_->findMany(folly::range(keys), folly::range(wrongSize)), + std::invalid_argument); +} + +// Lookups race with inserts and removes in the same buckets; every handle +// returned must be for the requested key. +TEST_F(FindManyTest, ConcurrentMutations) { + constexpr size_t kKeys = 4000; + for (size_t i = 0; i < kKeys; i += 2) { + insert(keyOf(i)); + } + + std::atomic stop{false}; + std::vector writers; + for (int t = 0; t < 2; ++t) { + writers.emplace_back([&] { + while (!stop.load(std::memory_order_relaxed)) { + const auto key = keyOf(folly::Random::rand32(kKeys)); + if (folly::Random::oneIn(2)) { + if (auto handle = cache_->allocate(pool_, key, kValueSize)) { + cache_->insertOrReplace(handle); + } + } else { + cache_->remove(key); + } + } + }); + } + + std::atomic found{0}; + std::vector readers; + for (int t = 0; t < 4; ++t) { + readers.emplace_back([&] { + for (int iter = 0; iter < 2000; ++iter) { + std::vector storage; + for (size_t i = 0; i < 40; ++i) { + storage.push_back(keyOf(folly::Random::rand32(kKeys))); + } + std::vector keys(storage.begin(), storage.end()); + std::vector handles(keys.size()); + cache_->findMany(folly::range(keys), folly::range(handles)); + for (size_t i = 0; i < keys.size(); ++i) { + if (handles[i]) { + ASSERT_EQ(keys[i], handles[i]->getKey()); + found.fetch_add(1, std::memory_order_relaxed); + } + } + } + }); + } + for (auto& t : readers) { + t.join(); + } + stop = true; + for (auto& t : writers) { + t.join(); + } + EXPECT_GT(found.load(), 0); +} + +} // namespace +} // namespace facebook::cachelib::tests diff --git a/cachelib/benchmarks/CMakeLists.txt b/cachelib/benchmarks/CMakeLists.txt index d417f1d55..eb3db748a 100644 --- a/cachelib/benchmarks/CMakeLists.txt +++ b/cachelib/benchmarks/CMakeLists.txt @@ -50,6 +50,7 @@ if (BUILD_TESTS) add_source_test (SListBench.cpp) add_source_test (ThreadLocalBench.cpp) add_source_test (EventTrackerPerf.cpp) + add_source_test (FindManyBench.cpp) add_source_test (StrictAliasingSafeReadBench.cpp) # Temporarily disabled test: require __rdstc() #add_source_test (CacheAllocatorOpsMicroBench.cpp) diff --git a/cachelib/benchmarks/FindManyBench.cpp b/cachelib/benchmarks/FindManyBench.cpp new file mode 100644 index 000000000..8b1d43dec --- /dev/null +++ b/cachelib/benchmarks/FindManyBench.cpp @@ -0,0 +1,272 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * 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. + */ + +// find() per key vs findMany() on a warmed DRAM cache holding all items: +// Zipfian keys in windows of --batch distinct keys, every hit reads its value. +// Reports CPU time per key for each pass. +// +// Example: +// FindManyBench --mode=find --num_keys=40000000 --size_gb=36 --threads=16 +// FindManyBench --mode=findMany --num_keys=40000000 --size_gb=36 --threads=16 + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cachelib/allocator/CacheAllocator.h" + +using namespace facebook::cachelib; + +DEFINE_string(mode, "findMany", "find or findMany"); +DEFINE_uint64(num_keys, 1 << 22, "Number of items in the cache"); +DEFINE_uint64(size_gb, 4, "Cache size in GB; must hold every item"); +DEFINE_uint32(value_size, 500, "Value size in bytes"); +DEFINE_uint32(threads, 4, "Lookup threads"); +DEFINE_uint32(batch, 32, "Distinct keys per lookup window"); +DEFINE_double(zipf_theta, 0.99, "Zipfian skew, in [0, 1); 0 is uniform"); +DEFINE_uint64(keys_per_pass, 1 << 21, "Keys looked up per thread per pass"); +DEFINE_uint32(passes, 4, "Measured passes after one warmup pass"); +DEFINE_uint64(seed, 1, "Random seed"); + +namespace { + +using Cache = LruAllocator; +using KeyBytes = std::array; + +KeyBytes makeKey(uint64_t id) { + uint64_t hash = id + 0x9e3779b97f4a7c15ULL; + hash = (hash ^ (hash >> 30)) * 0xbf58476d1ce4e5b9ULL; + hash = (hash ^ (hash >> 27)) * 0x94d049bb133111ebULL; + hash ^= hash >> 31; + KeyBytes key; + std::memcpy(key.data(), &hash, 8); + std::memcpy(key.data() + 8, &id, 8); + return key; +} + +Cache::Key asKey(const KeyBytes& key) { return {key.data(), key.size()}; } + +// YCSB ZipfianGenerator (Gray et al.); rank 0 is the most popular item. +class Zipf { + public: + Zipf(uint64_t items, double theta) : items_(items) { + double zetan = 0; + for (uint64_t i = 1; i <= items; ++i) { + zetan += 1.0 / std::pow(double(i), theta); + } + const double zeta2 = 1.0 + 1.0 / std::pow(2.0, theta); + zetan_ = zetan; + alpha_ = 1.0 / (1.0 - theta); + eta_ = (1.0 - std::pow(2.0 / items, 1.0 - theta)) / (1.0 - zeta2 / zetan); + half_ = 1.0 + std::pow(0.5, theta); + } + + template + uint64_t operator()(R& random) const { + const double u = std::uniform_real_distribution(0.0, 1.0)(random); + const double uz = u * zetan_; + if (uz < 1.0) { + return 0; + } + if (uz < half_) { + return 1; + } + return std::min( + items_ - 1, uint64_t(items_ * std::pow(eta_ * u - eta_ + 1.0, alpha_))); + } + + private: + uint64_t items_; + double zetan_ = 0, alpha_ = 0, eta_ = 0, half_ = 0; +}; + +uint64_t threadCpuNs() { + timespec value{}; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &value); + return value.tv_sec * 1'000'000'000ULL + value.tv_nsec; +} + +template +void parallelFor(uint32_t threads, uint64_t n, F&& f) { + std::vector pool; + for (uint32_t t = 0; t < threads; ++t) { + pool.emplace_back([&, t] { + for (uint64_t i = n * t / threads; i < n * (t + 1) / threads; ++i) { + f(i); + } + }); + } + for (auto& thread : pool) { + thread.join(); + } +} + +struct alignas(64) Worker { + std::vector cpuNs; + uint64_t hits = 0; + uint64_t checksum = 0; +}; + +void consume(const Cache::ReadHandle& handle, Worker& worker) { + if (!handle) { + return; + } + ++worker.hits; + const auto* data = static_cast(handle->getMemory()); + for (size_t i = 0; i < handle->getSize(); i += 64) { + worker.checksum += data[i]; + } +} + +} // namespace + +int main(int argc, char** argv) { + const folly::Init init(&argc, &argv); + const bool many = FLAGS_mode == "findMany"; + XCHECK(many || FLAGS_mode == "find") << "--mode must be find or findMany"; + XCHECK(FLAGS_batch > 0 && FLAGS_keys_per_pass % FLAGS_batch == 0); + XCHECK(FLAGS_zipf_theta >= 0 && FLAGS_zipf_theta < 1); + XCHECK_LT(FLAGS_num_keys, 1ULL << 32); + XCHECK_LE(FLAGS_batch, FLAGS_num_keys); + + Cache::Config config; + config.setCacheSize(FLAGS_size_gb << 30); + config.setAccessConfig(FLAGS_num_keys); + config.reaperInterval = std::chrono::seconds(0); + config.poolRebalanceInterval = std::chrono::seconds(0); + Cache cache(config); + Cache::MMConfig mmConfig; + mmConfig.defaultLruRefreshTime = 900; + mmConfig.lruRefreshTime = 900; + const auto pool = cache.addPool( + "default", cache.getCacheMemoryStats().ramCacheSize, {}, mmConfig); + + XLOG(INFO) << "populating " << FLAGS_num_keys << " items"; + std::atomic failures{0}; + parallelFor(FLAGS_threads, FLAGS_num_keys, [&](uint64_t id) { + auto handle = cache.allocate(pool, asKey(makeKey(id)), FLAGS_value_size); + if (!handle) { + failures.fetch_add(1, std::memory_order_relaxed); + return; + } + std::memset(handle->getMemory(), 'x', FLAGS_value_size); + cache.insertOrReplace(handle); + }); + XCHECK_EQ(0u, failures.load()) << "cache too small, increase --size_gb"; + // MMLru promotes an item on its first read regardless of the refresh time. + parallelFor(FLAGS_threads, FLAGS_num_keys, [&](uint64_t id) { + XCHECK(cache.find(asKey(makeKey(id)))) + << "item evicted, increase --size_gb"; + }); + + std::vector rankToItem(FLAGS_num_keys); + std::iota(rankToItem.begin(), rankToItem.end(), uint32_t(0)); + std::shuffle(rankToItem.begin(), rankToItem.end(), + std::mt19937_64(FLAGS_seed ^ 0x5eed)); + const Zipf zipf(FLAGS_num_keys, FLAGS_zipf_theta); + + const uint32_t passes = FLAGS_passes + 1; + std::vector workers(FLAGS_threads); + std::vector threads; + std::latch ready(FLAGS_threads), start(1); + for (uint32_t t = 0; t < FLAGS_threads; ++t) { + threads.emplace_back([&, t] { + auto& worker = workers[t]; + std::mt19937_64 random(FLAGS_seed * 1000003 + t); + std::vector sequence(FLAGS_keys_per_pass * passes); + std::vector window(FLAGS_batch); + for (uint64_t base = 0; base < sequence.size(); base += FLAGS_batch) { + for (uint32_t i = 0; i < FLAGS_batch; ++i) { + uint32_t item; + do { + item = rankToItem[zipf(random)]; + } while (std::find(window.begin(), window.begin() + i, item) != + window.begin() + i); + window[i] = item; + sequence[base + i] = makeKey(item); + } + } + std::vector keys; + keys.reserve(sequence.size()); + for (const auto& key : sequence) { + keys.push_back(asKey(key)); + } + std::vector handles(FLAGS_batch); + worker.cpuNs.resize(passes); + ready.count_down(); + start.wait(); + + const Cache::Key* next = keys.data(); + for (uint32_t pass = 0; pass < passes; ++pass) { + const auto started = threadCpuNs(); + for (uint64_t done = 0; done < FLAGS_keys_per_pass; + done += FLAGS_batch) { + if (many) { + cache.findMany(folly::Range(next, FLAGS_batch), + folly::range(handles)); + for (auto& handle : handles) { + consume(handle, worker); + handle.reset(); + } + } else { + for (uint32_t i = 0; i < FLAGS_batch; ++i) { + consume(cache.find(next[i]), worker); + } + } + next += FLAGS_batch; + } + worker.cpuNs[pass] = threadCpuNs() - started; + } + }); + } + ready.wait(); + XLOG(INFO) << "running " << FLAGS_mode; + start.count_down(); + for (auto& thread : threads) { + thread.join(); + } + + const uint64_t expected = FLAGS_keys_per_pass * passes; + const uint64_t perHit = ((FLAGS_value_size + 63) / 64) * uint64_t('x'); + for (const auto& worker : workers) { + XCHECK_EQ(expected, worker.hits); + XCHECK_EQ(expected * perHit, worker.checksum); + } + for (uint32_t pass = 1; pass < passes; ++pass) { + uint64_t cpuNs = 0; + for (const auto& worker : workers) { + cpuNs += worker.cpuNs[pass]; + } + const double keys = double(FLAGS_keys_per_pass) * FLAGS_threads; + XLOG(INFO) << fmt::format( + "{} pass {}: {:.1f} ns CPU per key, {:.2f}M keys/s per core", + FLAGS_mode, pass, cpuNs / keys, keys / cpuNs * 1e3); + } + return 0; +}