Prepopulate block cache during compaction (facebook#14445)

Summary:
When RocksDB operates with tiered or remote storage (e.g., Warm Storage, HDFS, S3), reading recently compacted data incurs high-latency remote reads because compaction output files are not present in the block cache. The existing `prepopulate_block_cache = kFlushOnly` avoids this for flush output but leaves compaction output cold until first access.

Add a new `PrepopulateBlockCache::kFlushAndCompaction` enum value that warms all block types (data, index, filter, compression dict) into the block cache during both flush and compaction. Flush-warmed blocks use `LOW` priority (unchanged from kFlushOnly behavior), while compaction-warmed blocks use `BOTTOM` priority — compaction data is less temporally local than freshly flushed data, so it should be the first to be evicted when the cache is full. This gives the remote-read avoidance benefit without risking cache thrashing.

The enum uses `kFlushAndCompaction` rather than separate `kCompactionOnly` + `kFlushAndCompaction` values because there is no practical use case for warming compaction output without also warming flush output. Flush output is by definition the hottest data (just written by the user), so if a workload benefits from warming the colder compaction output, it would always benefit from warming flush output too.

The implementation reuses the existing `InsertBlockInCacheHelper` / `WarmInCache` infrastructure in `BlockBasedTableBuilder`. The only internal change is adding a `warm_cache_priority` field to `Rep` alongside the existing `warm_cache` bool, and plumbing it through to the `WarmInCache` call instead of the previously hardcoded `Cache::Priority::LOW`.

### Key changes
- New `PrepopulateBlockCache::kFlushAndCompaction` enum value in table.h
- `Rep::warm_cache_priority` field in BlockBasedTableBuilder for per-reason priority control
- Serialization support ("kFlushAndCompaction" in string map)
- db_bench support (--prepopulate_block_cache=2)
- Crash test coverage (random choice includes new value)

**NOTE:** Unlike flush output (which is inherently hot — just written by the user), it is hard to distinguish hot from cold blocks in compaction output. Warming all compaction output therefore risks polluting the block cache and evicting genuinely hot entries. The kFlushAndCompaction mode is recommended only for use cases where most or all of the database is expected to reside in cache (e.g., the working set fits in cache). For workloads where only a fraction of the data is hot, kFlushOnly remains the safer choice.

Pull Request resolved: facebook#14445

Test Plan:
- New `WarmCacheWithDataBlocksDuringCompaction` test: verifies data blocks from compaction output are present in the block cache and served without misses
- Extended `DynamicOptions` test: verifies dynamic switching through kDisable -> kFlushAndCompaction -> kFlushOnly -> kDisable via SetOptions
- Existing `WarmCacheWithDataBlocksDuringFlush` and parameterized `WarmCacheWithBlocksDuringFlush` tests continue to pass (kFlushOnly behavior unchanged)
- db_block_cache_test: 81/81 passed
- options_test: 74/74 passed
- table_test: 6910/6910 passed

Reviewed By: xingbowang

Differential Revision: D95997952

Pulled By: mszeszko-meta

fbshipit-source-id: 4ad568264992532053947df298e63e343821ddb5

Author: mszeszko-meta

db/db_block_cache_test.cc

@@ -466,6 +466,90 @@ TEST_F(DBBlockCacheTest, WarmCacheWithDataBlocksDuringFlush) {
             options.statistics->getTickerCount(BLOCK_CACHE_DATA_ADD));
 }
 
+// Cache wrapper that tracks the priority of each Insert call.
+namespace {
+class PriorityTrackingCache : public CacheWrapper {
+ public:
+  explicit PriorityTrackingCache(std::shared_ptr<Cache> target)
+      : CacheWrapper(std::move(target)) {}
+
+  const char* Name() const override { return "PriorityTrackingCache"; }
+
+  Status Insert(const Slice& key, ObjectPtr value,
+                const CacheItemHelper* helper, size_t charge,
+                Handle** handle = nullptr, Priority priority = Priority::LOW,
+                const Slice& compressed_value = Slice(),
+                CompressionType type = kNoCompression) override {
+    insert_priorities_.push_back(priority);
+    return CacheWrapper::Insert(key, value, helper, charge, handle, priority,
+                                compressed_value, type);
+  }
+
+  void ResetPriorities() { insert_priorities_.clear(); }
+
+  bool HasPriority(Priority p) const {
+    return std::find(insert_priorities_.begin(), insert_priorities_.end(), p) !=
+           insert_priorities_.end();
+  }
+
+ private:
+  std::vector<Priority> insert_priorities_;
+};
+}  // namespace
+
+TEST_F(DBBlockCacheTest, WarmCacheWithDataBlocksDuringCompaction) {
+  Options options = CurrentOptions();
+  options.create_if_missing = true;
+  options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
+  options.disable_auto_compactions = true;
+
+  auto tracking_cache =
+      std::make_shared<PriorityTrackingCache>(NewLRUCache(1 << 25, 0, false));
+
+  BlockBasedTableOptions table_options;
+  table_options.block_cache = tracking_cache;
+  table_options.cache_index_and_filter_blocks = false;
+  table_options.prepopulate_block_cache =
+      BlockBasedTableOptions::PrepopulateBlockCache::kFlushAndCompaction;
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+  DestroyAndReopen(options);
+
+  std::string value(kValueSize, 'a');
+
+  // Flush warming: inserts should use LOW priority.
+  tracking_cache->ResetPriorities();
+  ASSERT_OK(Put("key", value));
+  ASSERT_OK(Flush());
+  EXPECT_TRUE(tracking_cache->HasPriority(Cache::Priority::LOW));
+  EXPECT_FALSE(tracking_cache->HasPriority(Cache::Priority::BOTTOM));
+
+  // Write overlapping key to force a real merge compaction (not trivial move).
+  ASSERT_OK(Put("key", value + "2"));
+  ASSERT_OK(Flush());
+
+  auto data_add_before =
+      options.statistics->getTickerCount(BLOCK_CACHE_DATA_ADD);
+
+  // Compaction warming: data block inserts should use BOTTOM priority.
+  // Internal cache bookkeeping (e.g., cache entry stats) may insert at HIGH,
+  // so we check for BOTTOM presence and LOW absence.
+  tracking_cache->ResetPriorities();
+  CompactRangeOptions cro;
+  cro.bottommost_level_compaction = BottommostLevelCompaction::kForceOptimized;
+  ASSERT_OK(db_->CompactRange(cro, /*begin=*/nullptr, /*end=*/nullptr));
+  EXPECT_GT(options.statistics->getTickerCount(BLOCK_CACHE_DATA_ADD),
+            data_add_before);
+  EXPECT_TRUE(tracking_cache->HasPriority(Cache::Priority::BOTTOM));
+  EXPECT_FALSE(tracking_cache->HasPriority(Cache::Priority::LOW));
+
+  // Compaction output is in cache — reads should have zero misses.
+  auto data_miss_before =
+      options.statistics->getTickerCount(BLOCK_CACHE_DATA_MISS);
+  ASSERT_EQ(value + "2", Get("key"));
+  EXPECT_EQ(data_miss_before,
+            options.statistics->getTickerCount(BLOCK_CACHE_DATA_MISS));
+}
+
 // This test cache data, index and filter blocks during flush.
 class DBBlockCacheTest1 : public DBTestBase,
                           public ::testing::WithParamInterface<uint32_t> {
@@ -620,6 +704,36 @@ TEST_F(DBBlockCacheTest, DynamicOptions) {
   ASSERT_EQ(0, st->getAndResetTickerCount(BLOCK_CACHE_DATA_MISS));
   ASSERT_EQ(1, st->getAndResetTickerCount(BLOCK_CACHE_DATA_HIT));
 
+  // Switch to kFlushAndCompaction
+  ++i;
+  ASSERT_OK(dbfull()->SetOptions(
+      {{"block_based_table_factory",
+        "{prepopulate_block_cache=kFlushAndCompaction;}"}}));
+
+  ASSERT_OK(Put(std::to_string(i), value));
+  ASSERT_OK(Flush());
+  // Flush warming still works
+  ASSERT_EQ(1, st->getAndResetTickerCount(BLOCK_CACHE_DATA_ADD));
+
+  ASSERT_EQ(value, Get(std::to_string(i)));
+  ASSERT_EQ(0, st->getAndResetTickerCount(BLOCK_CACHE_DATA_ADD));
+  ASSERT_EQ(0, st->getAndResetTickerCount(BLOCK_CACHE_DATA_MISS));
+  ASSERT_EQ(1, st->getAndResetTickerCount(BLOCK_CACHE_DATA_HIT));
+
+  // Switch back to kDisable
+  ++i;
+  ASSERT_OK(dbfull()->SetOptions(
+      {{"block_based_table_factory", "{prepopulate_block_cache=kDisable;}"}}));
+
+  ASSERT_OK(Put(std::to_string(i), value));
+  ASSERT_OK(Flush());
+  ASSERT_EQ(0, st->getAndResetTickerCount(BLOCK_CACHE_DATA_ADD));
+
+  ASSERT_EQ(value, Get(std::to_string(i)));
+  ASSERT_EQ(1, st->getAndResetTickerCount(BLOCK_CACHE_DATA_ADD));
+  ASSERT_EQ(1, st->getAndResetTickerCount(BLOCK_CACHE_DATA_MISS));
+  ASSERT_EQ(0, st->getAndResetTickerCount(BLOCK_CACHE_DATA_HIT));
+
   ++i;
   // NOT YET SUPPORTED
   // FIXME: find a way to make this fail again (until well supported)

include/rocksdb/table.h

@@ -739,17 +739,34 @@ struct BlockBasedTableOptions {
 
   // If enabled, prepopulate warm/hot blocks (data, uncompressed dict, index and
   // filter blocks) which are already in memory into block cache at the time of
-  // flush. On a flush, the block that is in memory (in memtables) get flushed
-  // to the device. If using Direct IO, additional IO is incurred to read this
-  // data back into memory again, which is avoided by enabling this option. This
+  // flush or compaction.
+  //
+  // On a flush, the data block that is in memory (in memtables) gets flushed to
+  // the device. If using Direct IO, additional IO is incurred to read this data
+  // back into memory again, which is avoided by enabling this option. This
   // further helps if the workload exhibits high temporal locality, where most
   // of the reads go to recently written data. This also helps in case of
   // Distributed FileSystem.
+  //
+  // On a compaction, output SST files are written to disk but not placed in the
+  // block cache by default. With tiered or remote storage (e.g., HDFS, S3),
+  // reading recently compacted data back incurs high latency.
+  // Enabling compaction warming avoids these cold reads. However, unlike flush
+  // output, it is hard to distinguish hot from cold blocks in compaction
+  // output, so warming all of it risks polluting the cache. To mitigate this,
+  // compaction-warmed blocks are inserted at BOTTOM priority (vs LOW for flush)
+  // so they are evicted first under cache pressure. Even so,
+  // kFlushAndCompaction is recommended only when most or all of the database is
+  // expected to reside in cache. For workloads where only a fraction of the
+  // data is hot, kFlushOnly is the safer choice.
   enum class PrepopulateBlockCache : char {
     // Disable prepopulate block cache.
     kDisable,
     // Prepopulate blocks during flush only.
     kFlushOnly,
+    // Prepopulate blocks during flush and compaction. Flush-warmed blocks are
+    // inserted at LOW priority, compaction-warmed blocks at BOTTOM priority.
+    kFlushAndCompaction,
   };
 
   PrepopulateBlockCache prepopulate_block_cache =

table/block_based/block_based_table_builder.cc

@@ -779,6 +779,35 @@ struct BlockBasedTableBuilder::ParallelCompressionRep {
 #endif  // BBTB_PC_WATCHDOG
 };
 
+struct WarmCacheConfig {
+  const bool enabled;
+  const Cache::Priority priority;
+
+  static WarmCacheConfig Compute(
+      BlockBasedTableOptions::PrepopulateBlockCache mode,
+      TableFileCreationReason reason) {
+    bool enabled = false;
+    Cache::Priority priority = Cache::Priority::LOW;
+    switch (mode) {
+      case BlockBasedTableOptions::PrepopulateBlockCache::kFlushOnly:
+        enabled = (reason == TableFileCreationReason::kFlush);
+        break;
+      case BlockBasedTableOptions::PrepopulateBlockCache::kFlushAndCompaction:
+        enabled = (reason == TableFileCreationReason::kFlush ||
+                   reason == TableFileCreationReason::kCompaction);
+        if (reason == TableFileCreationReason::kCompaction) {
+          priority = Cache::Priority::BOTTOM;
+        }
+        break;
+      case BlockBasedTableOptions::PrepopulateBlockCache::kDisable:
+        break;
+      default:
+        assert(false);
+    }
+    return {enabled, priority};
+  }
+};
+
 struct BlockBasedTableBuilder::Rep {
   const ImmutableOptions ioptions;
   // BEGIN from MutableCFOptions
@@ -819,7 +848,6 @@ struct BlockBasedTableBuilder::Rep {
   PartitionedIndexBuilder* p_index_builder_ = nullptr;
 
   std::string last_ikey;  // Internal key or empty (unset)
-  bool warm_cache = false;
   bool uses_explicit_compression_manager = false;
 
   uint64_t sample_for_compression;
@@ -921,6 +949,7 @@ struct BlockBasedTableBuilder::Rep {
 
   std::unique_ptr<ParallelCompressionRep> pc_rep;
   RelaxedAtomic<uint64_t> worker_cpu_micros{0};
+  const WarmCacheConfig warm_cache_config;
   BlockCreateContext create_context;
 
   // The size of the "tail" part of a SST file. "Tail" refers to
@@ -1066,6 +1095,8 @@ struct BlockBasedTableBuilder::Rep {
         flush_block_policy(
             table_options.flush_block_policy_factory->NewFlushBlockPolicy(
                 table_options, data_block)),
+        warm_cache_config(WarmCacheConfig::Compute(
+            table_options.prepopulate_block_cache, reason)),
         create_context(&table_options, &ioptions, ioptions.stats,
                        /*decompressor=*/nullptr,
                        tbo.moptions.block_protection_bytes_per_key,
@@ -1200,19 +1231,6 @@ struct BlockBasedTableBuilder::Rep {
       // the table properties with placeholder info
     }
 
-    switch (table_options.prepopulate_block_cache) {
-      case BlockBasedTableOptions::PrepopulateBlockCache::kFlushOnly:
-        warm_cache = (reason == TableFileCreationReason::kFlush);
-        break;
-      case BlockBasedTableOptions::PrepopulateBlockCache::kDisable:
-        warm_cache = false;
-        break;
-      default:
-        // missing case
-        assert(false);
-        warm_cache = false;
-    }
-
     const auto compress_dict_build_buffer_charged =
         table_options.cache_usage_options.options_overrides
             .at(CacheEntryRole::kCompressionDictionaryBuildingBuffer)
@@ -2141,7 +2159,7 @@ IOStatus BlockBasedTableBuilder::WriteMaybeCompressedBlockImpl(
     }
   }
 
-  if (r->warm_cache) {
+  if (r->warm_cache_config.enabled) {
     io_s = status_to_io_status(
         InsertBlockInCacheHelper(*uncompressed_block_data, handle, block_type));
     if (UNLIKELY(!io_s.ok())) {
@@ -2271,8 +2289,8 @@ Status BlockBasedTableBuilder::InsertBlockInCacheHelper(
     // (de)compression dictionary, which will clone and save a dict-based
     // decompressor from the corresponding non-dict decompressor.
     s = WarmInCache(block_cache, key.AsSlice(), block_contents,
-                    &rep_->create_context, helper, Cache::Priority::LOW,
-                    &charge);
+                    &rep_->create_context, helper,
+                    rep_->warm_cache_config.priority, &charge);
     if (LIKELY(s.ok())) {
       BlockBasedTable::UpdateCacheInsertionMetrics(
           block_type, nullptr /*get_context*/, charge, s.IsOkOverwritten(),

table/block_based/block_based_table_factory.cc

@@ -231,7 +231,9 @@ static std::unordered_map<std::string,
     block_base_table_prepopulate_block_cache_string_map = {
         {"kDisable", BlockBasedTableOptions::PrepopulateBlockCache::kDisable},
         {"kFlushOnly",
-         BlockBasedTableOptions::PrepopulateBlockCache::kFlushOnly}};
+         BlockBasedTableOptions::PrepopulateBlockCache::kFlushOnly},
+        {"kFlushAndCompaction",
+         BlockBasedTableOptions::PrepopulateBlockCache::kFlushAndCompaction}};
 
 static struct BlockBasedTableTypeInfo {
   std::unordered_map<std::string, OptionTypeInfo> info;

tools/db_bench_tool.cc

@@ -745,8 +745,9 @@ DEFINE_bool(separate_key_value_in_data_block,
             "If true, data blocks store keys and values separately.");
 
 DEFINE_int64(prepopulate_block_cache, 0,
-             "Pre-populate hot/warm blocks in block cache. 0 to disable and 1 "
-             "to insert during flush");
+             "Pre-populate hot/warm blocks in block cache. 0 to disable, 1 "
+             "to insert during flush, and 2 to insert during flush and "
+             "compaction");
 
 DEFINE_uint32(uncache_aggressiveness,
               ROCKSDB_NAMESPACE::ColumnFamilyOptions().uncache_aggressiveness,
@@ -4696,6 +4697,10 @@ class Benchmark {
           prepopulate_block_cache =
               BlockBasedTableOptions::PrepopulateBlockCache::kFlushOnly;
           break;
+        case 2:
+          prepopulate_block_cache = BlockBasedTableOptions::
+              PrepopulateBlockCache::kFlushAndCompaction;
+          break;
         default:
           fprintf(stderr, "Unknown prepopulate block cache mode\n");
       }

tools/db_crashtest.py

@@ -327,7 +327,7 @@ def apply_random_seed_per_iteration():
     "user_timestamp_size": 0,
     "secondary_cache_fault_one_in": lambda: random.choice([0, 0, 32]),
     "compressed_secondary_cache_size": lambda: random.choice([8388608, 16777216]),
-    "prepopulate_block_cache": lambda: random.choice([0, 1]),
+    "prepopulate_block_cache": lambda: random.choice([0, 1, 2]),
     "memtable_prefix_bloom_size_ratio": lambda: random.choice([0.001, 0.01, 0.1, 0.5]),
     "memtable_whole_key_filtering": lambda: random.randint(0, 1),
     "detect_filter_construct_corruption": lambda: random.choice([0, 1]),

unreleased_history/new_features/prepopulate_block_cache_on_compaction.md

@@ -0,0 +1 @@
+Added `BlockBasedTableOptions::PrepopulateBlockCache::kFlushAndCompaction` to prepopulate the block cache during both flush and compaction. Compaction-warmed blocks are inserted at `BOTTOM` priority (vs `LOW` for flush) so they are evicted first under cache pressure. Recommended only for use cases where most or all of the database is expected to reside in cache (e.g., tiered or remote storage where the working set fits in cache).