Add option to limit max prefetching in MultiScan (facebook#13920)

Summary:
Add a new option `MultiScanArgs::max_prefetch_size` that limits the memory usage of per file pinning of prefetched blocks. Note that this only accounts for compressed block size. This is intended to be a stopgap until we implement some kind of global prefetch manager that limits the global multiscan memory usage.

Pull Request resolved: facebook#13920

Test Plan: new unit test `./block_based_table_reader_test --gtest_filter="*MultiScanPrefetchSizeLimit/*"`

Reviewed By: xingbowang

Differential Revision: D81630629

Pulled By: cbi42

fbshipit-source-id: 9f66678915242fe1220620531a4b9fd22747cdea

Author: cbi42

include/rocksdb/options.h

@@ -1790,16 +1790,19 @@ class MultiScanArgs {
     comp_ = other.comp_;
     original_ranges_ = other.original_ranges_;
     io_coalesce_threshold = other.io_coalesce_threshold;
+    max_prefetch_size = other.max_prefetch_size;
   }
   MultiScanArgs(MultiScanArgs&& other) noexcept
       : io_coalesce_threshold(other.io_coalesce_threshold),
+        max_prefetch_size(other.max_prefetch_size),
         comp_(other.comp_),
         original_ranges_(std::move(other.original_ranges_)) {}
 
   MultiScanArgs& operator=(const MultiScanArgs& other) {
     comp_ = other.comp_;
     original_ranges_ = other.original_ranges_;
     io_coalesce_threshold = other.io_coalesce_threshold;
+    max_prefetch_size = other.max_prefetch_size;
     return *this;
   }
 
@@ -1808,6 +1811,7 @@ class MultiScanArgs {
       comp_ = other.comp_;
       original_ranges_ = std::move(other.original_ranges_);
       io_coalesce_threshold = other.io_coalesce_threshold;
+      max_prefetch_size = other.max_prefetch_size;
     }
     return *this;
   }
@@ -1849,6 +1853,18 @@ class MultiScanArgs {
 
   uint64_t io_coalesce_threshold = 16 << 10;  // 16KB by default
 
+  // Maximum size (in bytes) for the data blocks loaded by a MultiScan.
+  // This limits the amount of I/O and memory usage by pinned data blocks.
+  //
+  // When set to 0 (the default), there is no limit. When the limit is reached,
+  // the iterator will start returning Status::PrefetchLimitReached().
+  //
+  // Note that prefetching happens only once in Prepare(), which is different
+  // from ReadOptions::readahead_size, which applies any time the iterator does
+  // I/O.
+  // Note that this limit is per file and applies to compressed block size.
+  uint64_t max_prefetch_size = 0;
+
  private:
   // The comparator used for ordering ranges
   const Comparator* comp_;

include/rocksdb/status.h

@@ -115,6 +115,7 @@ class Status {
     kIOFenced = 14,
     kMergeOperatorFailed = 15,
     kMergeOperandThresholdExceeded = 16,
+    kPrefetchLimitReached = 17,
     kMaxSubCode
   };
 
@@ -318,6 +319,10 @@ class Status {
 
   static Status LockLimit() { return Status(kAborted, kLockLimit); }
 
+  static Status PrefetchLimitReached() {
+    return Status(kIncomplete, kPrefetchLimitReached);
+  }
+
   // Returns true iff the status indicates success.
   bool ok() const {
     MarkChecked();
@@ -486,6 +491,13 @@ class Status {
     return (code() == kIOError) && (subcode() == kIOFenced);
   }
 
+  // Returns true iff the status indicates prefetch limit reached during
+  // MultiScan.
+  bool IsPrefetchLimitReached() const {
+    MarkChecked();
+    return (code() == kIncomplete) && (subcode() == kPrefetchLimitReached);
+  }
+
   // Return a string representation of this status suitable for printing.
   // Returns the string "OK" for success.
   std::string ToString() const;

table/block_based/block_based_table_iterator.cc

@@ -982,7 +982,6 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
 
   // 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;
@@ -1042,11 +1041,26 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
   // 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());
+  std::vector<CachableEntry<Block>> pinned_data_blocks_guard(
+      blocks_to_prepare.size());
+  uint64_t total_prefetch_size = 0;
+
   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>(
+
+    // Check if we would exceed the prefetch size limit with this block
+    total_prefetch_size +=
+        BlockBasedTable::BlockSizeWithTrailer(data_block_handle);
+    if (multiscan_opts->max_prefetch_size > 0 &&
+        total_prefetch_size > multiscan_opts->max_prefetch_size) {
+      // All remaining blocks are by default empty.
+      for (size_t j = i; j < blocks_to_prepare.size(); ++j) {
+        assert(pinned_data_blocks_guard[j].IsEmpty());
+      }
+      break;
+    }
+
+    Status s = table_->LookupAndPinBlocksInCache<Block_kData>(
         read_options_, data_block_handle,
         &pinned_data_blocks_guard[i].As<Block_kData>());
 
@@ -1088,10 +1102,13 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
 
     // do IO
     IOOptions io_opts;
-    s = table_->get_rep()->file->PrepareIOOptions(read_options_, io_opts);
-    if (!s.ok()) {
-      // Abort: PrepareIOOptions failed
-      return;
+    {
+      Status s =
+          table_->get_rep()->file->PrepareIOOptions(read_options_, io_opts);
+      if (!s.ok()) {
+        // Abort: PrepareIOOptions failed
+        return;
+      }
     }
 
     // Init read requests for Multi-Read
@@ -1163,11 +1180,13 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
     }
 
     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;
+    {
+      Status 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()) {
@@ -1181,7 +1200,8 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
         table_->get_rep()->decompressor.get();
     CachableEntry<DecompressorDict> cached_dict;
     if (table_->get_rep()->uncompression_dict_reader) {
-      s = table_->get_rep()
+      Status s =
+          table_->get_rep()
               ->uncompression_dict_reader->GetOrReadUncompressionDictionary(
                   /* prefetch_buffer= */ nullptr, read_options_,
                   /* get_context= */ nullptr, /* lookup_context= */ nullptr,
@@ -1226,7 +1246,7 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
             table_->get_rep()->footer.GetBlockTrailerSize() > 0;
 #endif
         assert(pinned_data_blocks_guard[block_idx].IsEmpty());
-        s = table_->CreateAndPinBlockInCache<Block_kData>(
+        Status s = table_->CreateAndPinBlockInCache<Block_kData>(
             read_options_, block, decompressor, &tmp_contents,
             &(pinned_data_blocks_guard[block_idx].As<Block_kData>()));
         if (!s.ok()) {
@@ -1290,6 +1310,16 @@ bool BlockBasedTableIterator::SeekMultiScan(const Slice* target) {
       }
 
       ResetDataIter();
+
+      // Check if we've hit an empty entry indicating prefetch limit reached
+      if (multi_scan_->pinned_data_blocks[cur_scan_start_idx].IsEmpty()) {
+        multi_scan_->cur_data_block_idx = cur_scan_start_idx;
+        multi_scan_->prefetch_limit_reached = true;
+        assert(!Valid());
+        assert(status().IsPrefetchLimitReached());
+        return true;
+      }
+
       // 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.
@@ -1346,6 +1376,16 @@ void BlockBasedTableIterator::FindBlockForwardInMultiScan() {
     // Move to the next pinned data block
     ResetDataIter();
     ++multi_scan_->cur_data_block_idx;
+
+    // Check if we've hit an empty entry indicating prefetch limit reached
+    if (multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx]
+            .IsEmpty()) {
+      multi_scan_->prefetch_limit_reached = true;
+      assert(!Valid());
+      assert(status().IsPrefetchLimitReached());
+      return;
+    }
+
     table_->NewDataBlockIterator<DataBlockIter>(
         read_options_,
         multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx],

table/block_based/block_based_table_iterator.h

@@ -145,10 +145,14 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
       assert(!multi_scan_);
       return index_iter_->status();
     } else if (block_iter_points_to_real_block_) {
+      // This is the common case.
       return block_iter_.status();
     } else if (async_read_in_progress_) {
       assert(!multi_scan_);
       return Status::TryAgain("Async read in progress");
+    } else if (multi_scan_ && multi_scan_->prefetch_limit_reached) {
+      assert(!Valid());
+      return Status::PrefetchLimitReached();
     } else {
       return Status::OK();
     }
@@ -385,6 +389,10 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
     size_t next_scan_idx;
     size_t cur_data_block_idx;
 
+    // When true, the iterator will return
+    // Status::Incomplete(Status::kPrefetchLimitReached).
+    bool prefetch_limit_reached;
+
     MultiScanState(
         const MultiScanArgs* _scan_opts,
         std::vector<CachableEntry<Block>>&& _pinned_data_blocks,
@@ -393,7 +401,8 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
           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) {}
+          cur_data_block_idx(0),
+          prefetch_limit_reached(false) {}
   };
 
   std::unique_ptr<MultiScanState> multi_scan_;