Fix handling of out-of-range scan option (facebook#13995)

Summary:
currently BlockBasedTableIterator::Prepare() fails the iterator with non-ok status if an out-of-range scan option is detected. This is due to the interaction between LevelIterator and BlockBasedTableIterator, see added comment above BlockBasedTableIterator::Prepare(). This can fail stress test for L0 files since it doesn't use LevelIterator and scan options are not pruned. This PR fixes this by adding an internal option to MultiScanArgs that enables this check.

Pull Request resolved: facebook#13995

Test Plan:
- new unit test
- stress test that fails before this pr: `python3 -u ./tools/db_crashtest.py whitebox --iterpercent=60 --prefix_size=-1 --prefixpercent=0 --readpercent=0 --test_batches_snapshots=0 --use_multiscan=1 --read_fault_one_in=0 --kill_random_test=88888 --interval=60 --multiscan_use_async_io=0 --mmap_read=0 --level0_file_num_compaction_trigger=20`

Reviewed By: anand1976

Differential Revision: D83166088

Pulled By: cbi42

fbshipit-source-id: 241a7d43c8c00d9a98eea0cabb03d2174d51aae5

Author: cbi42

db/db_iterator_test.cc

@@ -4476,6 +4476,74 @@ TEST_P(DBMultiScanIteratorTest, RangeBetweenFiles) {
   iter.reset();
 }
 
+TEST_P(DBMultiScanIteratorTest, OutOfL0FileRange) {
+  // Test that prepare does not fail scan when a scan range
+  // is outside of a L0 file's key range.
+  auto options = CurrentOptions();
+  options.compression = kNoCompression;
+  DestroyAndReopen(options);
+
+  Random rnd(301);
+  // Create a Lmax file
+  // key01 ~ key99
+  for (int i = 0; i < 100; ++i) {
+    std::stringstream ss;
+    ss << std::setw(2) << std::setfill('0') << i;
+    ASSERT_OK(Put("k" + ss.str(), rnd.RandomString(1024)));
+  }
+  ASSERT_OK(Flush());
+  CompactRangeOptions cro;
+  cro.bottommost_level_compaction = BottommostLevelCompaction::kForce;
+  ASSERT_OK(db_->CompactRange(cro, nullptr, nullptr));
+
+  // Create a L0 file
+  // key00 ~ key09
+  for (int i = 0; i < 10; ++i) {
+    std::stringstream ss;
+    ss << std::setw(2) << std::setfill('0') << i;
+    ASSERT_OK(Put("k" + ss.str(), rnd.RandomString(1024)));
+  }
+  ASSERT_OK(Flush());
+  ASSERT_EQ(NumTableFilesAtLevel(0), 1);
+
+  // The second range is outside of L0 file's key range
+  std::vector<std::string> key_ranges({"k04", "k06", "k12", "k14"});
+  ReadOptions ro;
+  Slice ub;
+  ro.iterate_upper_bound = &ub;
+  ro.fill_cache = GetParam();
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.insert(key_ranges[0], key_ranges[1]);
+  scan_options.insert(key_ranges[2], key_ranges[3]);
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+  std::unique_ptr<Iterator> iter(dbfull()->NewIterator(ro, cfh));
+  ASSERT_NE(iter, nullptr);
+  iter->Prepare(scan_options);
+  int count = 0;
+  ub = key_ranges[1];
+  iter->Seek(key_ranges[0]);
+  while (iter->status().ok() && iter->Valid()) {
+    ASSERT_GE(iter->key().compare(key_ranges[0]), 0);
+    ASSERT_LT(iter->key().compare(key_ranges[1]), 0);
+    count++;
+    iter->Next();
+  }
+  ASSERT_OK(iter->status()) << iter->status().ToString();
+  ASSERT_EQ(count, 2);
+
+  ub = key_ranges[3];
+  count = 0;
+  iter->Seek(key_ranges[2]);
+  while (iter->status().ok() && iter->Valid()) {
+    ASSERT_GE(iter->key().compare(key_ranges[2]), 0);
+    ASSERT_LT(iter->key().compare(key_ranges[3]), 0);
+    count++;
+    iter->Next();
+  }
+  ASSERT_OK(iter->status()) << iter->status().ToString();
+  ASSERT_EQ(count, 2);
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

db/version_set.cc

@@ -1192,13 +1192,10 @@ class LevelIterator final : public InternalIterator {
         }
       }
     }
-    // Propagate io colaescing threshold
-    // TODO: This is error prone as we may forget to copy some fields. Think
-    // of a better way to do this.
+    // Propagate multiscan configs
     for (auto& file_to_arg : *file_to_scan_opts_) {
-      file_to_arg.second.io_coalesce_threshold = so->io_coalesce_threshold;
-      file_to_arg.second.max_prefetch_size = so->max_prefetch_size;
-      file_to_arg.second.use_async_io = so->use_async_io;
+      file_to_arg.second.CopyConfigFrom(*so);
+      file_to_arg.second.SetRequireFileOverlap(true);
     }
   }
 

include/rocksdb/options.h

@@ -1848,7 +1848,7 @@ class MultiScanArgs {
   operator std::vector<ScanOptions>*() { return &original_ranges_; }
 
   operator const std::vector<ScanOptions>*() const { return &original_ranges_; }
-  // Destructor
+
   ~MultiScanArgs() {}
 
   const std::vector<ScanOptions>& GetScanRanges() const {
@@ -1857,6 +1857,20 @@ class MultiScanArgs {
 
   const Comparator* GetComparator() const { return comp_; }
 
+  void SetRequireFileOverlap(bool require_overlap) {
+    require_file_overlap_ = require_overlap;
+  }
+
+  bool RequireFileOverlap() const { return require_file_overlap_; }
+
+  // Copies the configurations (excluding actual scan ranges) from another
+  // MultiScanArgs.
+  void CopyConfigFrom(const MultiScanArgs& other) {
+    io_coalesce_threshold = other.io_coalesce_threshold;
+    max_prefetch_size = other.max_prefetch_size;
+    use_async_io = other.use_async_io;
+  }
+
   uint64_t io_coalesce_threshold = 16 << 10;  // 16KB by default
 
   // Maximum size (in bytes) for the data blocks loaded by a MultiScan.
@@ -1880,6 +1894,11 @@ class MultiScanArgs {
   // The comparator used for ordering ranges
   const Comparator* comp_;
   std::vector<ScanOptions> original_ranges_;
+
+  // Internal use only.
+  // Fail the Prepare() on a file if a scan range does not overlap
+  // with the file.
+  bool require_file_overlap_{false};
 };
 
 // Options that control read operations

table/block_based/block_based_table_iterator.cc

@@ -963,6 +963,15 @@ BlockBasedTableIterator::MultiScanState::~MultiScanState() {
 // - After Prepare(), the iterator expects Seek to be called on the start key
 // of each ScanOption in order. If any other Seek is done, an error status is
 // returned
+// - Whenever all blocks of a scan opt are exhausted, the iterator will become
+// invalid and UpperBoundCheckResult() will return kOutOfBound. So that the
+// upper layer (LevelIterator) will stop scanning instead thinking EOF is
+// reached and continue into the next file. The only exception is for the last
+// scan opt. If we reach the end of the last scan opt, UpperBoundCheckResult()
+// will return kUnknown instead of kOutOfBound. This mechanism requires that
+// scan opts are properly pruned such that there is no scan opt that is after
+// this file's key range. This check can be enforeced by setting
+// MultiScanArgs::require_file_overlap to true.
 // 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
@@ -989,8 +998,9 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
   std::vector<BlockHandle> scan_block_handles;
   std::vector<std::tuple<size_t, size_t>> block_index_ranges_per_scan;
   const std::vector<ScanOptions>& scan_opts = multiscan_opts->GetScanRanges();
-  multi_scan_status_ = CollectBlockHandles(scan_opts, &scan_block_handles,
-                                           &block_index_ranges_per_scan);
+  multi_scan_status_ =
+      CollectBlockHandles(scan_opts, multiscan_opts->RequireFileOverlap(),
+                          &scan_block_handles, &block_index_ranges_per_scan);
   if (!multi_scan_status_.ok()) {
     return;
   }
@@ -1316,7 +1326,7 @@ Status BlockBasedTableIterator::ValidateScanOptions(
 }
 
 Status BlockBasedTableIterator::CollectBlockHandles(
-    const std::vector<ScanOptions>& scan_opts,
+    const std::vector<ScanOptions>& scan_opts, bool require_file_overlap,
     std::vector<BlockHandle>* scan_block_handles,
     std::vector<std::tuple<size_t, size_t>>* block_index_ranges_per_scan) {
   for (const auto& scan_opt : scan_opts) {
@@ -1368,12 +1378,14 @@ Status BlockBasedTableIterator::CollectBlockHandles(
       ++num_blocks;
     } else if (num_blocks == 0 && index_iter_->UpperBoundCheckResult() !=
                                       IterBoundCheck::kOutOfBound) {
-      // We should not have scan ranges that are completely after the file's
-      // range. This is important for FindBlockForwardInMultiScan() which only
-      // lets the upper layer (LevelIterator) advance to the next SST file when
-      // the last scan range is exhausted.
-      return Status::InvalidArgument("Scan does not intersect with file");
-      ;
+      // If require_file_overlap is set, then the scan ranges for this file
+      // must intersect with the file. Otherwise, allow empty intersection.
+      if (require_file_overlap) {
+        // This is important for FindBlockForwardInMultiScan() which only
+        // lets the upper layer (LevelIterator) advance to the next SST file
+        // when the last scan range is exhausted.
+        return Status::InvalidArgument("Scan does not intersect with file");
+      }
     }
     block_index_ranges_per_scan->emplace_back(
         scan_block_handles->size() - num_blocks, scan_block_handles->size());

table/block_based/block_based_table_iterator.h

@@ -669,7 +669,7 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
   Status ValidateScanOptions(const MultiScanArgs* multiscan_opts);
 
   Status CollectBlockHandles(
-      const std::vector<ScanOptions>& scan_opts,
+      const std::vector<ScanOptions>& scan_opts, bool require_file_overlap,
       std::vector<BlockHandle>* scan_block_handles,
       std::vector<std::tuple<size_t, size_t>>* block_index_ranges_per_scan);
 

table/block_based/block_based_table_reader_test.cc

@@ -1492,6 +1492,52 @@ TEST_P(BlockBasedTableReaderTest, MultiScanUnpinPreviousBlocks) {
   }
 }
 
+TEST_P(BlockBasedTableReaderTest, MultiScanOptFileOverlapChecking) {
+  std::vector<std::pair<std::string, std::string>> kv =
+      BlockBasedTableReaderBaseTest::GenerateKVMap(
+          20 /* num_block */,
+          true /* mixed_with_human_readable_string_value */);
+  std::vector<std::pair<std::string, std::string>> actual_kv(
+      kv.begin(), kv.begin() + 15 * kEntriesPerBlock);
+
+  std::string table_name = "BlockBasedTableReaderTest_UnpinPreviousBlocks" +
+                           CompressionTypeToString(compression_type_);
+  ImmutableOptions ioptions(options_);
+  CreateTable(table_name, ioptions, compression_type_, actual_kv,
+              compression_parallel_threads_, compression_dict_bytes_);
+
+  std::unique_ptr<BlockBasedTable> table;
+  FileOptions foptions;
+  foptions.use_direct_reads = use_direct_reads_;
+  InternalKeyComparator comparator(options_.comparator);
+  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table,
+                           true /* bool prefetch_index_and_filter_in_cache */,
+                           nullptr /* status */, persist_udt_);
+
+  ReadOptions read_opts;
+  std::unique_ptr<InternalIterator> iter;
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.SetRequireFileOverlap(false);
+  scan_options.insert(ExtractUserKey(kv[5 * kEntriesPerBlock].first),
+                      ExtractUserKey(kv[6 * kEntriesPerBlock].first));
+  scan_options.insert(ExtractUserKey(kv[16 * kEntriesPerBlock].first),
+                      ExtractUserKey(kv[17 * kEntriesPerBlock].first));
+
+  iter->Prepare(&scan_options);
+  ASSERT_OK(iter->status());
+
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+  scan_options.SetRequireFileOverlap(true);
+  iter->Prepare(&scan_options);
+  ASSERT_TRUE(iter->status().IsInvalidArgument());
+}
+
 // Param 1: compression type
 // Param 2: whether to use direct reads
 // Param 3: Block Based Table Index type, partitioned filters are also enabled