Add reverse order support to MultiScan (#14876)

Summary:
Pull Request resolved: #14876

Add reverse scan support to `MultiScanArgs`/`MultiScan` so callers can iterate bounded scan ranges in descending order. Plumb reverse scan state through `DBIter`, `BlockBasedTableIterator`, and `MultiScanIndexIterator`, including reverse prefetch windowing and async I/O support. Update `db_bench` `multiscanrandom` to honor `--reverse_iterator` and report rows scanned so reverse MultiScan and iterator baselines can be compared on row throughput.

Add `db_stress --multiscan_reverse` coverage so the prepared MultiScan path can be stressed with `SeekForPrev()`/`Prev()` against a control iterator.

Also fix the table child iterator path to accept reverse `SeekForPrev()` targets above a file's prepared ranges. `MergingIterator` and `LevelIterator` can legally seek every child to the current range limit; a child file may only have lower prepared ranges, so treating that as invalid caused intermittent `SeekForPrev target is outside prepared ranges` exceptions in the benchmark.

Reviewed By: xingbowang

Differential Revision: D109357730

fbshipit-source-id: eccece5ca7d8f5d6760d51ebe76801a332a49061

Author: jaykorean

db/db_iter.cc

@@ -1921,6 +1921,9 @@ Status DBIter::ValidateScanOptions(const MultiScanArgs& multiscan_opts) const {
 
   const std::vector<ScanOptions>& scan_opts = multiscan_opts.GetScanRanges();
   const bool has_limit = scan_opts.front().range.limit.has_value();
+  if (multiscan_opts.reverse && !has_limit) {
+    return Status::InvalidArgument("Reverse MultiScan requires limit");
+  }
   if (!has_limit && scan_opts.size() > 1) {
     return Status::InvalidArgument("Scan has no upper bound");
   }
@@ -2004,6 +2007,12 @@ void DBIter::Seek(const Slice& target) {
   StopWatch sw(clock_, statistics_, DB_SEEK);
 
   if (scan_opts_.has_value()) {
+    if (scan_opts_->reverse) {
+      status_ = Status::InvalidArgument("Seek called on reverse MultiScan");
+      valid_ = false;
+      return;
+    }
+
     // 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()) {
@@ -2110,6 +2119,61 @@ void DBIter::SeekForPrev(const Slice& target) {
   PERF_CPU_TIMER_GUARD(iter_seek_cpu_nanos, clock_);
   StopWatch sw(clock_, statistics_, DB_SEEK);
 
+  if (scan_opts_.has_value()) {
+    if (!scan_opts_->reverse) {
+      status_ =
+          Status::InvalidArgument("SeekForPrev called on forward MultiScan");
+      valid_ = false;
+      return;
+    }
+
+    auto const& scan_ranges = scan_opts_.value().GetScanRanges();
+    if (scan_index_ >= scan_ranges.size()) {
+      status_ = Status::InvalidArgument(
+          "SeekForPrev called after exhausting all of the scan ranges");
+      valid_ = false;
+      return;
+    }
+
+    auto const& range = scan_ranges[scan_index_];
+    auto const& limit = range.range.limit;
+    assert(limit.has_value());
+    if (!limit.has_value() ||
+        user_comparator_.CompareWithoutTimestamp(target, *limit) != 0) {
+      status_ = Status::InvalidArgument(
+          "SeekForPrev target does not match the limit of the next prepared "
+          "range at index " +
+          std::to_string(scan_index_));
+      valid_ = false;
+      return;
+    }
+
+    auto const& start = range.range.start;
+    assert(start.has_value());
+    if (iterate_lower_bound_ == nullptr ||
+        user_comparator_.CompareWithoutTimestamp(start.value(),
+                                                 *iterate_lower_bound_) != 0) {
+      status_ = Status::InvalidArgument(
+          "Lower bound is not set to the same start value of the next "
+          "prepared range at index " +
+          std::to_string(scan_index_));
+      valid_ = false;
+      return;
+    }
+    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 (has_trace_state_) {
     // TODO: What do we do if this returns an error?
     Slice lower_bound, upper_bound;

db/db_iterator_test.cc

@@ -5418,6 +5418,98 @@ TEST_P(DBMultiScanIteratorTest, AsyncPrefetchAcrossMultipleFiles) {
   iter.reset();
 }
 
+TEST_P(DBMultiScanIteratorTest, ReversePrefetchAcrossMultipleRanges) {
+  auto options = CurrentOptions();
+  options.target_file_size_base = 1 << 15;  // 32KiB
+  options.compaction_style = kCompactionStyleUniversal;
+  options.num_levels = 50;
+  options.compression = kNoCompression;
+  DestroyAndReopen(options);
+
+  Random rnd(303);
+  for (int i = 0; i < 1000; ++i) {
+    ASSERT_OK(Put(Key(i), rnd.RandomString(1 << 10)));
+  }
+  ASSERT_OK(Flush());
+  ASSERT_OK(db_->CompactRange({}, nullptr, nullptr));
+  ASSERT_GT(NumTableFilesAtLevel(49), 3);
+
+  ReadOptions ro;
+  ro.fill_cache = GetParam();
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+  auto tracking_dispatcher = std::make_shared<TrackingIODispatcher>();
+
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.use_async_io = false;
+  scan_options.reverse = true;
+  scan_options.io_dispatcher = tracking_dispatcher;
+  std::vector<std::string> key_ranges(
+      {Key(100), Key(200), Key(500), Key(600), Key(800), Key(900)});
+  scan_options.insert(key_ranges[0], key_ranges[1]);
+  scan_options.insert(key_ranges[2], key_ranges[3]);
+  scan_options.insert(key_ranges[4], key_ranges[5]);
+
+  std::unique_ptr<MultiScan> iter =
+      dbfull()->NewMultiScan(ro, cfh, scan_options);
+  ASSERT_NE(iter, nullptr);
+
+  std::vector<std::string> actual_keys;
+  try {
+    for (auto range : *iter) {
+      for (auto it : range) {
+        actual_keys.push_back(it.first.ToString());
+      }
+    }
+  } catch (MultiScanException& ex) {
+    FAIL() << "Iterator returned status " << ex.what();
+  } catch (std::logic_error& ex) {
+    FAIL() << "Iterator returned logic error " << ex.what();
+  }
+
+  std::vector<std::string> expected_keys;
+  for (int i = 199; i >= 100; --i) {
+    expected_keys.push_back(Key(i));
+  }
+  for (int i = 599; i >= 500; --i) {
+    expected_keys.push_back(Key(i));
+  }
+  for (int i = 899; i >= 800; --i) {
+    expected_keys.push_back(Key(i));
+  }
+  ASSERT_EQ(actual_keys, expected_keys);
+  ASSERT_GT(tracking_dispatcher->GetReadSets().size(), 0);
+}
+
+TEST_P(DBMultiScanIteratorTest, ReverseRequiresLimits) {
+  auto options = CurrentOptions();
+  options.compression = kNoCompression;
+  DestroyAndReopen(options);
+
+  ASSERT_OK(Put("a", "va"));
+  ASSERT_OK(Put("b", "vb"));
+  ASSERT_OK(Flush());
+
+  ReadOptions ro;
+  ro.fill_cache = GetParam();
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.reverse = true;
+  scan_options.insert("a");
+  scan_options.insert("b", "c");
+
+  std::unique_ptr<MultiScan> iter =
+      dbfull()->NewMultiScan(ro, cfh, scan_options);
+  ASSERT_NE(iter, nullptr);
+  try {
+    (void)iter->begin();
+    FAIL() << "Expected reverse MultiScan without an upper bound to fail";
+  } catch (MultiScanException& ex) {
+    ASSERT_NOK(ex.status());
+    ASSERT_TRUE(ex.status().IsInvalidArgument());
+  }
+}
+
 // Wrapper filesystem that does not support async IO.
 // Used to verify that MultiScan gracefully falls back to sync IO.
 class NoAsyncIOFS : public FileSystemWrapper {

db/multi_scan.cc

@@ -3,14 +3,56 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
+#include <string>
+#include <vector>
+
 #include "file/file_util.h"
 #include "rocksdb/db.h"
 #include "rocksdb/file_system.h"
+#include "rocksdb/iterator.h"
 
 namespace ROCKSDB_NAMESPACE {
 
 using MultiScanIterator = MultiScan::MultiScanIterator;
 
+namespace {
+
+Status ValidateScanOptionsForMultiScan(const MultiScanArgs& multiscan_opts) {
+  if (multiscan_opts.empty()) {
+    return Status::InvalidArgument("Empty MultiScanArgs");
+  }
+
+  const Comparator* comparator = multiscan_opts.GetComparator();
+  if (comparator == nullptr) {
+    return Status::InvalidArgument("MultiScanArgs requires comparator");
+  }
+
+  const std::vector<ScanOptions>& scan_opts = multiscan_opts.GetScanRanges();
+  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 (multiscan_opts.reverse && !scan_range.limit.has_value()) {
+      return Status::InvalidArgument(
+          "Reverse MultiScan requires limit at index " + std::to_string(i));
+    }
+    if (scan_range.limit.has_value() &&
+        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));
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace
+
 MultiScan::MultiScan(const ReadOptions& read_options,
                      const MultiScanArgs& scan_opts, DB* db,
                      ColumnFamilyHandle* cfh)
@@ -28,14 +70,28 @@ MultiScan::MultiScan(const ReadOptions& read_options,
       }()),
       db_(db),
       cfh_(cfh) {
+  Status validate_status = ValidateScanOptionsForMultiScan(scan_opts_);
+  if (!validate_status.ok()) {
+    db_iter_.reset(NewErrorIterator(validate_status));
+    return;
+  }
+
   bool slow_path = false;
-  // Setup read_options with iterate_uuper_bound based on the first scan.
-  // Subsequent scans will update and allocate a new DB iterator as necessary
-  if (scan_opts_.GetScanRanges()[0].range.limit) {
-    upper_bound_ = *scan_opts_.GetScanRanges()[0].range.limit;
-    read_options_.iterate_upper_bound = &upper_bound_;
-  } else {
-    read_options_.iterate_upper_bound = nullptr;
+  // Setup read_options bounds based on the first scan. Subsequent scans will
+  // update the bounds and allocate a new DB iterator as necessary.
+  {
+    const ScanOptions& first_scan = scan_opts_.GetScanRanges()[0];
+    if (scan_opts_.reverse) {
+      lower_bound_ = *first_scan.range.start;
+      read_options_.iterate_lower_bound = &lower_bound_;
+      upper_bound_ = *first_scan.range.limit;
+      read_options_.iterate_upper_bound = &upper_bound_;
+    } else if (first_scan.range.limit) {
+      upper_bound_ = *first_scan.range.limit;
+      read_options_.iterate_upper_bound = &upper_bound_;
+    } else {
+      read_options_.iterate_upper_bound = nullptr;
+    }
   }
   for (const auto& opts : scan_opts_.GetScanRanges()) {
     // Check that all the ScanOptions either specify an upper bound or not. If
@@ -54,36 +110,55 @@ MultiScan::MultiScan(const ReadOptions& read_options,
   }
 }
 
+void MultiScanIterator::SeekCurrentRange() {
+  const ScanOptions& scan_opt = scan_opts_[idx_];
+  if (scan_opt.range.limit) {
+    *upper_bound_ = *scan_opt.range.limit;
+    read_options_.iterate_upper_bound = upper_bound_;
+  } else {
+    read_options_.iterate_upper_bound = nullptr;
+  }
+
+  if (reverse_) {
+    *lower_bound_ = *scan_opt.range.start;
+    read_options_.iterate_lower_bound = lower_bound_;
+    assert(scan_opt.range.limit.has_value());
+    db_iter_->SeekForPrev(*scan_opt.range.limit);
+  } else {
+    db_iter_->Seek(*scan_opt.range.start);
+  }
+}
+
 MultiScanIterator& MultiScanIterator::operator++() {
   status_ = db_iter_->status();
   if (!status_.ok()) {
     throw MultiScanException(status_);
   }
 
+  if (!valid_) {
+    throw std::logic_error("Incrementing exhausted MultiScan iterator");
+  }
+  ++idx_;
   if (idx_ >= scan_opts_.size()) {
-    throw std::logic_error("Index out of range");
+    valid_ = false;
+    return *this;
   }
-  idx_++;
-  if (idx_ < scan_opts_.size()) {
-    // Check if we need to update read_options_
-    if (scan_opts_[idx_].range.limit.has_value() !=
-        (read_options_.iterate_upper_bound != nullptr)) {
-      if (scan_opts_[idx_].range.limit) {
-        *upper_bound_ = *scan_opts_[idx_].range.limit;
-        read_options_.iterate_upper_bound = upper_bound_;
-      } else {
-        read_options_.iterate_upper_bound = nullptr;
-      }
-      db_iter_.reset(db_->NewIterator(read_options_, cfh_));
-      scan_.Reset(db_iter_.get());
-    } else if (scan_opts_[idx_].range.limit) {
-      *upper_bound_ = *scan_opts_[idx_].range.limit;
-    }
-    db_iter_->Seek(*scan_opts_[idx_].range.start);
-    status_ = db_iter_->status();
-    if (!status_.ok()) {
-      throw MultiScanException(status_);
+
+  // Check if we need to update read_options_
+  if (scan_opts_[idx_].range.limit.has_value() !=
+      (read_options_.iterate_upper_bound != nullptr)) {
+    if (scan_opts_[idx_].range.limit) {
+      read_options_.iterate_upper_bound = upper_bound_;
+    } else {
+      read_options_.iterate_upper_bound = nullptr;
     }
+    db_iter_.reset(db_->NewIterator(read_options_, cfh_));
+    scan_.Reset(db_iter_.get(), reverse_);
+  }
+  SeekCurrentRange();
+  status_ = db_iter_->status();
+  if (!status_.ok()) {
+    throw MultiScanException(status_);
   }
   return *this;
 }

db_stress_tool/db_stress_common.h

@@ -470,6 +470,7 @@ DECLARE_uint32(min_tombstones_for_range_conversion);
 DECLARE_uint32(ingest_wbwi_one_in);
 DECLARE_bool(universal_reduce_file_locking);
 DECLARE_bool(use_multiscan);
+DECLARE_bool(multiscan_reverse);
 DECLARE_bool(multiscan_use_async_io);
 DECLARE_bool(read_scoped_block_buffer_provider);
 DECLARE_uint64(multiscan_max_prefetch_memory_bytes);

db_stress_tool/db_stress_gflags.cc

@@ -1723,6 +1723,10 @@ DEFINE_bool(
 DEFINE_bool(use_multiscan, false,
             "If set, use the batched MultiScan API for scans.");
 
+DEFINE_bool(multiscan_reverse, false,
+            "If set with use_multiscan, scan each MultiScan range in reverse "
+            "using SeekForPrev and Prev.");
+
 DEFINE_bool(multiscan_use_async_io, false,
             "If set, enable async_io for MultiScan operations.");