Fix a perf regression in ForwardIterator

Summary:
I noticed that memtable iterator usually crosses the `iterate_upper_bound`
threshold when tailing. Changes introduced in D43833 made `NeedToSeekImmutable`
always return true in such case, even when `Seek()` only needs to rewind the
memtable iterator. In a test I ran, this caused the "tailing efficiency"
(ratio of calls to `Seek()` that only affect the memtable versus all seeks)
to drop almost to zero.

This diff attempts to fix the regression by using a different flag to indicate
that `current_` is over the limit instead of resetting `valid_` in
`UpdateCurrent()`.

Test Plan: `DBTestTailingIterator.TailingIteratorUpperBound`

Reviewers: sdong, rven

Reviewed By: rven

Subscribers: dhruba, march

Differential Revision: https://reviews.facebook.net/D45909

db/db_tailing_iter_test.cc

@@ -358,6 +358,49 @@ TEST_F(DBTestTailingIterator, TailingIteratorSeekToSame) {
   ASSERT_EQ(found, iter->key().ToString());
 }
 
+// Sets iterate_upper_bound and verifies that ForwardIterator doesn't call
+// Seek() on immutable iterators when target key is >= prev_key and all
+// iterators, including the memtable iterator, are over the upper bound.
+TEST_F(DBTestTailingIterator, TailingIteratorUpperBound) {
+  CreateAndReopenWithCF({"pikachu"}, CurrentOptions());
+
+  const Slice upper_bound("20", 3);
+  ReadOptions read_options;
+  read_options.tailing = true;
+  read_options.iterate_upper_bound = &upper_bound;
+
+  ASSERT_OK(Put(1, "11", "11"));
+  ASSERT_OK(Put(1, "12", "12"));
+  ASSERT_OK(Put(1, "22", "22"));
+  ASSERT_OK(Flush(1));  // flush all those keys to an immutable SST file
+
+  // Add another key to the memtable.
+  ASSERT_OK(Put(1, "21", "21"));
+
+  std::unique_ptr<Iterator> it(db_->NewIterator(read_options, handles_[1]));
+  it->Seek("12");
+  ASSERT_TRUE(it->Valid());
+  ASSERT_EQ("12", it->key().ToString());
+
+  it->Next();
+  // Not valid since "21" is over the upper bound.
+  ASSERT_FALSE(it->Valid());
+
+  // This keeps track of the number of times NeedToSeekImmutable() was true.
+  int immutable_seeks = 0;
+  rocksdb::SyncPoint::GetInstance()->SetCallBack(
+      "ForwardIterator::SeekInternal:Immutable",
+      [&](void* arg) { ++immutable_seeks; });
+
+  // Seek to 13. This should not require any immutable seeks.
+  rocksdb::SyncPoint::GetInstance()->EnableProcessing();
+  it->Seek("13");
+  rocksdb::SyncPoint::GetInstance()->DisableProcessing();
+
+  ASSERT_FALSE(it->Valid());
+  ASSERT_EQ(0, immutable_seeks);
+}
+
 TEST_F(DBTestTailingIterator, ManagedTailingIteratorSingle) {
   ReadOptions read_options;
   read_options.tailing = true;

db/forward_iterator.cc

@@ -128,10 +128,11 @@ ForwardIterator::ForwardIterator(DBImpl* db, const ReadOptions& read_options,
       sv_(current_sv),
       mutable_iter_(nullptr),
       current_(nullptr),
+      valid_(false),
       status_(Status::OK()),
       immutable_status_(Status::OK()),
-      valid_(false),
       has_iter_trimmed_for_upper_bound_(false),
+      current_over_upper_bound_(false),
       is_prev_set_(false),
       is_prev_inclusive_(false) {
   if (sv_) {
@@ -179,7 +180,8 @@ void ForwardIterator::Cleanup(bool release_sv) {
 }
 
 bool ForwardIterator::Valid() const {
-  return valid_;
+  // See UpdateCurrent().
+  return valid_ ? !current_over_upper_bound_ : false;
 }
 
 void ForwardIterator::SeekToFirst() {
@@ -225,7 +227,7 @@ void ForwardIterator::SeekInternal(const Slice& internal_key,
   // an option to turn it off.
   if (seek_to_first || NeedToSeekImmutable(internal_key)) {
     immutable_status_ = Status::OK();
-    if (NeedToRebuildTrimmed(internal_key)) {
+    if (has_iter_trimmed_for_upper_bound_) {
       // Some iterators are trimmed. Need to rebuild.
       RebuildIterators(true);
       // Already seeked mutable iter, so seek again
@@ -375,6 +377,8 @@ void ForwardIterator::SeekInternal(const Slice& internal_key,
       is_prev_set_ = true;
       is_prev_inclusive_ = true;
     }
+
+    TEST_SYNC_POINT_CALLBACK("ForwardIterator::SeekInternal:Immutable", this);
   } else if (current_ && current_ != mutable_iter_) {
     // current_ is one of immutable iterators, push it back to the heap
     immutable_min_heap_.push(current_);
@@ -555,10 +559,13 @@ void ForwardIterator::UpdateCurrent() {
   if (!status_.ok()) {
     status_ = Status::OK();
   }
-  if (valid_ && IsOverUpperBound(current_->key())) {
-    valid_ = false;
-    current_ = nullptr;
-  }
+
+  // Upper bound doesn't apply to the memtable iterator. We want Valid() to
+  // return false when all iterators are over iterate_upper_bound, but can't
+  // just set valid_ to false, as that would effectively disable the tailing
+  // optimization (Seek() would be called on all immutable iterators regardless
+  // of whether the target key is greater than prev_key_).
+  current_over_upper_bound_ = valid_ && IsOverUpperBound(current_->key());
 }
 
 bool ForwardIterator::NeedToSeekImmutable(const Slice& target) {
@@ -594,26 +601,6 @@ bool ForwardIterator::NeedToSeekImmutable(const Slice& target) {
   return false;
 }
 
-bool ForwardIterator::NeedToRebuildTrimmed(const Slice& target) {
-  if (!has_iter_trimmed_for_upper_bound_) {
-    return false;
-  }
-  if (!valid_ || !current_ || !is_prev_set_ || !immutable_status_.ok()) {
-    return true;
-  }
-  Slice prev_key = prev_key_.GetKey();
-  if (prefix_extractor_ &&
-      prefix_extractor_->Transform(target)
-              .compare(prefix_extractor_->Transform(prev_key)) != 0) {
-    return true;
-  }
-  if (cfd_->internal_comparator().InternalKeyComparator::Compare(
-          prev_key, target) >= (is_prev_inclusive_ ? 1 : 0)) {
-    return true;
-  }
-  return false;
-}
-
 void ForwardIterator::DeleteCurrentIter() {
   const VersionStorageInfo* vstorage = sv_->current->storage_info();
   const std::vector<FileMetaData*>& l0 = vstorage->LevelFiles(0);

db/forward_iterator.h

@@ -79,7 +79,6 @@ class ForwardIterator : public Iterator {
   void SeekInternal(const Slice& internal_key, bool seek_to_first);
   void UpdateCurrent();
   bool NeedToSeekImmutable(const Slice& internal_key);
-  bool NeedToRebuildTrimmed(const Slice& internal_key);
   void DeleteCurrentIter();
   uint32_t FindFileInRange(
     const std::vector<FileMetaData*>& files, const Slice& internal_key,
@@ -100,16 +99,30 @@ class ForwardIterator : public Iterator {
   std::vector<Iterator*> l0_iters_;
   std::vector<LevelIterator*> level_iters_;
   Iterator* current_;
-  // internal iterator status
-  Status status_;
-  Status immutable_status_;
   bool valid_;
-  bool has_iter_trimmed_for_upper_bound_;
-  Slice smallest_file_key_bound;
 
+  // Internal iterator status; set only by one of the unsupported methods.
+  Status status_;
+  // Status of immutable iterators, maintained here to avoid iterating over
+  // all of them in status().
+  Status immutable_status_;
+  // Indicates that at least one of the immutable iterators pointed to a key
+  // larger than iterate_upper_bound and was therefore destroyed. Seek() may
+  // need to rebuild such iterators.
+  bool has_iter_trimmed_for_upper_bound_;
+  // Is current key larger than iterate_upper_bound? If so, makes Valid()
+  // return false.
+  bool current_over_upper_bound_;
+
+  // Left endpoint of the range of keys that immutable iterators currently
+  // cover. When Seek() is called with a key that's within that range, immutable
+  // iterators don't need to be moved; see NeedToSeekImmutable(). This key is
+  // included in the range after a Seek(), but excluded when advancing the
+  // iterator using Next().
   IterKey prev_key_;
   bool is_prev_set_;
   bool is_prev_inclusive_;
+
   Arena arena_;
 };