feat(cache): Add CacheEntrySizePolicy to reuse shorter cache entries as prefixes

velox/common/caching/AsyncDataCache.cpp

@@ -251,7 +251,8 @@ std::unique_ptr<AsyncDataCacheEntry> CacheShard::getFreeEntryLocked() {
 std::optional<CachePin> CacheShard::lookupLocked(
     RawFileCacheKey key,
     uint64_t size,
-    folly::SemiFuture<bool>* wait) {
+    folly::SemiFuture<bool>* wait,
+    CacheEntrySizePolicy sizePolicy) {
   ++eventCounter_;
   auto it = entryMap_.find(key);
   if (it == entryMap_.end()) {
@@ -266,7 +267,8 @@ std::optional<CachePin> 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<bool>* wait) {
+    folly::SemiFuture<bool>* wait,
+    CacheEntrySizePolicy sizePolicy) {
   AsyncDataCacheEntry* entryToInit = nullptr;
   {
     std::lock_guard<std::mutex> 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<bool>* wait) {
+    folly::SemiFuture<bool>* wait,
+    CacheEntrySizePolicy sizePolicy) {
   const int shard = std::hash<RawFileCacheKey>()(key) & shardMask_;
-  return shards_[shard]->findOrCreate(key, size, contiguous, wait);
+  return shards_[shard]->findOrCreate(key, size, contiguous, wait, sizePolicy);
 }
 
 std::optional<CachePin> AsyncDataCache::find(

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<bool>* readyFuture = nullptr);
+      folly::SemiFuture<bool>* 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<CachePin> lookupLocked(
       RawFileCacheKey key,
       uint64_t size,
-      folly::SemiFuture<bool>* waitFuture);
+      folly::SemiFuture<bool>* waitFuture,
+      CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast);
 
   void tryAddFreeEntry(std::unique_ptr<AsyncDataCacheEntry>&& 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<bool>* waitFuture = nullptr);
+      folly::SemiFuture<bool>* 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<RawFileCacheKey>& 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;
       }

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<bool> 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;