Optimize MultiScan dispatch for sorted blocks (#14783)

Summary:
- Propagate validated, sorted MultiScan range state from `DBIter::Prepare()` through `MultiScanArgs`.
- Mark block-based table IO jobs as already sorted when the public MultiScan ranges have been validated.
- Keep `IODispatcher::SubmitJob()` as the normalization boundary for unsorted callers, while allowing sorted callers to skip the defensive block-handle sort.
- Update private dispatcher coalescing helpers to consume sorted block indices and add debug assertions for that precondition.

## Testing
CI, new unit test

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

Reviewed By: anand1976

Differential Revision: D106301516

Pulled By: xingbowang

fbshipit-source-id: 99b7ffcaecbf27cb79f15feb4af8680ff1e422d9

db/db_iter_test.cc

@@ -216,12 +216,19 @@ class TestIterator : public InternalIterator {
   bool IsKeyPinned() const override { return true; }
   bool IsValuePinned() const override { return true; }
 
+  void Prepare(const MultiScanArgs* /*scan_opts*/) override {
+    ++prepare_call_count_;
+  }
+
+  size_t prepare_call_count() const { return prepare_call_count_; }
+
  private:
   bool initialized_;
   bool valid_;
   size_t sequence_number_;
   size_t iter_;
   size_t steps_ = 0;
+  size_t prepare_call_count_ = 0;
 
   InternalKeyComparator cmp;
   std::vector<std::pair<std::string, std::string>> data_;
@@ -243,6 +250,53 @@ class DBIteratorTest : public testing::Test {
   DBIteratorTest() : env_(Env::Default()) {}
 };
 
+TEST_F(DBIteratorTest, PrepareForwardsValidatedScanRanges) {
+  Options options;
+  ImmutableOptions ioptions = ImmutableOptions(options);
+  MutableCFOptions mutable_cf_options = MutableCFOptions(options);
+  TestIterator* internal_iter = new TestIterator(BytewiseComparator());
+  internal_iter->AddPut("a", "val_a");
+  internal_iter->AddPut("b", "val_b");
+  internal_iter->Finish();
+
+  ReadOptions ro;
+  std::unique_ptr<Iterator> db_iter(DBIter::NewIter(
+      env_, ro, ioptions, mutable_cf_options, BytewiseComparator(),
+      internal_iter, nullptr /* version */, 10 /* sequence */,
+      nullptr /* read_callback */, /*active_mem=*/nullptr));
+
+  MultiScanArgs scan_opts(BytewiseComparator());
+  scan_opts.insert("a", "c");
+
+  db_iter->Prepare(scan_opts);
+  ASSERT_OK(db_iter->status());
+  EXPECT_EQ(internal_iter->prepare_call_count(), 1);
+}
+
+TEST_F(DBIteratorTest, PrepareDoesNotForwardInvalidScanRanges) {
+  Options options;
+  ImmutableOptions ioptions = ImmutableOptions(options);
+  MutableCFOptions mutable_cf_options = MutableCFOptions(options);
+  TestIterator* internal_iter = new TestIterator(BytewiseComparator());
+  internal_iter->AddPut("a", "val_a");
+  internal_iter->AddPut("b", "val_b");
+  internal_iter->Finish();
+
+  ReadOptions ro;
+  std::unique_ptr<Iterator> db_iter(DBIter::NewIter(
+      env_, ro, ioptions, mutable_cf_options, BytewiseComparator(),
+      internal_iter, nullptr /* version */, 10 /* sequence */,
+      nullptr /* read_callback */, /*active_mem=*/nullptr));
+
+  MultiScanArgs scan_opts(BytewiseComparator());
+  scan_opts.insert("b", "d");
+  scan_opts.insert("a", "c");
+
+  db_iter->Prepare(scan_opts);
+  EXPECT_TRUE(db_iter->status().IsInvalidArgument());
+  EXPECT_EQ(internal_iter->prepare_call_count(), 0);
+}
+
 TEST_F(DBIteratorTest, DBIteratorPrevNext) {
   Options options;
   ImmutableOptions ioptions = ImmutableOptions(options);

db/db_iterator_test.cc

@@ -4395,6 +4395,71 @@ TEST_P(DBMultiScanIteratorTest, RangeAcrossFiles) {
   iter.reset();
 }
 
+TEST_P(DBMultiScanIteratorTest, SortedRangesSkipIODispatcherSort) {
+  auto options = CurrentOptions();
+  options.target_file_size_base = 100 << 10;
+  options.compaction_style = kCompactionStyleUniversal;
+  options.num_levels = 50;
+  options.compression = kNoCompression;
+  DestroyAndReopen(options);
+
+  auto key = [](int i) {
+    std::stringstream ss;
+    ss << "k" << std::setw(5) << std::setfill('0') << i;
+    return ss.str();
+  };
+
+  auto rnd = Random::GetTLSInstance();
+  for (int i = 0; i < 100; ++i) {
+    ASSERT_OK(Put(key(i), rnd->RandomString(2 << 10)));
+  }
+  ASSERT_OK(Flush());
+
+  ASSERT_OK(db_->CompactRange({}, nullptr, nullptr));
+  ASSERT_EQ(2, NumTableFilesAtLevel(49));
+
+  int sort_count = 0;
+  SyncPoint::GetInstance()->SetCallBack(
+      "IODispatcherImpl::SubmitJob:SortBlockHandles",
+      [&](void* /*arg*/) { ++sort_count; });
+  SyncPoint::GetInstance()->EnableProcessing();
+
+  auto tracking_dispatcher = std::make_shared<TrackingIODispatcher>();
+  std::vector<std::string> key_ranges({key(10), key(90)});
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.io_dispatcher = tracking_dispatcher;
+  scan_options.use_async_io = false;
+  scan_options.insert(key_ranges[0], key_ranges[1]);
+
+  ReadOptions ro;
+  ro.fill_cache = GetParam();
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+  std::unique_ptr<MultiScan> iter =
+      dbfull()->NewMultiScan(ro, cfh, scan_options);
+
+  try {
+    int i = 10;
+    for (auto range : *iter) {
+      for (auto it : range) {
+        ASSERT_EQ(it.first.ToString(), key(i));
+        ++i;
+      }
+    }
+    ASSERT_EQ(i, 90);
+  } catch (MultiScanException& ex) {
+    ASSERT_NOK(ex.status());
+    std::cerr << "Iterator returned status " << ex.what();
+    abort();
+  } catch (std::logic_error& ex) {
+    std::cerr << "Iterator returned logic error " << ex.what();
+    abort();
+  }
+
+  ASSERT_GT(tracking_dispatcher->GetReadSets().size(), 0);
+  EXPECT_EQ(sort_count, 0);
+  iter.reset();
+}
+
 TEST_P(DBMultiScanIteratorTest, FailureTest) {
   auto options = CurrentOptions();
   options.compression = kNoCompression;

include/rocksdb/io_dispatcher.h

@@ -152,6 +152,11 @@ class BlockBasedTable;
 struct JobOptions {
   uint64_t io_coalesce_threshold = 16 * 1024;
   ReadOptions read_options;
+
+  // True when IOJob::block_handles is already sorted by file offset. Callers
+  // that can guarantee sorted order may set this to let IODispatcher skip
+  // defensive sorting before cache lookup and coalescing.
+  bool block_handles_are_sorted = false;
 };
 
 class IOJob {

table/block_based/block_based_table_iterator.cc

@@ -1077,6 +1077,9 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
       multiscan_opts->io_coalesce_threshold;
   job->job_options.read_options = read_options_;
   job->job_options.read_options.async_io = multiscan_opts->use_async_io;
+  // CollectBlockHandles walks sorted scan ranges through the table index, whose
+  // key order matches data block offset order.
+  job->job_options.block_handles_are_sorted = true;
 
   std::shared_ptr<ReadSet> read_set;
   // IODispatcher should be provided by DBIter::Prepare() to enable sharing

util/io_dispatcher_imp.cc

@@ -44,6 +44,25 @@ struct IODispatcherImplData {
   virtual void ReleaseMemory(size_t bytes) = 0;
 };
 
+#ifndef NDEBUG
+static bool BlockIndicesAreSortedByOffset(
+    const std::vector<BlockHandle>& block_handles,
+    const std::vector<size_t>& block_indices) {
+  uint64_t prev_offset = 0;
+  for (size_t i = 0; i < block_indices.size(); ++i) {
+    if (block_indices[i] >= block_handles.size()) {
+      return false;
+    }
+    const uint64_t current_offset = block_handles[block_indices[i]].offset();
+    if (i > 0 && current_offset < prev_offset) {
+      return false;
+    }
+    prev_offset = current_offset;
+  }
+  return true;
+}
+#endif  // NDEBUG
+
 // Helper function to create and pin a block from a buffer
 // Used by both ReadSet::PollAndProcessAsyncIO and IODispatcherImpl::Impl
 static Status CreateAndPinBlockFromBuffer(
@@ -479,7 +498,8 @@ struct IODispatcherImpl::Impl : public IODispatcherImplData,
                         const std::vector<size_t>& block_indices);
 
   // Pre-coalesce blocks into groups, respecting max_group_bytes size limit.
-  // Returns groups ordered by first block index (earlier blocks first).
+  // block_indices must be sorted by block offset.
+  // Returns groups ordered by block offset.
   std::vector<CoalescedPrefetchGroup> PreCoalesceBlocks(
       const std::shared_ptr<IOJob>& job, const std::shared_ptr<ReadSet>& rs,
       const std::vector<size_t>& block_indices, size_t max_group_bytes);
@@ -697,16 +717,22 @@ Status IODispatcherImpl::Impl::SubmitJob(const std::shared_ptr<IOJob>& job,
   for (size_t i = 0; i < job->block_handles.size(); ++i) {
     rs->sorted_block_indices_[i] = i;
   }
-  std::sort(rs->sorted_block_indices_.begin(), rs->sorted_block_indices_.end(),
-            [&job](size_t a, size_t b) {
+  if (!job->job_options.block_handles_are_sorted) {
+    TEST_SYNC_POINT("IODispatcherImpl::SubmitJob:SortBlockHandles");
+    std::sort(rs->sorted_block_indices_.begin(),
+              rs->sorted_block_indices_.end(), [&job](size_t a, size_t b) {
                 return job->block_handles[a].offset() <
                        job->block_handles[b].offset();
               });
+  }
 
-  // Step 1: Check cache and pin cached blocks
+  // Step 1: Check cache and pin cached blocks. Iterate in offset order so all
+  // downstream private helpers can assume block index vectors are already
+  // sorted.
   std::vector<size_t> block_indices_to_read;
+  block_indices_to_read.reserve(job->block_handles.size());
 
-  for (size_t i = 0; i < job->block_handles.size(); ++i) {
+  for (size_t i : rs->sorted_block_indices_) {
     const auto& data_block_handle = job->block_handles[i];
 
     // Lookup and pin block in cache
@@ -816,17 +842,11 @@ void IODispatcherImpl::Impl::PrepareIORequests(
     const std::vector<BlockHandle>& block_handles,
     std::vector<FSReadRequest>* read_reqs,
     std::vector<std::vector<size_t>>* coalesced_block_indices) {
-  // This is necessary because block handles may not be in sorted order
-  std::vector<size_t> sorted_block_indices = block_indices_to_read;
-  std::sort(sorted_block_indices.begin(), sorted_block_indices.end(),
-            [&block_handles](size_t a, size_t b) {
-              return block_handles[a].offset() < block_handles[b].offset();
-            });
-
+  assert(BlockIndicesAreSortedByOffset(block_handles, block_indices_to_read));
   assert(coalesced_block_indices->empty());
   coalesced_block_indices->resize(1);
 
-  for (const auto& block_idx : sorted_block_indices) {
+  for (const auto& block_idx : block_indices_to_read) {
     if (!coalesced_block_indices->back().empty()) {
       // Check if we can coalesce with previous block
       const auto& last_block_handle =
@@ -883,17 +903,12 @@ std::vector<CoalescedPrefetchGroup> IODispatcherImpl::Impl::PreCoalesceBlocks(
   const auto& block_handles = job->block_handles;
   const uint64_t coalesce_threshold = job->job_options.io_coalesce_threshold;
 
-  // Sort block indices by offset for coalescing
-  std::vector<size_t> sorted_indices = block_indices;
-  std::sort(sorted_indices.begin(), sorted_indices.end(),
-            [&block_handles](size_t a, size_t b) {
-              return block_handles[a].offset() < block_handles[b].offset();
-            });
+  assert(BlockIndicesAreSortedByOffset(block_handles, block_indices));
 
   // Build coalesced groups respecting max_group_bytes
   groups.emplace_back();
 
-  for (size_t idx : sorted_indices) {
+  for (size_t idx : block_indices) {
     size_t block_size = rs->block_sizes_[idx];
 
     // Skip blocks that are individually larger than the memory budget

util/io_dispatcher_test.cc

@@ -10,6 +10,7 @@
 
 #include "rocksdb/io_dispatcher.h"
 
+#include <algorithm>
 #include <memory>
 #include <mutex>
 #include <thread>
@@ -836,6 +837,85 @@ TEST_F(IODispatcherTest, VerifyCoalescing) {
   }
 }
 
+TEST_F(IODispatcherTest, SortedBlockHandlesSkipDispatcherSorts) {
+  struct SortCounters {
+    int submit_job = 0;
+  };
+
+  auto install_callbacks = [](SortCounters* counters) {
+    SyncPoint::GetInstance()->SetCallBack(
+        "IODispatcherImpl::SubmitJob:SortBlockHandles",
+        [counters](void*) { ++counters->submit_job; });
+    SyncPoint::GetInstance()->EnableProcessing();
+  };
+
+  auto submit_job = [](IODispatcher* dispatcher, BlockBasedTable* table,
+                       const std::vector<BlockHandle>& block_handles,
+                       bool sorted) {
+    auto job = std::make_shared<IOJob>();
+    job->block_handles = block_handles;
+    job->table = table;
+    job->job_options.block_handles_are_sorted = sorted;
+    job->job_options.read_options.async_io = false;
+    job->job_options.io_coalesce_threshold = 1024 * 1024;
+
+    std::shared_ptr<ReadSet> read_set;
+    ASSERT_OK(dispatcher->SubmitJob(job, &read_set));
+    ASSERT_NE(read_set, nullptr);
+    for (size_t i = 0; i < block_handles.size(); ++i) {
+      CachableEntry<Block> block;
+      ASSERT_OK(read_set->ReadIndex(i, &block));
+      ASSERT_NE(block.GetValue(), nullptr);
+      ASSERT_GT(block.GetValue()->size(), 0);
+
+      InternalKeyComparator internal_comparator(BytewiseComparator());
+      std::unique_ptr<DataBlockIter> iter(block.GetValue()->NewDataIterator(
+          internal_comparator.user_comparator(), kDisableGlobalSequenceNumber));
+      iter->SeekToFirst();
+      ASSERT_TRUE(iter->Valid());
+
+      ParsedInternalKey parsed_key;
+      ASSERT_OK(
+          ParseInternalKey(iter->key(), &parsed_key, true /* log_err_key */));
+      ASSERT_TRUE(parsed_key.user_key.starts_with("key"));
+      ASSERT_OK(iter->status());
+    }
+  };
+
+  std::unique_ptr<BlockBasedTable> sorted_table;
+  std::vector<BlockHandle> sorted_handles;
+  ASSERT_OK(CreateAndOpenSST(10, &sorted_table, &sorted_handles));
+  ASSERT_GE(sorted_handles.size(), 5);
+  sorted_handles.resize(5);
+
+  SortCounters sorted_counters;
+  install_callbacks(&sorted_counters);
+  std::unique_ptr<IODispatcher> sorted_dispatcher(NewIODispatcher());
+  submit_job(sorted_dispatcher.get(), sorted_table.get(), sorted_handles,
+             true /* sorted */);
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->ClearAllCallBacks();
+
+  EXPECT_EQ(sorted_counters.submit_job, 0);
+
+  std::unique_ptr<BlockBasedTable> unsorted_table;
+  std::vector<BlockHandle> unsorted_handles;
+  ASSERT_OK(CreateAndOpenSST(10, &unsorted_table, &unsorted_handles));
+  ASSERT_GE(unsorted_handles.size(), 5);
+  unsorted_handles.resize(5);
+  std::reverse(unsorted_handles.begin(), unsorted_handles.end());
+
+  SortCounters unsorted_counters;
+  install_callbacks(&unsorted_counters);
+  std::unique_ptr<IODispatcher> unsorted_dispatcher(NewIODispatcher());
+  submit_job(unsorted_dispatcher.get(), unsorted_table.get(), unsorted_handles,
+             false /* sorted */);
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->ClearAllCallBacks();
+
+  EXPECT_EQ(unsorted_counters.submit_job, 1);
+}
+
 // Test that verifies the read request offsets and lengths match the
 // expected block handles.
 TEST_F(IODispatcherTest, VerifyReadRequestDetails) {