Move the MultiScan seek key check to upper layer (facebook#14040)

Summary:
The current seek key validation is too strict. This change relaxes it at block iterator level, and add additional check at DB iterator level. The new contract is that when MultiScan is used, after prepared is called, each following seek must seek the start key of the prepared scan range in order. Otherwise, the iterator is set with error status.

Pull Request resolved: facebook#14040

Test Plan: Unit test

Reviewed By: anand1976

Differential Revision: D84292297

Pulled By: xingbowang

fbshipit-source-id: 7b31f727e67e7c0bfc53c2f9a6552e0c3d324869

Author: Xingbo Wang

db/db_iter.cc

@@ -1565,11 +1565,115 @@ void DBIter::SetSavedKeyToSeekForPrevTarget(const Slice& target) {
   }
 }
 
+Status DBIter::ValidateScanOptions(const MultiScanArgs& multiscan_opts) const {
+  if (multiscan_opts.empty()) {
+    return Status::InvalidArgument("Empty MultiScanArgs");
+  }
+
+  const std::vector<ScanOptions>& scan_opts = multiscan_opts.GetScanRanges();
+  const bool has_limit = scan_opts.front().range.limit.has_value();
+  if (!has_limit && scan_opts.size() > 1) {
+    return Status::InvalidArgument("Scan has no upper bound");
+  }
+
+  for (size_t i = 0; i < scan_opts.size(); ++i) {
+    const auto& scan_range = scan_opts[i].range;
+    if (!scan_range.start.has_value()) {
+      return Status::InvalidArgument("Scan has no start key at index " +
+                                     std::to_string(i));
+    }
+
+    if (scan_range.limit.has_value()) {
+      if (user_comparator_.CompareWithoutTimestamp(
+              scan_range.start.value(), /*a_has_ts=*/false,
+              scan_range.limit.value(), /*b_has_ts=*/false) >= 0) {
+        return Status::InvalidArgument(
+            "Scan start key is large or equal than limit at index " +
+            std::to_string(i));
+      }
+    }
+
+    if (i > 0) {
+      if (!scan_range.limit.has_value()) {
+        // multiple scan without limit scan ranges
+        return Status::InvalidArgument("Scan has no upper bound at index " +
+                                       std::to_string(i));
+      }
+
+      const auto& last_end_key = scan_opts[i - 1].range.limit.value();
+      if (user_comparator_.CompareWithoutTimestamp(
+              scan_range.start.value(), /*a_has_ts=*/false, last_end_key,
+              /*b_has_ts=*/false) < 0) {
+        return Status::InvalidArgument("Overlapping ranges at index " +
+                                       std::to_string(i));
+      }
+    }
+  }
+  return Status::OK();
+}
+
+void DBIter::Prepare(const MultiScanArgs& scan_opts) {
+  status_ = ValidateScanOptions(scan_opts);
+  if (!status_.ok()) {
+    return;
+  }
+  std::optional<MultiScanArgs> new_scan_opts;
+  new_scan_opts.emplace(scan_opts);
+  scan_opts_.swap(new_scan_opts);
+  scan_index_ = 0;
+  if (!scan_opts.empty()) {
+    iter_.Prepare(&scan_opts_.value());
+  } else {
+    iter_.Prepare(nullptr);
+  }
+}
+
 void DBIter::Seek(const Slice& target) {
   PERF_COUNTER_ADD(iter_seek_count, 1);
   PERF_CPU_TIMER_GUARD(iter_seek_cpu_nanos, clock_);
   StopWatch sw(clock_, statistics_, DB_SEEK);
 
+  if (scan_opts_.has_value()) {
+    // Validate the seek target is as expected in the previously prepared range
+    auto const& scan_ranges = scan_opts_.value().GetScanRanges();
+    if (scan_index_ >= scan_ranges.size()) {
+      status_ = Status::InvalidArgument(
+          "Seek called after exhausting all of the scan ranges");
+      valid_ = false;
+      return;
+    }
+
+    // Validate start key of next prepare range matches the seek target
+    auto const& range = scan_ranges[scan_index_];
+    auto const& start = range.range.start;
+    assert(start.has_value());
+    if (user_comparator_.CompareWithoutTimestamp(target, *start) != 0) {
+      status_ = Status::InvalidArgument(
+          "Seek target does not match the start of the next prepared range at "
+          "index " +
+          std::to_string(scan_index_));
+      valid_ = false;
+      return;
+    }
+
+    // validate the upper bound is set to the same value of limit, if limit
+    // exists
+    auto const& limit = range.range.limit;
+    if (limit.has_value()) {
+      if (iterate_upper_bound_ == nullptr ||
+          user_comparator_.CompareWithoutTimestamp(
+              limit.value(), *iterate_upper_bound_) != 0) {
+        status_ = Status::InvalidArgument(
+            "Upper bound is not set to the same limit value of the next "
+            "prepared range at index " +
+            std::to_string(scan_index_));
+        valid_ = false;
+        return;
+      }
+    }
+    scan_index_++;
+  }
+
   if (cfh_ != nullptr) {
     // TODO: What do we do if this returns an error?
     Slice lower_bound, upper_bound;

db/db_iter.h

@@ -240,16 +240,8 @@ class DBIter final : public Iterator {
 
   bool PrepareValue() override;
 
-  void Prepare(const MultiScanArgs& scan_opts) override {
-    std::optional<MultiScanArgs> new_scan_opts;
-    new_scan_opts.emplace(scan_opts);
-    scan_opts_.swap(new_scan_opts);
-    if (!scan_opts.empty()) {
-      iter_.Prepare(&scan_opts_.value());
-    } else {
-      iter_.Prepare(nullptr);
-    }
-  }
+  void Prepare(const MultiScanArgs& scan_opts) override;
+  Status ValidateScanOptions(const MultiScanArgs& multiscan_opts) const;
 
  private:
   DBIter(Env* _env, const ReadOptions& read_options,
@@ -506,6 +498,7 @@ class DBIter final : public Iterator {
   const size_t timestamp_size_;
   std::string saved_timestamp_;
   std::optional<MultiScanArgs> scan_opts_;
+  size_t scan_index_{0};
   ReadOnlyMemTable* const active_mem_;
   SequenceNumber memtable_seqno_lb_;
   uint32_t memtable_op_scan_flush_trigger_;

db/db_iterator_test.cc

@@ -4382,18 +4382,6 @@ TEST_P(DBMultiScanIteratorTest, FailureTest) {
   iter->Seek(key_ranges[2]);
   ASSERT_NOK(iter->status());
   iter.reset();
-
-  // Test the case of overlapping ranges
-  iter.reset(dbfull()->NewIterator(ro, cfh));
-  ASSERT_NE(iter, nullptr);
-  (*scan_options).clear();
-  scan_options.insert(key_ranges[0]);
-  scan_options.insert(key_ranges[2], key_ranges[3]);
-  iter->Prepare(scan_options);
-  ub = key_ranges[3];
-  iter->Seek(key_ranges[2]);
-  ASSERT_NOK(iter->status());
-  iter.reset();
 }
 
 TEST_P(DBMultiScanIteratorTest, RangeBetweenFiles) {

db/multi_scan.cc

@@ -40,6 +40,11 @@ MultiScan::MultiScan(const ReadOptions& read_options,
 }
 
 MultiScanIterator& MultiScanIterator::operator++() {
+  status_ = db_iter_->status();
+  if (!status_.ok()) {
+    throw MultiScanException(status_);
+  }
+
   if (idx_ >= scan_opts_.size()) {
     throw std::logic_error("Index out of range");
   }

db/version_set.cc

@@ -1668,8 +1668,6 @@ void LevelIterator::SetFileIterator(InternalIterator* iter) {
     if (FileHasMultiScanArg(file_index_)) {
       const MultiScanArgs& new_opts = GetMultiScanArgForFile(file_index_);
       file_iter_.Prepare(&new_opts);
-    } else {
-      file_iter_.Prepare(scan_opts_);
     }
   }
 

include/rocksdb/multi_scan.h

@@ -191,6 +191,10 @@ class MultiScan {
       if (scan_opts_.empty()) {
         throw std::logic_error("Zero scans in multi-scan");
       }
+      status_ = db_iter_->status();
+      if (!status_.ok()) {
+        throw MultiScanException(status_);
+      }
       db_iter_->Seek(*scan_opts_[idx_].range.start);
       status_ = db_iter_->status();
       if (!status_.ok()) {