Fix prefix_same_as_start iterator correctness gaps (#15606)

db/db_iter.cc

@@ -114,10 +114,13 @@ DBIter::DBIter(Env* _env, const ReadOptions& read_options,
       is_key_seqnum_zero_(false),
       prefix_same_as_start_(
           prefix_extractor_ ? read_options.prefix_same_as_start : false),
+      strict_prefix_same_as_start_(prefix_same_as_start_ &&
+                                   ioptions.prefix_seek_opt_in_only),
+      unbounded_prefix_scan_(false),
       pin_thru_lifetime_(read_options.pin_data),
-      expect_total_order_inner_iter_(prefix_extractor_ == nullptr ||
-                                     read_options.total_order_seek ||
-                                     read_options.auto_prefix_mode),
+      expect_total_order_inner_iter_(
+          GetIteratorPrefixMode(read_options, prefix_extractor_ != nullptr) !=
+          IteratorPrefixMode::kPrefix),
       // Read-path range conversion assumes the scan can observe all interior
       // live keys. table_filter can hide whole SSTs, read tiers other than
       // kReadAllTier can skip SSTs or uncached blocks, and timestamp filtering
@@ -1779,7 +1782,8 @@ void DBIter::FlushPendingTombstoneRun(const Slice& end_key,
 }
 
 void DBIter::MaybeInsertRangeTombstone(const Slice& end_key) {
-  if (contiguous_tombstone_count_ < min_tombstones_for_range_conversion_) {
+  if (contiguous_tombstone_count_ < min_tombstones_for_range_conversion_ ||
+      unbounded_prefix_scan_) {
     return;
   }
 
@@ -2031,6 +2035,9 @@ 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 (!CheckSeekTargetPrefix(target)) {
+    return;
+  }
 
   if (scan_opts_.has_value()) {
     if (scan_opts_->reverse) {
@@ -2147,6 +2154,9 @@ void DBIter::SeekForPrev(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 (!CheckSeekTargetPrefix(target)) {
+    return;
+  }
 
   if (scan_opts_.has_value()) {
     if (!scan_opts_->reverse) {
@@ -2267,11 +2277,35 @@ void DBIter::SeekForPrev(const Slice& target) {
   }
 }
 
+bool DBIter::CheckSeekTargetPrefix(const Slice& target) {
+  unbounded_prefix_scan_ = false;
+  if (expect_total_order_inner_iter_ || prefix_extractor_->InDomain(target)) {
+    return true;
+  }
+  if (strict_prefix_same_as_start_) {
+    valid_ = false;
+    status_ = Status::InvalidArgument(
+        "Seek target outside the prefix extractor's domain is not supported "
+        "with prefix_same_as_start under prefix_seek_opt_in_only");
+    return false;
+  }
+  unbounded_prefix_scan_ = true;
+  return true;
+}
+
 void DBIter::SeekToFirst() {
+  if (strict_prefix_same_as_start_) {
+    valid_ = false;
+    status_ = Status::NotSupported(
+        "SeekToFirst() is not supported with prefix_same_as_start under "
+        "prefix_seek_opt_in_only; Seek() to a key in the wanted prefix");
+    return;
+  }
   if (iterate_lower_bound_ != nullptr) {
     Seek(*iterate_lower_bound_);
     return;
   }
+  unbounded_prefix_scan_ = !expect_total_order_inner_iter_;
   PERF_COUNTER_ADD(iter_seek_count, 1);
   PERF_CPU_TIMER_GUARD(iter_seek_cpu_nanos, clock_);
   // Don't use iter_::Seek() if we set a prefix extractor
@@ -2322,6 +2356,14 @@ void DBIter::SeekToFirst() {
 }
 
 void DBIter::SeekToLast() {
+  if (strict_prefix_same_as_start_) {
+    valid_ = false;
+    status_ = Status::NotSupported(
+        "SeekToLast() is not supported with prefix_same_as_start under "
+        "prefix_seek_opt_in_only; SeekForPrev() to a key in the wanted "
+        "prefix");
+    return;
+  }
   if (iterate_upper_bound_ != nullptr) {
     // Seek to last key strictly less than ReadOptions.iterate_upper_bound.
     SeekForPrev(*iterate_upper_bound_);
@@ -2336,6 +2378,7 @@ void DBIter::SeekToLast() {
 #endif
     return;
   }
+  unbounded_prefix_scan_ = !expect_total_order_inner_iter_;
 
   PERF_COUNTER_ADD(iter_seek_count, 1);
   PERF_CPU_TIMER_GUARD(iter_seek_cpu_nanos, clock_);

db/db_iter.h

@@ -612,6 +612,12 @@ class DBIter final : public Iterator {
            prefix_extractor_->InDomain(target);
   }
 
+  // Called at the start of Seek() and SeekForPrev(). Returns false, with
+  // status_ set, if target is outside the prefix extractor's domain under
+  // strict_prefix_same_as_start_. Otherwise sets unbounded_prefix_scan_ when
+  // target gives no prefix to bound a prefix seek.
+  bool CheckSeekTargetPrefix(const Slice& target);
+
   void ResetSeekState() {
     ReleaseTempPinnedData();
     ResetBlobData();
@@ -738,6 +744,19 @@ class DBIter final : public Iterator {
   // user key.
   bool is_key_seqnum_zero_;
   const bool prefix_same_as_start_;
+  // prefix_same_as_start_ under DBOptions::prefix_seek_opt_in_only. Every scan
+  // must take its prefix from an in-domain seek key: SeekToFirst() and
+  // SeekToLast() set NotSupported, and Seek() and SeekForPrev() with a target
+  // outside the prefix extractor's domain set InvalidArgument.
+  const bool strict_prefix_same_as_start_;
+  // True while the current scan uses prefix seek in the child iterators but
+  // has no prefix to bound it: after SeekToFirst() without
+  // iterate_lower_bound, SeekToLast() without iterate_upper_bound, or a
+  // Seek()/SeekForPrev() target outside the prefix extractor's domain. Prefix
+  // seek in the child iterators can then skip live keys, so read-path range
+  // tombstone conversion is disabled so that a skipped live key cannot be
+  // covered by a new tombstone.
+  bool unbounded_prefix_scan_;
   // Means that we will pin all data blocks we read as long the Iterator
   // is not deleted, will be true if ReadOptions::pin_data is true
   const bool pin_thru_lifetime_;

db/db_iterator_test.cc

@@ -87,6 +87,147 @@ class DBIteratorBaseTest : public DBTestBase {
 };
 
 #ifdef OS_LINUX
+TEST_F(DBIteratorBaseTest, PrefixSameAsStartDirectionChange) {
+  // prefix_same_as_start uses prefix seek even with total_order_seek=true (as
+  // forced by prefix_seek_opt_in_only) or auto_prefix_mode. On Next() after
+  // Prev(), the inner iterator stands on a key with another prefix (bbbb1).
+  // DBIter must reseek it to the current key; moving it without a reseek
+  // seeks the other children to bbbb1, where a prefix filter drops the child
+  // holding kkkk1 and kkkk3.
+  struct ReadCase {
+    const char* name;
+    bool opt_in;
+    bool total_order_seek;
+    bool auto_prefix_mode;
+  };
+  const std::vector<ReadCase> read_cases = {
+      {"prefix_same_as_start", false, false, false},
+      {"+total_order_seek", false, true, false},
+      {"+auto_prefix_mode", false, false, true},
+      {"+prefix_seek_opt_in_only", true, false, false},
+  };
+  for (bool memtable_bloom : {false, true}) {
+    for (const ReadCase& rc : read_cases) {
+      SCOPED_TRACE(std::string(rc.name) +
+                   (memtable_bloom ? " memtable_bloom" : " sst_filter"));
+      Options options = CurrentOptions();
+      options.disable_auto_compactions = true;
+      options.prefix_extractor.reset(NewFixedPrefixTransform(4));
+      options.prefix_seek_opt_in_only = rc.opt_in;
+      BlockBasedTableOptions table_options;
+      if (memtable_bloom) {
+        options.memtable_prefix_bloom_size_ratio = 0.1;
+      } else {
+        table_options.filter_policy.reset(NewBloomFilterPolicy(10, false));
+        table_options.whole_key_filtering = false;
+      }
+      options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+      DestroyAndReopen(options);
+      // The child that gets filtered has prefix kkkk only: the SST for the
+      // SST filter case, the memtable for the memtable bloom case.
+      const std::vector<std::string> kkkk_only = {"kkkk1", "kkkk3"};
+      const std::vector<std::string> mixed = {"bbbb1", "kkkk2"};
+      for (const std::string& key : memtable_bloom ? mixed : kkkk_only) {
+        ASSERT_OK(Put(key, "v"));
+      }
+      ASSERT_OK(Flush());
+      for (const std::string& key : memtable_bloom ? kkkk_only : mixed) {
+        ASSERT_OK(Put(key, "v"));
+      }
+
+      ReadOptions ro;
+      ro.prefix_same_as_start = true;
+      ro.total_order_seek = rc.total_order_seek;
+      ro.auto_prefix_mode = rc.auto_prefix_mode;
+      std::unique_ptr<Iterator> iter(db_->NewIterator(ro));
+      iter->Seek("kkkk1");
+      std::vector<std::string> keys;
+      // Seek, Next, Prev, then Next until the end of the prefix.
+      for (int step = 0; iter->Valid(); ++step) {
+        keys.push_back(iter->key().ToString());
+        if (step == 1) {
+          iter->Prev();
+        } else {
+          iter->Next();
+        }
+      }
+      ASSERT_OK(iter->status());
+      const std::vector<std::string> expected = {"kkkk1", "kkkk2", "kkkk1",
+                                                 "kkkk2", "kkkk3"};
+      EXPECT_EQ(expected, keys);
+    }
+  }
+}
+
+TEST_F(DBIteratorBaseTest, PrefixSameAsStartWithExplicitBounds) {
+  // prefix_same_as_start implies a bound at the end of the seek prefix. With
+  // or without auto_prefix_mode, the iterator returns only keys within both
+  // that implied bound and the explicit iterate_lower_bound /
+  // iterate_upper_bound, whichever is narrower.
+  Options options = CurrentOptions();
+  options.disable_auto_compactions = true;
+  options.prefix_extractor.reset(NewFixedPrefixTransform(4));
+  BlockBasedTableOptions table_options;
+  table_options.filter_policy.reset(NewBloomFilterPolicy(10, false));
+  table_options.whole_key_filtering = false;
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+  DestroyAndReopen(options);
+  for (const char* key : {"kkkj9", "kkkk1", "kkkk5"}) {
+    ASSERT_OK(Put(key, "v"));
+  }
+  ASSERT_OK(Flush());
+  for (const char* key : {"kkkk3", "kkkl1"}) {
+    ASSERT_OK(Put(key, "v"));
+  }
+
+  struct BoundCase {
+    const char* name;
+    const char* bound;
+    std::vector<std::string> expected;
+  };
+  const std::vector<BoundCase> forward_cases = {
+      {"upper bound inside prefix", "kkkk4", {"kkkk1", "kkkk3"}},
+      {"upper bound past prefix", "kkkm", {"kkkk1", "kkkk3", "kkkk5"}},
+  };
+  const std::vector<BoundCase> backward_cases = {
+      {"lower bound inside prefix", "kkkk2", {"kkkk5", "kkkk3"}},
+      {"lower bound before prefix", "kkkj", {"kkkk5", "kkkk3", "kkkk1"}},
+  };
+  for (bool auto_prefix_mode : {false, true}) {
+    ReadOptions ro;
+    ro.prefix_same_as_start = true;
+    ro.auto_prefix_mode = auto_prefix_mode;
+    for (const BoundCase& bc : forward_cases) {
+      SCOPED_TRACE(std::string(bc.name) +
+                   (auto_prefix_mode ? " auto_prefix_mode" : ""));
+      Slice upper_bound(bc.bound);
+      ro.iterate_upper_bound = &upper_bound;
+      std::unique_ptr<Iterator> iter(db_->NewIterator(ro));
+      std::vector<std::string> keys;
+      for (iter->Seek("kkkk1"); iter->Valid(); iter->Next()) {
+        keys.push_back(iter->key().ToString());
+      }
+      ASSERT_OK(iter->status());
+      EXPECT_EQ(bc.expected, keys);
+    }
+    ro.iterate_upper_bound = nullptr;
+    for (const BoundCase& bc : backward_cases) {
+      SCOPED_TRACE(std::string(bc.name) +
+                   (auto_prefix_mode ? " auto_prefix_mode" : ""));
+      Slice lower_bound(bc.bound);
+      ro.iterate_lower_bound = &lower_bound;
+      std::unique_ptr<Iterator> iter(db_->NewIterator(ro));
+      std::vector<std::string> keys;
+      for (iter->SeekForPrev("kkkk5"); iter->Valid(); iter->Prev()) {
+        keys.push_back(iter->key().ToString());
+      }
+      ASSERT_OK(iter->status());
+      EXPECT_EQ(bc.expected, keys);
+    }
+    ro.iterate_lower_bound = nullptr;
+  }
+}
+
 TEST_F(DBIteratorBaseTest, PrevChecksQueryAbort) {
   ASSERT_OK(Put("a", "va"));
   ASSERT_OK(Put("b", "vb"));
@@ -8101,10 +8242,10 @@ TEST_P(ReadPathRangeTombstoneTest,
 
 TEST_P(ReadPathRangeTombstoneTest, PrefixFilterOutOfDomainSeek) {
   // With an out-of-domain seek target, prefix_same_as_start cannot establish
-  // a seek-prefix bound. The iterator therefore behaves like an unrestricted
-  // scan while still avoiding Transform() on the short seek target, so the
-  // valid run [bbbb, bbdd) should be converted even though it spans multiple
-  // distinct prefixes.
+  // a seek-prefix bound. The iterator scans past the seek prefix while the
+  // child iterators still use prefix seek, which can skip live keys, so the
+  // run [bbbb, bbdd) must not be converted. The scan must still avoid
+  // Transform() on the short seek target.
   Options options = CurrentOptions();
   options.min_tombstones_for_range_conversion = 2;
   options.statistics = CreateDBStatistics();
@@ -8144,8 +8285,7 @@ TEST_P(ReadPathRangeTombstoneTest, PrefixFilterOutOfDomainSeek) {
     }
   }
   ASSERT_OK(it->status());
-  ASSERT_EQ(attempted_insert_ranges_.size(), 1u);
-  AssertRange(0, "bbbb", "bbdd");
+  ASSERT_TRUE(attempted_insert_ranges_.empty());
 }
 
 TEST_P(ReadPathRangeTombstoneTest, TableFilterNotAllowed) {