diff --git a/velox/common/caching/AsyncDataCache.cpp b/velox/common/caching/AsyncDataCache.cpp index c02b6f95928..6b9a63229d8 100644 --- a/velox/common/caching/AsyncDataCache.cpp +++ b/velox/common/caching/AsyncDataCache.cpp @@ -251,7 +251,8 @@ std::unique_ptr CacheShard::getFreeEntryLocked() { std::optional CacheShard::lookupLocked( RawFileCacheKey key, uint64_t size, - folly::SemiFuture* wait) { + folly::SemiFuture* wait, + CacheEntrySizePolicy sizePolicy) { ++eventCounter_; auto it = entryMap_.find(key); if (it == entryMap_.end()) { @@ -266,7 +267,8 @@ std::optional CacheShard::lookupLocked( return CachePin{}; } // size=0 (from find()) always passes since entry size is non-negative. - if (foundEntry->size() < size) { + if (sizePolicy == CacheEntrySizePolicy::kRequireAtLeast && + foundEntry->size() < size) { // This can happen if different load quanta apply to access via different // connectors. This is not an error but still worth logging. VELOX_CACHE_LOG_EVERY_MS(WARNING, 1'000) @@ -296,11 +298,12 @@ CachePin CacheShard::findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous, - folly::SemiFuture* wait) { + folly::SemiFuture* wait, + CacheEntrySizePolicy sizePolicy) { AsyncDataCacheEntry* entryToInit = nullptr; { std::lock_guard l(mutex_); - auto result = lookupLocked(key, size, wait); + auto result = lookupLocked(key, size, wait, sizePolicy); if (result.has_value()) { return std::move(result.value()); } @@ -874,9 +877,10 @@ CachePin AsyncDataCache::findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous, - folly::SemiFuture* wait) { + folly::SemiFuture* wait, + CacheEntrySizePolicy sizePolicy) { const int shard = std::hash()(key) & shardMask_; - return shards_[shard]->findOrCreate(key, size, contiguous, wait); + return shards_[shard]->findOrCreate(key, size, contiguous, wait, sizePolicy); } std::optional AsyncDataCache::find( diff --git a/velox/common/caching/AsyncDataCache.h b/velox/common/caching/AsyncDataCache.h index 40d9a4b4ed1..9bd8c59da94 100644 --- a/velox/common/caching/AsyncDataCache.h +++ b/velox/common/caching/AsyncDataCache.h @@ -120,6 +120,13 @@ struct RawFileCacheKey { return offset == other.offset && fileNum == other.fileNum; } }; + +/// Controls whether an entry that begins at the requested key but is shorter +/// than the requested size is stale or is a reusable prefix. +enum class CacheEntrySizePolicy : uint8_t { + kRequireAtLeast, + kAllowSmaller, +}; } // namespace facebook::velox::cache namespace std { @@ -625,12 +632,15 @@ class CacheShard { : cache_(cache), maxWriteRatio_(maxWriteRatio) {} /// See AsyncDataCache::findOrCreate. If 'contiguous' is true, the - /// entry's data is allocated as a single contiguous region. + /// entry's data is allocated as a single contiguous region. 'sizePolicy' + /// decides how an existing entry at 'key' that is shorter than 'size' is + /// treated. See AsyncDataCache::findOrCreate for details. CachePin findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous = false, - folly::SemiFuture* readyFuture = nullptr); + folly::SemiFuture* readyFuture = nullptr, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast); /// Finds a cache entry for 'key'. Returns a shared-mode pin if the entry /// exists and is not exclusive. Returns an empty pin (inside optional) if @@ -731,7 +741,8 @@ class CacheShard { std::optional lookupLocked( RawFileCacheKey key, uint64_t size, - folly::SemiFuture* waitFuture); + folly::SemiFuture* waitFuture, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast); void tryAddFreeEntry(std::unique_ptr&& entry); @@ -895,11 +906,22 @@ class AsyncDataCache : public memory::Cache { /// When the future is realized, the caller may retry findOrCreate(). /// runtime error with code kNoCacheSpace if there is no space to /// create the new entry after evicting any unpinned content. + /// + /// 'sizePolicy' decides how an existing shared entry at 'key' that is + /// shorter than 'size' is treated. With kRequireAtLeast (the default) the + /// entry is stale: it is evicted and a new exclusive entry of 'size' is + /// created for the caller to fill. With kAllowSmaller the shorter entry is + /// returned as a valid prefix and counted as a hit. The cache does not + /// track the uncovered suffix. The caller must read the returned entry's + /// size(), and request [key.offset + size(), key.offset + 'size') under a + /// separate key, which may in turn hit, miss, or return another shorter + /// prefix. CachePin findOrCreate( RawFileCacheKey key, uint64_t size, bool contiguous = false, - folly::SemiFuture* waitFuture = nullptr); + folly::SemiFuture* waitFuture = nullptr, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast); /// Finds a cache entry for 'key'. Returns a shared-mode pin if the entry /// exists and is not exclusive. Returns an empty pin (inside optional) if @@ -978,9 +1000,11 @@ class AsyncDataCache : public memory::Cache { const std::vector& keys, const SizeFunc& sizeFunc, const ProcessPin& processPin, - bool contiguous = false) { + bool contiguous = false, + CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast) { for (size_t i = 0; i < keys.size(); ++i) { - auto pin = findOrCreate(keys[i], sizeFunc(i), contiguous); + auto pin = + findOrCreate(keys[i], sizeFunc(i), contiguous, nullptr, sizePolicy); if (pin.empty() || pin.checkedEntry()->isShared()) { continue; } diff --git a/velox/common/caching/tests/AsyncDataCacheTest.cpp b/velox/common/caching/tests/AsyncDataCacheTest.cpp index 7df81482038..849cf1bc7d7 100644 --- a/velox/common/caching/tests/AsyncDataCacheTest.cpp +++ b/velox/common/caching/tests/AsyncDataCacheTest.cpp @@ -1304,6 +1304,87 @@ TEST_P(AsyncDataCacheTest, staleEntry) { ASSERT_EQ(stats.numHit, 1); } +TEST_P(AsyncDataCacheTest, cacheEntrySizePolicy) { + constexpr uint64_t kRamBytes = 1UL << 30; + initializeCache(kRamBytes, 0, 0); + StringIdLease file(fileIds(), std::string_view("cacheEntrySizePolicy")); + const uint64_t size = 200; + folly::SemiFuture wait(false); + RawFileCacheKey key{file.id(), 1000}; + + // A miss creates an exclusive entry of the requested size. + auto before = cache_->refreshStats(); + auto pin = cache_->findOrCreate( + key, + size, + /*contiguous=*/false, + &wait, + CacheEntrySizePolicy::kAllowSmaller); + ASSERT_FALSE(pin.empty()); + ASSERT_TRUE(wait.isReady()); + ASSERT_TRUE(pin.entry()->isExclusive()); + ASSERT_EQ(pin.entry()->size(), size); + pin.entry()->setExclusiveToShared(); + pin.clear(); + auto delta = cache_->refreshStats() - before; + ASSERT_EQ(delta.numHit, 0); + ASSERT_EQ(delta.numStales, 0); + + // kAllowSmaller returns a shorter entry as a shared prefix and counts a hit. + before = cache_->refreshStats(); + auto prefixPin = cache_->findOrCreate( + key, + 2 * size, + /*contiguous=*/false, + &wait, + CacheEntrySizePolicy::kAllowSmaller); + ASSERT_FALSE(prefixPin.empty()); + ASSERT_TRUE(wait.isReady()); + ASSERT_FALSE(prefixPin.entry()->isExclusive()); + ASSERT_EQ(prefixPin.entry()->size(), size); + prefixPin.clear(); + auto stats = cache_->refreshStats(); + ASSERT_EQ(stats.numEntries, 1); + delta = stats - before; + ASSERT_EQ(delta.numHit, 1); + ASSERT_EQ(delta.numStales, 0); + + // kAllowSmaller with a request that fits the entry is an ordinary hit. + before = cache_->refreshStats(); + auto exactPin = cache_->findOrCreate( + key, + size, + /*contiguous=*/false, + &wait, + CacheEntrySizePolicy::kAllowSmaller); + ASSERT_FALSE(exactPin.empty()); + ASSERT_FALSE(exactPin.entry()->isExclusive()); + ASSERT_EQ(exactPin.entry()->size(), size); + exactPin.clear(); + delta = cache_->refreshStats() - before; + ASSERT_EQ(delta.numHit, 1); + ASSERT_EQ(delta.numStales, 0); + + // kRequireAtLeast treats the same shorter entry as stale and replaces it. + before = cache_->refreshStats(); + auto stalePin = cache_->findOrCreate( + key, + 2 * size, + /*contiguous=*/false, + &wait, + CacheEntrySizePolicy::kRequireAtLeast); + ASSERT_FALSE(stalePin.empty()); + ASSERT_TRUE(wait.isReady()); + ASSERT_TRUE(stalePin.entry()->isExclusive()); + ASSERT_EQ(stalePin.entry()->size(), 2 * size); + stalePin.entry()->setExclusiveToShared(); + stats = cache_->refreshStats(); + ASSERT_EQ(stats.numEntries, 1); + delta = stats - before; + ASSERT_EQ(delta.numHit, 0); + ASSERT_EQ(delta.numStales, 1); +} + TEST_P(AsyncDataCacheTest, shrinkCache) { constexpr uint64_t kRamBytes = 128UL << 20; constexpr uint64_t kSsdBytes = 512UL << 20;