Support Prepare() in BlockBasedTableIterator For MultiScan (#13778)

Summary:
initial support for Prepare() to optimize the performance of MultiScan when using block-based tables. In Prepare(), we do the following:
1. Load all data blocks that will be read in multiscan to block cache
2. Pin the data blocks during the scan
3. if I/O is needed, coalesce I/Os when they are adjacent.

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

Test Plan:
Added a new unit test.

Benchmark:
1. Set up the DB, I use FIFO here so that files will be in L0 and iterator will use BlockBasedTableIterator directly instead of LevelIterator, where Prepare() call is not implemented yet.
```
./db_bench --benchmarks="fillseq,compact" --disable_wal=1 --threads=1 --num_levels=1 --compaction_style=2 --fifo_compaction_max_table_files_size_mb=1000 --write_buffer_size=268435456
```

2. Multi-scan: based on https://github.com/facebook/rocksdb/issues/13765
```
./db_bench --db="/tmp/rocksdbtest-543376/dbbench" --use_existing_db=1 --benchmarks=multiscan --disable_auto_compactions=1 --seek_nexts=100 --threads=32 --duration=10 --statistics=1

multiscan_stride = 100
multiscan_size = 10
seek_nexts = 100

Main:
multiscan    :     449.386 micros/op 70562 ops/sec 10.359 seconds 730968 operations; (multscans:22999)
multiscan    :     453.606 micros/op 69433 ops/sec 10.369 seconds 719968 operations; (multscans:22999)
rocksdb.non.last.level.read.bytes COUNT : 47763519421
rocksdb.non.last.level.read.count COUNT : 21573878
Branch:
multiscan    :     332.670 micros/op 94698 ops/sec 10.285 seconds 973968 operations; (multscans:29999)
rocksdb.non.last.level.read.bytes COUNT : 111791308336
rocksdb.non.last.level.read.count COUNT : 1062942

With direct-IO:
./db_bench --db="/tmp/rocksdbtest-543376/dbbench" --use_existing_db=1 --benchmarks=multiscan --disable_auto_compactions=1 --seek_nexts=100 --threads=32 --duration=10 --statistics=1 --use_direct_reads=1

Main:
multiscan    :     586.045 micros/op 53825 ops/sec 10.366 seconds 557968 operations; (multscans:14999)
rocksdb.non.last.level.read.bytes COUNT : 69107458693
rocksdb.non.last.level.read.count COUNT : 6724651
Branch:
multiscan    :     386.679 micros/op 81282 ops/sec 10.359 seconds 841968 operations; (multscans:25999)
rocksdb.non.last.level.read.bytes COUNT : 96605800558
rocksdb.non.last.level.read.count COUNT : 918973
```

Throughput is 36% higher with non-direct IO and 50% higher with direct IO. The improvement is likely from doing less number of I/Os due to I/O coalescing during Prepare(), as shown in `rocksdb.non.last.level.read.count`. The total bytes read is more with this PR for the same reason.

3. Regular iterator:
```
./db_bench --use_existing_db=1 --db="/tmp/rocksdbtest-543376/dbbench" --benchmarks=seekrandom --disable_auto_compactions=1 --seek_nexts=10 --threads=32 --duration=10

Main:
seekrandom   :      13.014 micros/op 2456735 ops/sec 10.014 seconds 24602968 operations; 2717.8 MB/s (773999 of 773999 found)
Branch:
seekrandom   :      13.048 micros/op 2450554 ops/sec 10.013 seconds 24537968 operations; 2710.9 MB/s (772999 of 772999 found)
```
The result fluctuates but without noticeable regression.

Reviewed By: anand1976

Differential Revision: D78440807

Pulled By: cbi42

fbshipit-source-id: 80ac6fd222696fa65ac0b4b5441748be5ee0b979

table/block_based/block_based_table_iterator.cc

@@ -37,6 +37,14 @@ void BlockBasedTableIterator::SeekImpl(const Slice* target,
                                        bool async_prefetch) {
   // TODO(hx235): set `seek_key_prefix_for_readahead_trimming_`
   // even when `target == nullptr` that is when `SeekToFirst()` is called
+  if (multi_scan_) {
+    if (SeekMultiScan(target)) {
+      return;
+    }
+  }
+
+  assert(!multi_scan_);
+
   if (target != nullptr && prefix_extractor_ &&
       read_options_.prefix_same_as_start) {
     const Slice& seek_user_key = ExtractUserKey(*target);
@@ -56,7 +64,7 @@ void BlockBasedTableIterator::SeekImpl(const Slice* target,
   ResetBlockCacheLookupVar();
 
   bool autotune_readaheadsize =
-      is_first_pass && read_options_.auto_readahead_size &&
+      read_options_.auto_readahead_size &&
       (read_options_.iterate_upper_bound || read_options_.prefix_same_as_start);
 
   if (autotune_readaheadsize &&
@@ -181,6 +189,7 @@ void BlockBasedTableIterator::SeekImpl(const Slice* target,
 }
 
 void BlockBasedTableIterator::SeekForPrev(const Slice& target) {
+  multi_scan_.reset();
   direction_ = IterDirection::kBackward;
   ResetBlockCacheLookupVar();
   is_out_of_bound_ = false;
@@ -255,6 +264,7 @@ void BlockBasedTableIterator::SeekForPrev(const Slice& target) {
 }
 
 void BlockBasedTableIterator::SeekToLast() {
+  multi_scan_.reset();
   direction_ = IterDirection::kBackward;
   ResetBlockCacheLookupVar();
   is_out_of_bound_ = false;
@@ -278,7 +288,9 @@ void BlockBasedTableIterator::SeekToLast() {
 }
 
 void BlockBasedTableIterator::Next() {
+  assert(Valid());
   if (is_at_first_key_from_index_ && !MaterializeCurrentBlock()) {
+    assert(!multi_scan_);
     return;
   }
   assert(block_iter_points_to_real_block_);
@@ -299,7 +311,9 @@ bool BlockBasedTableIterator::NextAndGetResult(IterateResult* result) {
 }
 
 void BlockBasedTableIterator::Prev() {
-  if (readahead_cache_lookup_ && !IsIndexAtCurr()) {
+  assert(!multi_scan_);
+  if ((readahead_cache_lookup_ && !IsIndexAtCurr()) || multi_scan_) {
+    multi_scan_.reset();
     // In case of readahead_cache_lookup_, index_iter_ has moved forward. So we
     // need to reseek the index_iter_ to point to current block by using
     // block_iter_'s key.
@@ -566,6 +580,10 @@ void BlockBasedTableIterator::FindKeyForward() {
 }
 
 void BlockBasedTableIterator::FindBlockForward() {
+  if (multi_scan_) {
+    FindBlockForwardInMultiScan();
+    return;
+  }
   // TODO the while loop inherits from two-level-iterator. We don't know
   // whether a block can be empty so it can be replaced by an "if".
   do {
@@ -901,4 +919,351 @@ void BlockBasedTableIterator::BlockCacheLookupForReadAheadSize(
   ResetPreviousBlockOffset();
 }
 
+// Note:
+// - Iterator should not be reused for multiple multiscans or mixing
+// multiscan with regular iterator usage.
+// - scan ranges should be non-overlapping, and have increasing start keys.
+// If a scan range's limit is not set, then there should only be one scan range.
+// - After Prepare(), the iterator expects Seek to be called on the start key
+// of each ScanOption in order. If any other seek is done, the optimization here
+// is aborted and fall back to vanilla iterator.
+// FIXME: DBIter and MergingIterator may
+// internally do Seek() on child iterators, e.g. due to
+// ReadOptions::max_skippable_internal_keys or reseeking into range deletion
+// end key. So these Seeks can cause iterator to fall back to normal
+// (non-prepared) iterator and ignore the optimizations done in Prepare().
+void BlockBasedTableIterator::Prepare(
+    const std::vector<ScanOptions>* scan_opts) {
+  index_iter_->Prepare(scan_opts);
+
+  assert(!multi_scan_);
+  if (multi_scan_) {
+    multi_scan_.reset();
+    return;
+  }
+  if (scan_opts == nullptr || scan_opts->empty()) {
+    return;
+  }
+  const bool has_limit = scan_opts->front().range.limit.has_value();
+  if (!has_limit && scan_opts->size() > 1) {
+    // Abort: overlapping ranges
+    return;
+  }
+
+  // Validate scan ranges to be increasing and with limit.
+  for (size_t i = 0; i < scan_opts->size(); ++i) {
+    const auto& scan_range = (*scan_opts)[i].range;
+    if (!scan_range.start.has_value()) {
+      // Abort: no start key
+      return;
+    }
+
+    // Assume for each scan range start <= limit.
+    if (scan_range.limit.has_value()) {
+      assert(user_comparator_.Compare(scan_range.start.value(),
+                                      scan_range.limit.value()) <= 0);
+    }
+
+    if (i > 0) {
+      if (!scan_range.limit.has_value()) {
+        // multiple no limit scan ranges
+        return;
+      }
+
+      const auto& last_end_key = (*scan_opts)[i - 1].range.limit.value();
+      if (user_comparator_.Compare(scan_range.start.value(), last_end_key) <
+          0) {
+        // Abort: overlapping ranges
+        return;
+      }
+    }
+  }
+
+  // Gather all relevant data block handles
+  std::vector<BlockHandle> blocks_to_prepare;
+  Status s;
+  std::vector<std::tuple<size_t, size_t>> block_ranges_per_scan;
+  for (const auto& scan_opt : *scan_opts) {
+    size_t num_blocks = 0;
+    // Current scan overlap the last block of the previous scan.
+    bool check_overlap = !blocks_to_prepare.empty();
+
+    // Scan range is specified in user key, here we seek to the minimum internal
+    // key with this user key.
+    InternalKey start_key(scan_opt.range.start.value(), kMaxSequenceNumber,
+                          kValueTypeForSeek);
+    index_iter_->Seek(start_key.Encode());
+    while (index_iter_->Valid() &&
+           (!scan_opt.range.limit.has_value() ||
+            user_comparator_.CompareWithoutTimestamp(
+                index_iter_->user_key(),
+                /*a_has_ts*/ true, *scan_opt.range.limit,
+                /*b_has_ts=*/false) <= 0)) {
+      if (check_overlap &&
+          blocks_to_prepare.back() == index_iter_->value().handle) {
+        // Skip the current block since it's already in the list
+      } else {
+        blocks_to_prepare.push_back(index_iter_->value().handle);
+      }
+      ++num_blocks;
+      index_iter_->Next();
+      check_overlap = false;
+    }
+    // Stop until index->key > limit
+    // Include the current block since it can still contain keys <= limit
+    if (index_iter_->Valid()) {
+      if (check_overlap &&
+          blocks_to_prepare.back() == index_iter_->value().handle) {
+        // Skip adding the current block since it's already in the list
+      } else {
+        blocks_to_prepare.push_back(index_iter_->value().handle);
+      }
+      ++num_blocks;
+    }
+
+    if (!index_iter_->status().ok()) {
+      // Abort: index iterator error
+      return;
+    }
+
+    block_ranges_per_scan.emplace_back(blocks_to_prepare.size() - num_blocks,
+                                       blocks_to_prepare.size());
+  }
+
+  // blocks_to_prepare has all the blocks that need to be read.
+  // Look up entries in cache and pin if exist.
+  // Store indices of blocks to read.
+  std::vector<size_t> blocks_to_read;
+  std::vector<CachableEntry<Block>> pinned_data_blocks_guard;
+  pinned_data_blocks_guard.resize(blocks_to_prepare.size());
+  for (size_t i = 0; i < blocks_to_prepare.size(); ++i) {
+    const auto& data_block_handle = blocks_to_prepare[i];
+    s = table_->LookupAndPinBlocksInCache<Block_kData>(
+        read_options_, data_block_handle,
+        &pinned_data_blocks_guard[i].As<Block_kData>());
+
+    if (!s.ok()) {
+      // Abort: block cache look up failed.
+      return;
+    }
+    if (!pinned_data_blocks_guard[i].GetValue()) {
+      // Block not in cache, will read it below.
+      blocks_to_read.emplace_back(i);
+    }
+  }
+
+  // Coalesce IOs
+  // TODO: limit prefetching size to bound memory usage.
+  if (!blocks_to_read.empty()) {
+    // Each vector correspond to blocks to read in a single read request.
+    // Each member in the vector is an index into blocks_to_prepare.
+    std::vector<std::vector<size_t>> collapsed_blocks_to_read(1);
+
+    // TODO: make this threshold configurable
+    constexpr size_t kCoalesceThreshold = 16 << 10;  // 16KB
+
+    for (const auto& block_idx : blocks_to_read) {
+      if (!collapsed_blocks_to_read.back().empty()) {
+        // Check if we can coalesce.
+        const auto& last_block =
+            blocks_to_prepare[collapsed_blocks_to_read.back().back()];
+        uint64_t last_block_end =
+            last_block.offset() +
+            BlockBasedTable::BlockSizeWithTrailer(last_block);
+        uint64_t current_start = blocks_to_prepare[block_idx].offset();
+
+        if (current_start > last_block_end + kCoalesceThreshold) {
+          // new IO
+          collapsed_blocks_to_read.emplace_back();
+        }
+      }
+      collapsed_blocks_to_read.back().emplace_back(block_idx);
+    }
+
+    // do IO
+    IOOptions io_opts;
+    s = table_->get_rep()->file->PrepareIOOptions(read_options_, io_opts);
+    if (!s.ok()) {
+      // Abort: PrepareIOOptions failed
+      return;
+    }
+
+    // Init read requests for Multi-Read
+    std::vector<FSReadRequest> read_reqs;
+    read_reqs.reserve(collapsed_blocks_to_read.size());
+    size_t total_len = 0;
+    for (const auto& blocks : collapsed_blocks_to_read) {
+      assert(blocks.size());
+      const auto& first_block = blocks_to_prepare[blocks[0]];
+      const auto& last_block = blocks_to_prepare[blocks.back()];
+
+      const auto start_offset = first_block.offset();
+      const auto end_offset = last_block.offset() +
+                              BlockBasedTable::BlockSizeWithTrailer(last_block);
+      assert(end_offset > start_offset);
+      FSReadRequest read_req;
+      read_req.offset = start_offset;
+      read_req.len = end_offset - start_offset;
+      total_len += read_req.len;
+      read_reqs.emplace_back(std::move(read_req));
+    }
+
+    // Init buffer for read
+    std::unique_ptr<char[]> buf;
+    const bool direct_io = table_->get_rep()->file->use_direct_io();
+    if (direct_io) {
+      for (auto& read_req : read_reqs) {
+        read_req.scratch = nullptr;
+      }
+    } else {
+      // TODO: optimize if FSSupportedOps::kFSBuffer is supported.
+      buf.reset(new char[total_len]);
+      size_t offset = 0;
+      for (auto& read_req : read_reqs) {
+        read_req.scratch = buf.get() + offset;
+        offset += read_req.len;
+      }
+    }
+
+    AlignedBuf aligned_buf;
+    s = table_->get_rep()->file.get()->MultiRead(
+        io_opts, read_reqs.data(), read_reqs.size(),
+        direct_io ? &aligned_buf : nullptr);
+    if (!s.ok()) {
+      return;
+    }
+    for (auto& req : read_reqs) {
+      if (!req.status.ok()) {
+        return;
+      }
+    }
+
+    // Init blocks and pin them in block cache.
+    MemoryAllocator* memory_allocator =
+        table_->get_rep()->table_options.block_cache->memory_allocator();
+    for (size_t i = 0; i < collapsed_blocks_to_read.size(); i++) {
+      const auto& blocks = collapsed_blocks_to_read[i];
+      const auto& read_req = read_reqs[i];
+      for (const auto& block_idx : blocks) {
+        const auto& block = blocks_to_prepare[block_idx];
+        const auto block_size_with_trailer =
+            BlockBasedTable::BlockSizeWithTrailer(block);
+        const auto block_offset_in_buffer = block.offset() - read_req.offset;
+
+        CacheAllocationPtr data =
+            AllocateBlock(block_size_with_trailer, memory_allocator);
+        memcpy(data.get(), read_req.result.data() + block_offset_in_buffer,
+               block_size_with_trailer);
+        BlockContents tmp_contents(std::move(data), block.size());
+
+#ifndef NDEBUG
+        tmp_contents.has_trailer =
+            table_->get_rep()->footer.GetBlockTrailerSize() > 0;
+#endif
+        assert(pinned_data_blocks_guard[block_idx].IsEmpty());
+        s = table_->CreateAndPinBlockInCache<Block_kData>(
+            read_options_, block, &tmp_contents,
+            &(pinned_data_blocks_guard[block_idx].As<Block_kData>()));
+        if (!s.ok()) {
+          // Abort: failed to create and pin block in cache
+          return;
+        }
+      }
+    }
+  }
+
+  // Successful Prepare, init related states so the iterator reads from prepared
+  // blocks
+  multi_scan_.reset(new MultiScanState(scan_opts,
+                                       std::move(pinned_data_blocks_guard),
+                                       std::move(block_ranges_per_scan)));
+  is_index_at_curr_block_ = false;
+  block_iter_points_to_real_block_ = false;
+}
+
+bool BlockBasedTableIterator::SeekMultiScan(const Slice* target) {
+  assert(multi_scan_);
+  // This is a MultiScan and Preapre() has been called.
+  //
+  // Validate seek key with scan options
+  if (multi_scan_->next_scan_idx >= multi_scan_->scan_opts->size()) {
+    multi_scan_.reset();
+  } else if (!target) {
+    // start key must be set for multi-scan
+    multi_scan_.reset();
+  } else if (user_comparator_.CompareWithoutTimestamp(
+                 ExtractUserKey(*target), /*a_has_ts=*/true,
+                 (*multi_scan_->scan_opts)[multi_scan_->next_scan_idx]
+                     .range.start.value(),
+                 /*b_has_ts=*/false) != 0) {
+    // Unexpected seek key
+    multi_scan_.reset();
+  } else {
+    auto [cur_scan_start_idx, cur_scan_end_idx] =
+        multi_scan_->block_ranges_per_scan[multi_scan_->next_scan_idx];
+    // We should have the data block already loaded
+    ++multi_scan_->next_scan_idx;
+    if (cur_scan_start_idx >= cur_scan_end_idx) {
+      is_out_of_bound_ = true;
+      assert(!Valid());
+      return true;
+    } else {
+      is_out_of_bound_ = false;
+    }
+
+    if (!block_iter_points_to_real_block_ ||
+        multi_scan_->cur_data_block_idx != cur_scan_start_idx) {
+      if (block_iter_points_to_real_block_) {
+        // Should be scan in increasing key range.
+        // All blocks before cur_data_block_idx_ are not pinned anymore.
+        assert(multi_scan_->cur_data_block_idx < cur_scan_start_idx);
+      }
+
+      ResetDataIter();
+      // Note that the block_iter_ takes ownership of the pinned data block
+      // TODO: we can delegate the clean up like with pinned_iters_mgr_ if
+      // need to pin blocks longer.
+      table_->NewDataBlockIterator<DataBlockIter>(
+          read_options_, multi_scan_->pinned_data_blocks[cur_scan_start_idx],
+          &block_iter_, Status::OK());
+    }
+    multi_scan_->cur_data_block_idx = cur_scan_start_idx;
+    block_iter_points_to_real_block_ = true;
+    block_iter_.Seek(*target);
+    FindKeyForward();
+    return true;
+  }
+
+  return false;
+}
+
+void BlockBasedTableIterator::FindBlockForwardInMultiScan() {
+  assert(multi_scan_);
+  assert(multi_scan_->next_scan_idx >= 1);
+  const auto cur_scan_end_idx = std::get<1>(
+      multi_scan_->block_ranges_per_scan[multi_scan_->next_scan_idx - 1]);
+  do {
+    if (!block_iter_.status().ok()) {
+      return;
+    }
+
+    if (multi_scan_->cur_data_block_idx + 1 >= cur_scan_end_idx) {
+      // We don't ResetDataIter() here since next scan might be reading from
+      // the same block. ResetDataIter() will free the underlying block cache
+      // handle and we don't want the block to be unpinned.
+      is_out_of_bound_ = true;
+      assert(!Valid());
+      return;
+    }
+    // Move to the next pinned data block
+    ResetDataIter();
+    ++multi_scan_->cur_data_block_idx;
+    table_->NewDataBlockIterator<DataBlockIter>(
+        read_options_,
+        multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx],
+        &block_iter_, Status::OK());
+    block_iter_points_to_real_block_ = true;
+    block_iter_.SeekToFirst();
+  } while (!block_iter_.Valid());
+}
 }  // namespace ROCKSDB_NAMESPACE

table/block_based/block_based_table_iterator.h

@@ -41,11 +41,11 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
             compaction_readahead_size,
             table_->get_rep()->table_options.initial_auto_readahead_size),
         allow_unprepared_value_(allow_unprepared_value),
-        block_iter_points_to_real_block_(false),
         check_filter_(check_filter),
         need_upper_bound_check_(need_upper_bound_check),
         async_read_in_progress_(false),
-        is_last_level_(table->IsLastLevel()) {}
+        is_last_level_(table->IsLastLevel()),
+        block_iter_points_to_real_block_(false) {}
 
   ~BlockBasedTableIterator() override { ClearBlockHandles(); }
 
@@ -69,6 +69,7 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
   Slice key() const override {
     assert(Valid());
     if (is_at_first_key_from_index_) {
+      assert(!multi_scan_);
       return index_iter_->value().first_internal_key;
     } else {
       return block_iter_.key();
@@ -141,10 +142,12 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
     // Prefix index set status to NotFound when the prefix does not exist.
     if (IsIndexAtCurr() && !index_iter_->status().ok() &&
         !index_iter_->status().IsNotFound()) {
+      assert(!multi_scan_);
       return index_iter_->status();
     } else if (block_iter_points_to_real_block_) {
       return block_iter_.status();
     } else if (async_read_in_progress_) {
+      assert(!multi_scan_);
       return Status::TryAgain("Async read in progress");
     } else {
       return Status::OK();
@@ -222,9 +225,7 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
     }
   }
 
-  void Prepare(const std::vector<ScanOptions>* scan_opts) override {
-    index_iter_->Prepare(scan_opts);
-  }
+  void Prepare(const std::vector<ScanOptions>* scan_opts) override;
 
   FilePrefetchBuffer* prefetch_buffer() {
     return block_prefetcher_.prefetch_buffer();
@@ -312,12 +313,20 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
 
   BlockPrefetcher block_prefetcher_;
 
+  // It stores all the block handles that are lookuped in cache ahead when
+  // BlockCacheLookupForReadAheadSize is called. Since index_iter_ may point to
+  // different blocks when readahead_size is calculated in
+  // BlockCacheLookupForReadAheadSize, to avoid index_iter_ reseek,
+  // block_handles_ is used.
+  // `block_handles_` is lazily constructed to save CPU when it is unused
+  std::unique_ptr<std::deque<BlockHandleInfo>> block_handles_;
+
+  // The prefix of the key called with SeekImpl().
+  // This is for readahead trimming so no data blocks containing keys of a
+  // different prefix are prefetched
+  std::string seek_key_prefix_for_readahead_trimming_ = "";
+
   const bool allow_unprepared_value_;
-  // True if block_iter_ is initialized and points to the same block
-  // as index iterator.
-  bool block_iter_points_to_real_block_;
-  // See InternalIteratorBase::IsOutOfBound().
-  bool is_out_of_bound_ = false;
   // How current data block's boundary key with the next block is compared with
   // iterate upper bound.
   BlockUpperBound block_upper_bound_check_ = BlockUpperBound::kUnknown;
@@ -337,18 +346,6 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
   // size based on cache hit and miss.
   bool readahead_cache_lookup_ = false;
 
-  // It stores all the block handles that are lookuped in cache ahead when
-  // BlockCacheLookupForReadAheadSize is called. Since index_iter_ may point to
-  // different blocks when readahead_size is calculated in
-  // BlockCacheLookupForReadAheadSize, to avoid index_iter_ reseek,
-  // block_handles_ is used.
-  // `block_handles_` is lazily constructed to save CPU when it is unused
-  std::unique_ptr<std::deque<BlockHandleInfo>> block_handles_;
-
-  // During cache lookup to find readahead size, index_iter_ is iterated and it
-  // can point to a different block. is_index_at_curr_block_ keeps track of
-  // that.
-  bool is_index_at_curr_block_ = true;
   bool is_index_out_of_bound_ = false;
 
   // Used in case of auto_readahead_size to disable the block_cache lookup if
@@ -357,10 +354,48 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
   // is used to disable the lookup.
   IterDirection direction_ = IterDirection::kForward;
 
-  // The prefix of the key called with SeekImpl().
-  // This is for readahead trimming so no data blocks containing keys of a
-  // different prefix are prefetched
-  std::string seek_key_prefix_for_readahead_trimming_ = "";
+  //*** BEGIN States used by both regular scan and multiscan
+
+  // True if block_iter_ is initialized and points to the same block
+  // as index iterator.
+  bool block_iter_points_to_real_block_;
+  // See InternalIteratorBase::IsOutOfBound().
+  bool is_out_of_bound_ = false;
+  // During cache lookup to find readahead size, index_iter_ is iterated and it
+  // can point to a different block.
+  // If Prepare() is called, index_iter_ is used to prefetch data blocks for the
+  // multiscan, so is_index_at_curr_block_ will be false.
+  // Whether index is expected to match the current data_block_iter_.
+  bool is_index_at_curr_block_ = true;
+
+  // *** END States used by both regular scan and multiscan
+
+  // *** BEGIN MultiScan related states ***
+  struct MultiScanState {
+    // bool prepared_ = false;
+    const std::vector<ScanOptions>* scan_opts;
+    std::vector<CachableEntry<Block>> pinned_data_blocks;
+
+    // Indicies into multiscan_pinned_data_blocks_ for data blocks that are
+    // relevant for each scan range.
+    // inclusive start, exclusive end
+    std::vector<std::tuple<size_t, size_t>> block_ranges_per_scan;
+    size_t next_scan_idx;
+    size_t cur_data_block_idx;
+
+    MultiScanState(
+        const std::vector<ScanOptions>* _scan_opts,
+        std::vector<CachableEntry<Block>>&& _pinned_data_blocks,
+        std::vector<std::tuple<size_t, size_t>>&& _block_ranges_per_scan)
+        : scan_opts(_scan_opts),
+          pinned_data_blocks(std::move(_pinned_data_blocks)),
+          block_ranges_per_scan(std::move(_block_ranges_per_scan)),
+          next_scan_idx(0),
+          cur_data_block_idx(0) {}
+  };
+
+  std::unique_ptr<MultiScanState> multi_scan_;
+  // *** END MultiScan related APIs and states ***
 
   void SeekSecondPass(const Slice* target);
 
@@ -476,5 +511,12 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
                                     uint64_t& end_updated_offset,
                                     size_t& prev_handles_size);
   // *** END APIs relevant to auto tuning of readahead_size ***
+
+  // *** BEGIN APIs relevant to multiscan ***
+  // Returns true iff seek is successful.
+  bool SeekMultiScan(const Slice* target);
+
+  void FindBlockForwardInMultiScan();
+  // *** END APIs relevant to multiscan ***
 };
 }  // namespace ROCKSDB_NAMESPACE

table/block_based/block_based_table_reader.cc

@@ -106,7 +106,11 @@ CacheAllocationPtr CopyBufferToHeap(MemoryAllocator* allocator, Slice& buf) {
       bool use_block_cache_for_lookup) const;                                  \
   template Status BlockBasedTable::LookupAndPinBlocksInCache<T>(               \
       const ReadOptions& ro, const BlockHandle& handle,                        \
-      CachableEntry<T>* out_parsed_block) const;
+      CachableEntry<T>* out_parsed_block) const;                               \
+  template Status BlockBasedTable::CreateAndPinBlockInCache<T>(                \
+      const ReadOptions& ro, const BlockHandle& handle,                        \
+      BlockContents* block_contents, CachableEntry<T>* out_parsed_block)       \
+      const;
 
 INSTANTIATE_BLOCKLIKE_TEMPLATES(ParsedFullFilterBlock);
 INSTANTIATE_BLOCKLIKE_TEMPLATES(DecompressorDict);
@@ -1735,6 +1739,17 @@ Status BlockBasedTable::LookupAndPinBlocksInCache(
   return s;
 }
 
+template <typename TBlocklike>
+Status BlockBasedTable::CreateAndPinBlockInCache(
+    const ReadOptions& ro, const BlockHandle& handle, BlockContents* contents,
+    CachableEntry<TBlocklike>* out_parsed_block) const {
+  return MaybeReadBlockAndLoadToCache(
+      nullptr, ro, handle, rep_->decompressor.get(),
+      /*for_compaction=*/false, out_parsed_block, nullptr, nullptr, contents,
+      /*async_read=*/false,
+      /*use_block_cache_for_lookup=*/true);
+}
+
 // If contents is nullptr, this function looks up the block caches for the
 // data block referenced by handle, and read the block from disk if necessary.
 // If contents is non-null, it skips the cache lookup and disk read, since

table/block_based/block_based_table_reader.h

@@ -299,11 +299,21 @@ class BlockBasedTable : public TableReader {
   Status GetKVPairsFromDataBlocks(const ReadOptions& read_options,
                                   std::vector<KVPairBlock>* kv_pair_blocks);
 
+  // Look up the block cache for the specified block.
+  // out_parsed_block is set to nullptr if the block is not found in the cache.
   template <typename TBlocklike>
   Status LookupAndPinBlocksInCache(
       const ReadOptions& ro, const BlockHandle& handle,
       CachableEntry<TBlocklike>* out_parsed_block) const;
 
+  // Create the block given in `block_contents` and insert it into block cache.
+  // `out_parsed_block` points to the inserted block if successful.
+  template <typename TBlocklike>
+  Status CreateAndPinBlockInCache(
+      const ReadOptions& ro, const BlockHandle& handle,
+      BlockContents* block_contents,
+      CachableEntry<TBlocklike>* out_parsed_block) const;
+
   struct Rep;
 
   Rep* get_rep() { return rep_; }

table/block_based/block_based_table_reader_test.cc

@@ -173,7 +173,7 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
         0 /* _tail_size */, user_defined_timestamps_persisted);
 
     std::unique_ptr<RandomAccessFileReader> file;
-    NewFileReader(table_name, foptions, &file);
+    NewFileReader(table_name, foptions, &file, ioptions.statistics.get());
 
     uint64_t file_size = 0;
     ASSERT_OK(env_->GetFileSize(Path(table_name), &file_size));
@@ -222,12 +222,15 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
   }
 
   void NewFileReader(const std::string& filename, const FileOptions& opt,
-                     std::unique_ptr<RandomAccessFileReader>* reader) {
+                     std::unique_ptr<RandomAccessFileReader>* reader,
+                     Statistics* stats = nullptr) {
     std::string path = Path(filename);
     std::unique_ptr<FSRandomAccessFile> f;
     ASSERT_OK(fs_->NewRandomAccessFile(path, opt, &f, nullptr));
     reader->reset(new RandomAccessFileReader(std::move(f), path,
-                                             env_->GetSystemClock().get()));
+                                             env_->GetSystemClock().get(),
+                                             /*io_tracer=*/nullptr,
+                                             /*stats=*/stats));
   }
 };
 
@@ -990,6 +993,167 @@ TEST_P(BlockBasedTableReaderTestVerifyChecksum, ChecksumMismatch) {
   ASSERT_EQ(s.code(), Status::kCorruption);
 }
 
+TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
+  Options options;
+  options.statistics = CreateDBStatistics();
+  ReadOptions read_opts;
+  size_t ts_sz = options.comparator->timestamp_size();
+  std::vector<std::pair<std::string, std::string>> kv =
+      BlockBasedTableReaderBaseTest::GenerateKVMap(
+          100 /* num_block */,
+          true /* mixed_with_human_readable_string_value */, ts_sz);
+
+  std::string table_name = "BlockBasedTableReaderTest_NewIterator" +
+                           CompressionTypeToString(compression_type_);
+
+  ImmutableOptions ioptions(options);
+  CreateTable(table_name, ioptions, compression_type_, kv,
+              compression_parallel_threads_, compression_dict_bytes_);
+
+  std::unique_ptr<BlockBasedTable> table;
+  FileOptions foptions;
+  foptions.use_direct_reads = true;
+  InternalKeyComparator comparator(options.comparator);
+  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table,
+                           true /* bool prefetch_index_and_filter_in_cache */,
+                           nullptr /* status */, persist_udt_);
+
+  std::unique_ptr<InternalIterator> iter;
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+
+  // Should coalesce into a single I/O
+  std::vector<ScanOptions> scan_options(
+      {ScanOptions(ExtractUserKey(kv[0].first),
+                   ExtractUserKey(kv[kEntriesPerBlock].first)),
+       ScanOptions(ExtractUserKey(kv[2 * kEntriesPerBlock].first),
+                   ExtractUserKey(kv[3 * kEntriesPerBlock].first))});
+
+  auto read_count_before =
+      options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT);
+  iter->Prepare(&scan_options);
+  auto read_count_after =
+      options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT);
+  ASSERT_EQ(read_count_before + 1, read_count_after);
+  iter->Seek(kv[0].first);
+  for (size_t i = 0; i < kEntriesPerBlock + 1; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  // Iter may still be valid after scan range. Upper layer (DBIter) handles
+  // exact upper bound checking. So we don't check !iter->Valid() here.
+  ASSERT_OK(iter->status());
+  iter->Seek(kv[2 * kEntriesPerBlock].first);
+  for (size_t i = 2 * kEntriesPerBlock; i < 3 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_OK(iter->status());
+
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+  // No IO coalesce, should do MultiRead with 2 read requests.
+  scan_options = {ScanOptions(ExtractUserKey(kv[70 * kEntriesPerBlock].first),
+                              ExtractUserKey(kv[75 * kEntriesPerBlock].first)),
+                  ScanOptions(ExtractUserKey(kv[90 * kEntriesPerBlock].first),
+                              ExtractUserKey(kv[95 * kEntriesPerBlock].first))};
+  read_count_before =
+      options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT);
+  iter->Prepare(&scan_options);
+  read_count_after =
+      options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT);
+  ASSERT_EQ(read_count_before + 2, read_count_after);
+
+  iter->Seek(kv[70 * kEntriesPerBlock].first);
+  for (size_t i = 70 * kEntriesPerBlock; i < 75 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_OK(iter->status());
+  iter->Seek(kv[90 * kEntriesPerBlock].first);
+  for (size_t i = 90 * kEntriesPerBlock; i < 95 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_OK(iter->status());
+
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+  // Should do two I/Os since blocks 80-81 and 90-95 are already in block cache,
+  // reads from blocks 50-79 and 82-.. are co
+  scan_options = {ScanOptions(ExtractUserKey(kv[50 * kEntriesPerBlock].first))};
+  read_count_before =
+      options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT);
+  iter->Prepare(&scan_options);
+  read_count_after =
+      options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT);
+  ASSERT_EQ(read_count_before + 3, read_count_after);
+  iter->Seek(kv[50 * kEntriesPerBlock].first);
+  for (size_t i = 50 * kEntriesPerBlock; i < 100 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_FALSE(iter->Valid());
+  ASSERT_OK(iter->status());
+
+  // Check cases when Seek key does not match start key in ScanOptions
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+  scan_options = {ScanOptions(ExtractUserKey(kv[10 * kEntriesPerBlock].first),
+                              ExtractUserKey(kv[20 * kEntriesPerBlock].first)),
+                  ScanOptions(ExtractUserKey(kv[30 * kEntriesPerBlock].first),
+                              ExtractUserKey(kv[40 * kEntriesPerBlock].first))};
+  iter->Prepare(&scan_options);
+  // Match start key
+  iter->Seek(kv[10 * kEntriesPerBlock].first);
+  for (size_t i = 10 * kEntriesPerBlock; i < 20 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_OK(iter->status());
+  // Does not match start key of the second ScanOptions.
+  iter->Seek(kv[50 * kEntriesPerBlock + 1].first);
+  for (size_t i = 50 * kEntriesPerBlock + 1; i < 100 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_FALSE(iter->Valid());
+  ASSERT_OK(iter->status());
+
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+  scan_options = {ScanOptions(ExtractUserKey(kv[10 * kEntriesPerBlock].first)),
+                  ScanOptions(ExtractUserKey(kv[11 * kEntriesPerBlock].first))};
+  iter->Prepare(&scan_options);
+  // Does not match the first ScanOptions.
+  iter->SeekToFirst();
+  for (size_t i = 0; i < kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_OK(iter->status());
+  iter->Seek(kv[10 * kEntriesPerBlock].first);
+  for (size_t i = 10 * kEntriesPerBlock; i < 12 * kEntriesPerBlock; ++i) {
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ(iter->key().ToString(), kv[i].first);
+    iter->Next();
+  }
+  ASSERT_OK(iter->status());
+}
+
 // Param 1: compression type
 // Param 2: whether to use direct reads
 // Param 3: Block Based Table Index type, partitioned filters are also enabled

table/format.h

@@ -55,7 +55,7 @@ class BlockHandle {
   uint64_t offset() const { return offset_; }
   void set_offset(uint64_t _offset) { offset_ = _offset; }
 
-  // The size of the stored block
+  // The size of the stored block, this size does not include the block trailer.
   uint64_t size() const { return size_; }
   void set_size(uint64_t _size) { size_ = _size; }
 

unreleased_history/performance_improvements/bbiter-multiscan.md

@@ -0,0 +1 @@
+* Optimized MultiScan using BlockBasedTable to coalesce I/Os and prefetch all data blocks.