From 1b48ecc2c6f319e590e3e4cd0e8adf863ef92236 Mon Sep 17 00:00:00 2001
From: Changyu Bi <changyubi@meta.com>
Date: Tue, 21 Feb 2023 11:57:58 -0800
Subject: [PATCH] Fix an assertion failure in DBIter::SeekToLast() when
 user-defined timestamp is enabled (#11223)

Summary:
in DBIter::SeekToLast(), key() can be called when iter is invalid and fails the following assertion:
```
./db/db_iter.h:153: virtual rocksdb::Slice rocksdb::DBIter::key() const: Assertion `valid_' failed.
```
This happens when `iterate_upper_bound` and timestamp_lb_ are set. SeekForPrev(*iterate_upper_bound_) positions the iterator on the same user key as *iterate_upper_bound_. A subsequent PrevInternal() call makes the iterator invalid just be the call to key().

This PR fixes this issue by setting updating the seek key to have max sequence number AND max timestamp when the seek key has the same user key as *iterate_upper_bound_.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/11223

Test Plan: - Added a unit test that would fail the above assertion before this fix.

Reviewed By: jowlyzhang

Differential Revision: D43283600

Pulled By: cbi42

fbshipit-source-id: 0dd3999845b722584679bbc95be2664b266005ba
---
 db/db_iter.cc                      | 26 ++++++++------------------
 db/db_with_timestamp_basic_test.cc | 26 ++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/db/db_iter.cc b/db/db_iter.cc
index 8626a4f810..dff251a409 100644
--- a/db/db_iter.cc
+++ b/db/db_iter.cc
@@ -1433,9 +1433,8 @@ void DBIter::SetSavedKeyToSeekForPrevTarget(const Slice& target) {
     if (timestamp_size_ > 0) {
       const std::string kTsMax(timestamp_size_, '\xff');
       Slice ts = kTsMax;
-      saved_key_.UpdateInternalKey(
-          kMaxSequenceNumber, kValueTypeForSeekForPrev,
-          timestamp_lb_ != nullptr ? timestamp_lb_ : &ts);
+      saved_key_.UpdateInternalKey(kMaxSequenceNumber, kValueTypeForSeekForPrev,
+                                   &ts);
     }
   }
 }
@@ -1637,24 +1636,15 @@ void DBIter::SeekToLast() {
   if (iterate_upper_bound_ != nullptr) {
     // Seek to last key strictly less than ReadOptions.iterate_upper_bound.
     SeekForPrev(*iterate_upper_bound_);
-    const bool is_ikey = (timestamp_size_ > 0 && timestamp_lb_ != nullptr);
+#ifndef NDEBUG
     Slice k = Valid() ? key() : Slice();
-    if (is_ikey && Valid()) {
+    if (Valid() && timestamp_size_ > 0 && timestamp_lb_) {
       k.remove_suffix(kNumInternalBytes + timestamp_size_);
     }
-    while (Valid() && 0 == user_comparator_.CompareWithoutTimestamp(
-                               *iterate_upper_bound_, /*a_has_ts=*/false, k,
-                               /*b_has_ts=*/false)) {
-      ReleaseTempPinnedData();
-      ResetBlobValue();
-      ResetValueAndColumns();
-      PrevInternal(nullptr);
-
-      k = key();
-      if (is_ikey) {
-        k.remove_suffix(kNumInternalBytes + timestamp_size_);
-      }
-    }
+    assert(!Valid() || user_comparator_.CompareWithoutTimestamp(
+                           k, /*a_has_ts=*/false, *iterate_upper_bound_,
+                           /*b_has_ts=*/false) < 0);
+#endif
     return;
   }
 
diff --git a/db/db_with_timestamp_basic_test.cc b/db/db_with_timestamp_basic_test.cc
index 31ebb25c2a..4b8132df3f 100644
--- a/db/db_with_timestamp_basic_test.cc
+++ b/db/db_with_timestamp_basic_test.cc
@@ -3892,6 +3892,32 @@ TEST_F(DBBasicTestWithTimestamp, RangeTombstoneApproximateSize) {
       std::numeric_limits<uint64_t>::max() /* max_file_num_to_ignore */,
       "" /*trim_ts*/));
 }
+
+TEST_F(DBBasicTestWithTimestamp, IterSeekToLastWithIterateUpperbound) {
+  // Test for a bug fix where DBIter::SeekToLast() could fail when
+  // iterate_upper_bound and iter_start_ts are both set.
+  Options options = CurrentOptions();
+  const size_t kTimestampSize = Timestamp(0, 0).size();
+  TestComparator test_cmp(kTimestampSize);
+  options.comparator = &test_cmp;
+  DestroyAndReopen(options);
+
+  ASSERT_OK(db_->Put(WriteOptions(), Key(1), Timestamp(2, 0), "val"));
+  ReadOptions ro;
+  std::string k = Key(1);
+  Slice k_slice = k;
+  ro.iterate_upper_bound = &k_slice;
+  std::string ts = Timestamp(3, 0);
+  Slice read_ts = ts;
+  ro.timestamp = &read_ts;
+  std::string start_ts = Timestamp(0, 0);
+  Slice start_ts_slice = start_ts;
+  ro.iter_start_ts = &start_ts_slice;
+  std::unique_ptr<Iterator> iter{db_->NewIterator(ro)};
+  iter->SeekToLast();
+  ASSERT_FALSE(iter->Valid());
+  ASSERT_OK(iter->status());
+}
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {