Restore MultiScan prefetch stats dropped in the IODispatcher refactor (#14991)

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

D91705195 ("Replace Prefetch Logic in BlockBasedTableIterator with IODispatcher") moved MultiScan block prefetching out of BlockBasedTableIterator and into the IODispatcher. The IO-level RecordTick/RecordInHistogram calls lived in the deleted prefetch code, and the IODispatcher was not given equivalent instrumentation, so these tickers were still defined but never populated (they published as absent):

- MULTISCAN_BLOCKS_PREFETCHED
- MULTISCAN_BLOCKS_FROM_CACHE
- MULTISCAN_PREFETCH_BYTES
- MULTISCAN_IO_REQUESTS
- MULTISCAN_IO_COALESCED_NONADJACENT
- MULTISCAN_BLOCKS_PER_PREPARE (histogram)

The PREPARE_*, SEEK_ERRORS, and PREFETCH_BLOCKS_WASTED tickers were unaffected (recorded on the prepare/teardown path, which was not replaced).

Fix:
- ReadSet now exposes generic, caller-agnostic prefetch counters (GetNumBlocksPrefetched, GetNumIORequests, GetPrefetchBytes, GetNumCoalescedNonadjacent), incremented once per block in DispatchPrefetch -- the single choke point for both immediate and memory-deferred dispatch. Counting at dispatch time (not completion) keeps async and deferred prefetches accounted for without double counting the async->sync fallback.
- PrepareIORequests now outputs prefetch_bytes (sum of request len) and num_coalesced_nonadjacent (within-threshold gap count) while building coalesced requests, per review feedback to avoid the extra pass over read_reqs/coalesced_block_indices at former line 836. DispatchPrefetch now gets stats directly from PrepareIORequests.
- The MULTISCAN_* ticker names stay at the caller: MultiScanIndexIterator's destructor records the aggregate counters at end-of-scan (alongside the existing PREFETCH_BLOCKS_WASTED), so async/deferred reads that complete during iteration are included. MULTISCAN_BLOCKS_PER_PREPARE is recorded in BlockBasedTableIterator::Prepare. This keeps the IODispatcher a generic mechanism rather than leaking a caller's metric names into it.
- Also plumb the same metrics into PerfContext for per-operation visibility.

Reviewed By: xingbowang

Differential Revision: D113356436

fbshipit-source-id: 660e9755d6b1b9aecccd112df94eb6838061b298

db/db_iterator_test.cc

@@ -6268,6 +6268,95 @@ TEST_P(DBMultiScanIteratorTest, WastedBlocksTracking) {
   ASSERT_GT(wasted, 0);
 }
 
+TEST_P(DBMultiScanIteratorTest, PrefetchStatsRecorded) {
+  // Regression test: the MultiScan prefetch stats must still be published after
+  // the prefetch path moved into the IODispatcher. A cold scan populates the
+  // prefetch counters; a warm re-scan registers cache hits.
+  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);  // 10MB cache
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+  DestroyAndReopen(options);
+
+  const int kNumKeys = 50;
+  std::string value(100, 'v');
+  for (int i = 0; i < kNumKeys; ++i) {
+    std::stringstream ss;
+    ss << "k" << std::setw(3) << std::setfill('0') << i;
+    ASSERT_OK(Put(ss.str(), value));
+  }
+  ASSERT_OK(Flush());
+  ASSERT_OK(db_->CompactRange({}, nullptr, nullptr));
+
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+
+  auto run_scan = [&]() {
+    MultiScanArgs scan_options(BytewiseComparator());
+    scan_options.use_async_io = false;  // deterministic sync path
+    scan_options.insert("k000", "k049");
+    ReadOptions ro;
+    ro.fill_cache = true;
+    std::unique_ptr<MultiScan> iter =
+        dbfull()->NewMultiScan(ro, cfh, scan_options);
+    ASSERT_NE(iter, nullptr);
+    if (iter == nullptr) {
+      return;
+    }
+    int count = 0;
+    try {
+      for (auto range : *iter) {
+        for (auto it : range) {
+          it.first.ToString();
+          count++;
+        }
+      }
+    } catch (MultiScanException& ex) {
+      FAIL() << "Scan failed: " << ex.what();
+    }
+    ASSERT_GT(count, 0);
+  };
+
+  // Cold scan: force blocks to be read from disk so the prefetch path runs.
+  table_options.block_cache->EraseUnRefEntries();
+  SetPerfLevel(kEnableCount);
+  PerfContext* perf = get_perf_context();
+  ASSERT_NE(perf, nullptr);
+  if (perf == nullptr) {
+    return;
+  }
+  perf->Reset();
+  run_scan();  // MultiScanIndexIterator destroyed here -> stats recorded
+
+  ASSERT_GT(TestGetTickerCount(options, MULTISCAN_PREPARE_CALLS), 0);
+  ASSERT_GT(TestGetTickerCount(options, MULTISCAN_BLOCKS_PREFETCHED), 0);
+  ASSERT_GT(TestGetTickerCount(options, MULTISCAN_PREFETCH_BYTES), 0);
+  ASSERT_GT(TestGetTickerCount(options, MULTISCAN_IO_REQUESTS), 0);
+  HistogramData blocks_per_prepare;
+  options.statistics->histogramData(MULTISCAN_BLOCKS_PER_PREPARE,
+                                    &blocks_per_prepare);
+  ASSERT_GT(blocks_per_prepare.count, 0);
+
+  // PerfContext mirrors the same metrics scoped to this operation/thread.
+  ASSERT_GT(perf->multiscan_prepare_count, 0);
+  ASSERT_GT(perf->multiscan_blocks_prefetched, 0);
+  ASSERT_GT(perf->multiscan_prefetch_bytes, 0);
+  ASSERT_GT(perf->multiscan_io_requests, 0);
+  SetPerfLevel(kDisable);
+
+  // Warm scan: the just-cached blocks now count as cache hits.
+  uint64_t cache_hits_before =
+      TestGetTickerCount(options, MULTISCAN_BLOCKS_FROM_CACHE);
+  run_scan();
+  ASSERT_GT(TestGetTickerCount(options, MULTISCAN_BLOCKS_FROM_CACHE),
+            cache_hits_before);
+}
+
 class ReadPathRangeTombstoneTest : public DBIteratorBaseTest,
                                    public ::testing::WithParamInterface<bool> {
  protected:

include/rocksdb/io_dispatcher.h

@@ -221,6 +221,15 @@ class ReadSet {
   uint64_t GetNumSyncReads() const { return num_sync_reads_; }
   uint64_t GetNumAsyncReads() const { return num_async_reads_; }
   uint64_t GetNumCacheHits() const { return num_cache_hits_; }
+  // Blocks for which a prefetch IO was dispatched (sync + async), the number of
+  // (coalesced) IO requests issued for them, the total bytes those requests
+  // span, and how many blocks were coalesced across a non-adjacent gap.
+  uint64_t GetNumBlocksPrefetched() const { return num_blocks_prefetched_; }
+  uint64_t GetNumIORequests() const { return num_io_requests_; }
+  uint64_t GetPrefetchBytes() const { return prefetch_bytes_; }
+  uint64_t GetNumCoalescedNonadjacent() const {
+    return num_coalesced_nonadjacent_;
+  }
 
  private:
   friend class IODispatcherImpl;
@@ -255,6 +264,12 @@ class ReadSet {
   std::atomic<uint64_t> num_sync_reads_ = 0;
   std::atomic<uint64_t> num_async_reads_ = 0;
   std::atomic<uint64_t> num_cache_hits_ = 0;
+  // Prefetch IO accounting, incremented as prefetch IO is dispatched (not on
+  // completion) so async and memory-deferred prefetches are all captured.
+  std::atomic<uint64_t> num_blocks_prefetched_ = 0;
+  std::atomic<uint64_t> num_io_requests_ = 0;
+  std::atomic<uint64_t> prefetch_bytes_ = 0;
+  std::atomic<uint64_t> num_coalesced_nonadjacent_ = 0;
 
   // Poll and process a specific async IO request
   Status PollAndProcessAsyncIO(

include/rocksdb/perf_context.h

@@ -300,6 +300,24 @@ struct PerfContextBase {
   uint64_t filter_block_read_byte;
   uint64_t compression_dict_block_read_byte;
   uint64_t metadata_block_read_byte;
+
+  // MultiScan (scan Prepare) prefetch metrics. Populated by
+  // BlockBasedTableIterator::Prepare and its IODispatcher prefetch path. These
+  // mirror the rocksdb.multiscan.* tickers but are scoped to the current
+  // thread/operation. Average blocks per prepare is derivable as
+  // (multiscan_blocks_prefetched + multiscan_blocks_from_cache) /
+  // multiscan_prepare_count.
+  uint64_t multiscan_prepare_count;
+  // Blocks for which a prefetch IO was dispatched (read from disk).
+  uint64_t multiscan_blocks_prefetched;
+  // Blocks that were already in the block cache at Prepare time.
+  uint64_t multiscan_blocks_from_cache;
+  // Total bytes spanned by the (coalesced) prefetch IO requests.
+  uint64_t multiscan_prefetch_bytes;
+  // Number of (coalesced) IO requests issued for the prefetch.
+  uint64_t multiscan_io_requests;
+  // Blocks coalesced into an IO request across a non-adjacent gap.
+  uint64_t multiscan_io_coalesced_nonadjacent;
 };
 
 struct PerfContext : public PerfContextBase {

monitoring/perf_context.cc

@@ -162,7 +162,13 @@ struct PerfContextByLevelInt {
   defCmd(index_block_read_byte)                    \
   defCmd(filter_block_read_byte)                   \
   defCmd(compression_dict_block_read_byte)         \
-  defCmd(metadata_block_read_byte)
+  defCmd(metadata_block_read_byte)                 \
+  defCmd(multiscan_prepare_count)                  \
+  defCmd(multiscan_blocks_prefetched)             \
+  defCmd(multiscan_blocks_from_cache)             \
+  defCmd(multiscan_prefetch_bytes)                \
+  defCmd(multiscan_io_requests)                   \
+  defCmd(multiscan_io_coalesced_nonadjacent)
 // clang-format on
 
 struct PerfContextInt {

table/block_based/block_based_table_iterator.cc

@@ -8,6 +8,8 @@
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
 #include "table/block_based/block_based_table_iterator.h"
 
+#include "monitoring/perf_context_imp.h"
+
 namespace ROCKSDB_NAMESPACE {
 
 void BlockBasedTableIterator::SeekToFirst() { SeekImpl(nullptr, false); }
@@ -1029,6 +1031,7 @@ void BlockBasedTableIterator::BlockCacheLookupForReadAheadSize(
 void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
   assert(!multi_scan_read_set_);
   RecordTick(table_->GetStatistics(), MULTISCAN_PREPARE_CALLS);
+  PERF_COUNTER_ADD(multiscan_prepare_count, 1);
   StopWatch sw(table_->get_rep()->ioptions.clock, table_->GetStatistics(),
                MULTISCAN_PREPARE_MICROS);
 
@@ -1118,6 +1121,10 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
     return;
   }
 
+  // Number of data blocks submitted to the dispatcher for this Prepare.
+  RecordInHistogram(table_->GetStatistics(), MULTISCAN_BLOCKS_PER_PREPARE,
+                    job->block_handles.size());
+
   // Successful Prepare. Create MultiScanIndexIterator and swap it in as
   // the index iterator. The original index_iter_ is saved for restoration
   // on backward operations.

table/block_based/multi_scan_index_iterator.cc

@@ -7,6 +7,7 @@
 
 #include <algorithm>
 
+#include "monitoring/perf_context_imp.h"
 #include "monitoring/statistics_impl.h"
 #include "rocksdb/options.h"
 
@@ -40,6 +41,31 @@ MultiScanIndexIterator::~MultiScanIndexIterator() {
     RecordTick(statistics_, MULTISCAN_PREFETCH_BLOCKS_WASTED,
                wasted_blocks_count_);
   }
+  // Report the prefetch IO the ReadSet performed for this Prepare. Recorded at
+  // teardown (not in Prepare) so async and memory-deferred prefetches, which
+  // complete during iteration, are included in the counts.
+  if (read_set_) {
+    const uint64_t blocks_from_cache = read_set_->GetNumCacheHits();
+    const uint64_t blocks_prefetched = read_set_->GetNumBlocksPrefetched();
+    const uint64_t prefetch_bytes = read_set_->GetPrefetchBytes();
+    const uint64_t io_requests = read_set_->GetNumIORequests();
+    const uint64_t io_coalesced_nonadjacent =
+        read_set_->GetNumCoalescedNonadjacent();
+    if (statistics_) {
+      RecordTick(statistics_, MULTISCAN_BLOCKS_FROM_CACHE, blocks_from_cache);
+      RecordTick(statistics_, MULTISCAN_BLOCKS_PREFETCHED, blocks_prefetched);
+      RecordTick(statistics_, MULTISCAN_PREFETCH_BYTES, prefetch_bytes);
+      RecordTick(statistics_, MULTISCAN_IO_REQUESTS, io_requests);
+      RecordTick(statistics_, MULTISCAN_IO_COALESCED_NONADJACENT,
+                 io_coalesced_nonadjacent);
+    }
+    PERF_COUNTER_ADD(multiscan_blocks_from_cache, blocks_from_cache);
+    PERF_COUNTER_ADD(multiscan_blocks_prefetched, blocks_prefetched);
+    PERF_COUNTER_ADD(multiscan_prefetch_bytes, prefetch_bytes);
+    PERF_COUNTER_ADD(multiscan_io_requests, io_requests);
+    PERF_COUNTER_ADD(multiscan_io_coalesced_nonadjacent,
+                     io_coalesced_nonadjacent);
+  }
   // Release any remaining pinned blocks
   if (read_set_) {
     for (size_t i = first_unreleased_prefetch_idx_;

util/io_dispatcher_imp.cc

@@ -611,7 +611,9 @@ struct IODispatcherImpl::Impl : public IODispatcherImplData,
       const std::vector<size_t>& block_indices_to_read,
       const std::vector<BlockHandle>& block_handles,
       std::vector<FSReadRequest>* read_reqs,
-      std::vector<std::vector<size_t>>* coalesced_block_indices);
+      std::vector<std::vector<size_t>>* coalesced_block_indices,
+      uint64_t* out_prefetch_bytes = nullptr,
+      uint64_t* out_num_coalesced_nonadjacent = nullptr);
 
   // Surface actual async IO errors to caller, but allow fallback for
   // unsupported cases. Returns block indices that need sync fallback.
@@ -806,8 +808,19 @@ void IODispatcherImpl::Impl::DispatchPrefetch(
   // Prepare and execute IO for the given blocks
   std::vector<FSReadRequest> read_reqs;
   std::vector<std::vector<size_t>> coalesced_block_indices;
+  uint64_t dispatched_bytes = 0;
+  uint64_t nonadjacent = 0;
   PrepareIORequests(job, block_indices, job->block_handles, &read_reqs,
-                    &coalesced_block_indices);
+                    &coalesced_block_indices, &dispatched_bytes, &nonadjacent);
+
+  // Account for the prefetch IO issued by this dispatch. Counted here (once per
+  // dispatched block, before the async/sync split) rather than at completion so
+  // async and memory-deferred prefetches are all captured, and so blocks that
+  // fall back from async to sync are not double counted.
+  read_set->num_blocks_prefetched_ += block_indices.size();
+  read_set->num_io_requests_ += read_reqs.size();
+  read_set->prefetch_bytes_ += dispatched_bytes;
+  read_set->num_coalesced_nonadjacent_ += nonadjacent;
 
   if (job->job_options.read_options.async_io) {
     Status async_status;
@@ -987,11 +1000,13 @@ void IODispatcherImpl::Impl::PrepareIORequests(
     const std::vector<size_t>& block_indices_to_read,
     const std::vector<BlockHandle>& block_handles,
     std::vector<FSReadRequest>* read_reqs,
-    std::vector<std::vector<size_t>>* coalesced_block_indices) {
+    std::vector<std::vector<size_t>>* coalesced_block_indices,
+    uint64_t* out_prefetch_bytes, uint64_t* out_num_coalesced_nonadjacent) {
   assert(BlockIndicesAreSortedByOffset(block_handles, block_indices_to_read));
   assert(coalesced_block_indices->empty());
   coalesced_block_indices->resize(1);
 
+  uint64_t nonadjacent = 0;
   for (const auto& block_idx : block_indices_to_read) {
     if (!coalesced_block_indices->back().empty()) {
       // Check if we can coalesce with previous block
@@ -1006,6 +1021,9 @@ void IODispatcherImpl::Impl::PrepareIORequests(
           last_block_end + job->job_options.io_coalesce_threshold) {
         // Gap too large - start new IO request
         coalesced_block_indices->emplace_back();
+      } else if (current_start > last_block_end) {
+        // Within coalesce threshold but not adjacent
+        ++nonadjacent;
       }
     }
     coalesced_block_indices->back().emplace_back(block_idx);
@@ -1015,6 +1033,7 @@ void IODispatcherImpl::Impl::PrepareIORequests(
   assert(read_reqs->empty());
   read_reqs->reserve(coalesced_block_indices->size());
 
+  uint64_t total_bytes = 0;
   for (const auto& block_indices : *coalesced_block_indices) {
     assert(!block_indices.empty());
 
@@ -1034,6 +1053,14 @@ void IODispatcherImpl::Impl::PrepareIORequests(
     read_reqs->back().offset = start_offset;
     read_reqs->back().len = end_offset - start_offset;
     read_reqs->back().scratch = nullptr;
+    total_bytes += read_reqs->back().len;
+  }
+
+  if (out_prefetch_bytes) {
+    *out_prefetch_bytes = total_bytes;
+  }
+  if (out_num_coalesced_nonadjacent) {
+    *out_num_coalesced_nonadjacent = nonadjacent;
   }
 }