Count unreleased MultiScan prefetch blocks as wasted (#14926)

Summary:
Pull Request resolved: https://github.com/facebook/rocksdb/pull/14926

This fixes a gap in the MultiScan prefetch waste accounting introduced by https://github.com/facebook/rocksdb/pull/14876. That change added reverse MultiScan support and introduced `unreleased_prefetch_blocks_` to track prefetched blocks still owned by the `ReadSet`, because reverse scans can prefetch a suffix of the collected block list and release blocks in non-prefix order. However, those unreleased prefetched blocks were only released during `MultiScanIndexIterator` destruction; they were not included in `MULTISCAN_PREFETCH_BLOCKS_WASTED`. As a result, scans that stopped before consuming all prefetched blocks underreported wasted prefetch blocks.

This change counts unreleased prefetched blocks as wasted before destructor cleanup, while avoiding counting the currently positioned block as wasted when it has been reached.

Reviewed By: xingbowang

Differential Revision: D111531409

fbshipit-source-id: 5e605ce686037256bf6847ab7ca363c3f7701401

db/db_iterator_test.cc

@@ -5480,6 +5480,65 @@ TEST_P(DBMultiScanIteratorTest, ReversePrefetchAcrossMultipleRanges) {
   ASSERT_GT(tracking_dispatcher->GetReadSets().size(), 0);
 }
 
+TEST_P(DBMultiScanIteratorTest, ReverseUnreleasedPrefetchBlocksCountAsWasted) {
+  auto options = CurrentOptions();
+  options.compression = kNoCompression;
+  options.disable_auto_compactions = true;
+  options.statistics = CreateDBStatistics();
+
+  BlockBasedTableOptions table_options;
+  table_options.flush_block_policy_factory =
+      std::make_shared<FlushBlockEveryKeyPolicyFactory>();
+  table_options.block_cache = NewLRUCache(10 * 1024 * 1024);
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+  DestroyAndReopen(options);
+
+  constexpr int kNumKeys = 40;
+  const std::string value(100, 'v');
+  for (int i = 0; i < kNumKeys; ++i) {
+    ASSERT_OK(Put(Key(i), value));
+  }
+  ASSERT_OK(Flush());
+  ASSERT_OK(db_->CompactRange({}, nullptr, nullptr));
+
+  ReadOptions ro;
+  ro.fill_cache = GetParam();
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+  auto tracking_dispatcher = std::make_shared<TrackingIODispatcher>();
+
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.use_async_io = false;
+  scan_options.reverse = true;
+  scan_options.io_dispatcher = tracking_dispatcher;
+  const std::string start_key = Key(0);
+  const std::string limit_key = Key(kNumKeys);
+  scan_options.insert(start_key, limit_key);
+
+  {
+    std::unique_ptr<MultiScan> iter =
+        dbfull()->NewMultiScan(ro, cfh, scan_options);
+    ASSERT_NE(iter, nullptr);
+
+    try {
+      auto range_it = iter->begin();
+      ASSERT_NE(range_it, iter->end());
+      auto range = *range_it;
+      auto row_it = range.begin();
+      ASSERT_NE(row_it, range.end());
+      ASSERT_EQ((*row_it).first.ToString(), Key(kNumKeys - 1));
+    } catch (MultiScanException& ex) {
+      FAIL() << "Iterator returned status " << ex.what();
+    } catch (std::logic_error& ex) {
+      FAIL() << "Iterator returned logic error " << ex.what();
+    }
+
+    ASSERT_GT(tracking_dispatcher->GetReadSets().size(), 0);
+  }
+
+  ASSERT_GT(
+      options.statistics->getTickerCount(MULTISCAN_PREFETCH_BLOCKS_WASTED), 0);
+}
+
 TEST_P(DBMultiScanIteratorTest, ReverseRequiresLimits) {
   auto options = CurrentOptions();
   options.compression = kNoCompression;
@@ -6206,13 +6265,7 @@ TEST_P(DBMultiScanIteratorTest, WastedBlocksTracking) {
   uint64_t wasted =
       options.statistics->getTickerCount(MULTISCAN_PREFETCH_BLOCKS_WASTED);
 
-  // We expect some wasted blocks due to the gap between ranges
-  // The exact number depends on prefetch behavior, but should be > 0
-  // if blocks between k020-k050 were prefetched
-  std::cout << "Wasted blocks: " << wasted << std::endl;
-
-  // Note: The test verifies the tracking mechanism works.
-  // The actual count depends on prefetch heuristics which may vary.
+  ASSERT_GT(wasted, 0);
 }
 
 class ReadPathRangeTombstoneTest : public DBIteratorBaseTest,

table/block_based/multi_scan_index_iterator.cc

@@ -35,6 +35,7 @@ MultiScanIndexIterator::MultiScanIndexIterator(
       statistics_(statistics) {}
 
 MultiScanIndexIterator::~MultiScanIndexIterator() {
+  CountUnreleasedPrefetchedBlocksAsWasted();
   if (statistics_ && wasted_blocks_count_ > 0) {
     RecordTick(statistics_, MULTISCAN_PREFETCH_BLOCKS_WASTED,
                wasted_blocks_count_);
@@ -50,6 +51,24 @@ MultiScanIndexIterator::~MultiScanIndexIterator() {
   }
 }
 
+void MultiScanIndexIterator::CountUnreleasedPrefetchedBlocksAsWasted() {
+  if (unreleased_prefetch_blocks_ == 0) {
+    return;
+  }
+
+  for (size_t i = first_unreleased_prefetch_idx_;
+       i < one_past_last_unreleased_prefetch_idx_; ++i) {
+    if (!IsPrefetchedBlock(i) ||
+        released_prefetch_blocks_[i - prefetch_start_idx_]) {
+      continue;
+    }
+    if (valid_ && !scan_range_exhausted_ && i == cur_idx_) {
+      continue;
+    }
+    ++wasted_blocks_count_;
+  }
+}
+
 bool MultiScanIndexIterator::IsPrefetchedBlock(size_t block_idx) const {
   return block_idx >= prefetch_start_idx_ && block_idx < prefetch_max_idx_;
 }

table/block_based/multi_scan_index_iterator.h

@@ -111,6 +111,7 @@ class MultiScanIndexIterator : public InternalIteratorBase<IndexValue> {
   bool ReleasePrefetchedBlock(size_t block_idx);
   void ShrinkUnreleasedPrefetchBounds();
   bool IsPrefetchedBlock(size_t block_idx) const;
+  void CountUnreleasedPrefetchedBlocksAsWasted();
 
   // Find the correct scan range and block for an unexpected seek target
   // (target doesn't match expected scan range start).