Clean up InternalIterator upper bound logic a little bit (#7200)

Summary:
IteratorIterator::IsOutOfBound() and IteratorIterator::MayBeOutOfUpperBound() are two functions that related to upper bound check. It is hard for users to reason about this complexity. Consolidate the two functions into one and assign an enum as results to improve readability.

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

Test Plan: Run all existing test. Would run crash test with atomic for a while.

Reviewed By: anand1976

Differential Revision: D22833181

fbshipit-source-id: a0c724267056adbd0476bde74650e6c7226077e6

db/db_iter.cc

@@ -223,11 +223,13 @@ bool DBIter::FindNextUserEntryInternal(bool skipping_saved_key,
 
     is_key_seqnum_zero_ = (ikey_.sequence == 0);
 
-    assert(iterate_upper_bound_ == nullptr || iter_.MayBeOutOfUpperBound() ||
+    assert(iterate_upper_bound_ == nullptr ||
+           iter_.UpperBoundCheckResult() != IterBoundCheck::kInbound ||
            user_comparator_.CompareWithoutTimestamp(
                ikey_.user_key, /*a_has_ts=*/true, *iterate_upper_bound_,
                /*b_has_ts=*/false) < 0);
-    if (iterate_upper_bound_ != nullptr && iter_.MayBeOutOfUpperBound() &&
+    if (iterate_upper_bound_ != nullptr &&
+        iter_.UpperBoundCheckResult() != IterBoundCheck::kInbound &&
         user_comparator_.CompareWithoutTimestamp(
             ikey_.user_key, /*a_has_ts=*/true, *iterate_upper_bound_,
             /*b_has_ts=*/false) >= 0) {

db/memtable.cc

@@ -382,7 +382,7 @@ class MemTableIterator : public InternalIterator {
     bool is_valid = valid_;
     if (is_valid) {
       result->key = key();
-      result->may_be_out_of_upper_bound = true;
+      result->bound_check_result = IterBoundCheck::kUnknown;
       result->value_prepared = true;
     }
     return is_valid;

db/version_set.cc

@@ -942,9 +942,12 @@ class LevelIterator final : public InternalIterator {
     return may_be_out_of_lower_bound_ && file_iter_.MayBeOutOfLowerBound();
   }
 
-  inline bool MayBeOutOfUpperBound() override {
-    assert(Valid());
-    return file_iter_.MayBeOutOfUpperBound();
+  inline IterBoundCheck UpperBoundCheckResult() override {
+    if (Valid()) {
+      return file_iter_.UpperBoundCheckResult();
+    } else {
+      return IterBoundCheck::kUnknown;
+    }
   }
 
   void SetPinnedItersMgr(PinnedIteratorsManager* pinned_iters_mgr) override {
@@ -1148,7 +1151,7 @@ bool LevelIterator::NextAndGetResult(IterateResult* result) {
     is_valid = Valid();
     if (is_valid) {
       result->key = key();
-      result->may_be_out_of_upper_bound = MayBeOutOfUpperBound();
+      result->bound_check_result = file_iter_.UpperBoundCheckResult();
       // Ideally, we should return the real file_iter_.value_prepared but the
       // information is not here. It would casue an extra PrepareValue()
       // for the first key of a file.
@@ -1168,7 +1171,8 @@ bool LevelIterator::SkipEmptyFileForward() {
   bool seen_empty_file = false;
   while (file_iter_.iter() == nullptr ||
          (!file_iter_.Valid() && file_iter_.status().ok() &&
-          !file_iter_.iter()->IsOutOfBound())) {
+          file_iter_.iter()->UpperBoundCheckResult() !=
+              IterBoundCheck::kOutOfBound)) {
     seen_empty_file = true;
     // Move to next file
     if (file_index_ >= flevel_->num_files - 1) {

table/block_based/block_based_table_iterator.cc

@@ -189,7 +189,7 @@ bool BlockBasedTableIterator::NextAndGetResult(IterateResult* result) {
   bool is_valid = Valid();
   if (is_valid) {
     result->key = key();
-    result->may_be_out_of_upper_bound = MayBeOutOfUpperBound();
+    result->bound_check_result = UpperBoundCheckResult();
     result->value_prepared = !is_at_first_key_from_index_;
   }
   return is_valid;

table/block_based/block_based_table_iterator.h

@@ -99,13 +99,16 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
     }
   }
 
-  // Whether iterator invalidated for being out of bound.
-  bool IsOutOfBound() override { return is_out_of_bound_; }
-
-  inline bool MayBeOutOfUpperBound() override {
-    assert(Valid());
-    return block_upper_bound_check_ !=
-           BlockUpperBound::kUpperBoundBeyondCurBlock;
+  inline IterBoundCheck UpperBoundCheckResult() override {
+    if (is_out_of_bound_) {
+      return IterBoundCheck::kOutOfBound;
+    } else if (block_upper_bound_check_ ==
+               BlockUpperBound::kUpperBoundBeyondCurBlock) {
+      assert(!is_out_of_bound_);
+      return IterBoundCheck::kInbound;
+    } else {
+      return IterBoundCheck::kUnknown;
+    }
   }
 
   void SetPinnedItersMgr(PinnedIteratorsManager* pinned_iters_mgr) override {

table/block_based/partitioned_index_iterator.h

@@ -78,18 +78,10 @@ class ParititionedIndexIterator : public InternalIteratorBase<IndexValue> {
       return Status::OK();
     }
   }
-
-  // Whether iterator invalidated for being out of bound.
-  bool IsOutOfBound() override {
-    // Shoulldn't be called
-    assert(false);
-    return false;
-  }
-
-  inline bool MayBeOutOfUpperBound() override {
+  inline IterBoundCheck UpperBoundCheckResult() override {
     // Shouldn't be called.
     assert(false);
-    return true;
+    return IterBoundCheck::kUnknown;
   }
   void SetPinnedItersMgr(PinnedIteratorsManager*) override {
     // Shouldn't be called.

table/internal_iterator.h

@@ -17,9 +17,15 @@ namespace ROCKSDB_NAMESPACE {
 
 class PinnedIteratorsManager;
 
+enum class IterBoundCheck : char {
+  kUnknown = 0,
+  kOutOfBound,
+  kInbound,
+};
+
 struct IterateResult {
   Slice key;
-  bool may_be_out_of_upper_bound;
+  IterBoundCheck bound_check_result = IterBoundCheck::kUnknown;
   // If false, PrepareValue() needs to be called before value().
   bool value_prepared = true;
 };
@@ -69,7 +75,7 @@ class InternalIteratorBase : public Cleanable {
 
   // Moves to the next entry in the source, and return result. Iterator
   // implementation should override this method to help methods inline better,
-  // or when MayBeOutOfUpperBound() is non-trivial.
+  // or when UpperBoundCheckResult() is non-trivial.
   // REQUIRES: Valid()
   virtual bool NextAndGetResult(IterateResult* result) {
     Next();
@@ -77,11 +83,11 @@ class InternalIteratorBase : public Cleanable {
     if (is_valid) {
       result->key = key();
       // Default may_be_out_of_upper_bound to true to avoid unnecessary virtual
-      // call. If an implementation has non-trivial MayBeOutOfUpperBound(),
+      // call. If an implementation has non-trivial UpperBoundCheckResult(),
       // it should also override NextAndGetResult().
-      result->may_be_out_of_upper_bound = true;
+      result->bound_check_result = IterBoundCheck::kUnknown;
       result->value_prepared = false;
-      assert(MayBeOutOfUpperBound());
+      assert(UpperBoundCheckResult() != IterBoundCheck::kOutOfBound);
     }
     return is_valid;
   }
@@ -126,19 +132,17 @@ class InternalIteratorBase : public Cleanable {
   // REQUIRES: Valid()
   virtual bool PrepareValue() { return true; }
 
-  // True if the iterator is invalidated because it reached a key that is above
-  // the iterator upper bound. Used by LevelIterator to decide whether it should
-  // stop or move on to the next file.
-  // Important: if iterator reached the end of the file without encountering any
-  // keys above the upper bound, IsOutOfBound() must return false.
-  virtual bool IsOutOfBound() { return false; }
-
   // Keys return from this iterator can be smaller than iterate_lower_bound.
   virtual bool MayBeOutOfLowerBound() { return true; }
 
-  // Keys return from this iterator can be larger or equal to
-  // iterate_upper_bound.
-  virtual bool MayBeOutOfUpperBound() { return true; }
+  // If the iterator has checked the key against iterate_upper_bound, returns
+  // the result here. The function can be used by user of the iterator to skip
+  // their own checks. If Valid() = true, IterBoundCheck::kUnknown is always
+  // a valid value. If Valid() = false, IterBoundCheck::kOutOfBound indicates
+  // that the iterator is filtered out by upper bound checks.
+  virtual IterBoundCheck UpperBoundCheckResult() {
+    return IterBoundCheck::kUnknown;
+  }
 
   // Pass the PinnedIteratorsManager to the Iterator, most Iterators don't
   // communicate with PinnedIteratorsManager so default implementation is no-op

table/iterator_wrapper.h

@@ -127,9 +127,9 @@ class IteratorWrapperBase {
     return iter_->MayBeOutOfLowerBound();
   }
 
-  bool MayBeOutOfUpperBound() {
+  IterBoundCheck UpperBoundCheckResult() {
     assert(Valid());
-    return result_.may_be_out_of_upper_bound;
+    return result_.bound_check_result;
   }
 
   void SetPinnedItersMgr(PinnedIteratorsManager* pinned_iters_mgr) {
@@ -155,7 +155,7 @@ class IteratorWrapperBase {
     if (valid_) {
       assert(iter_->status().ok());
       result_.key = iter_->key();
-      result_.may_be_out_of_upper_bound = true;
+      result_.bound_check_result = IterBoundCheck::kUnknown;
       result_.value_prepared = false;
     }
   }

table/merging_iterator.cc

@@ -194,7 +194,7 @@ class MergingIterator : public InternalIterator {
     bool is_valid = Valid();
     if (is_valid) {
       result->key = key();
-      result->may_be_out_of_upper_bound = MayBeOutOfUpperBound();
+      result->bound_check_result = UpperBoundCheckResult();
       result->value_prepared = current_->IsValuePrepared();
     }
     return is_valid;
@@ -261,9 +261,9 @@ class MergingIterator : public InternalIterator {
     return current_->MayBeOutOfLowerBound();
   }
 
-  bool MayBeOutOfUpperBound() override {
+  IterBoundCheck UpperBoundCheckResult() override {
     assert(Valid());
-    return current_->MayBeOutOfUpperBound();
+    return current_->UpperBoundCheckResult();
   }
 
   void SetPinnedItersMgr(PinnedIteratorsManager* pinned_iters_mgr) override {

table/table_test.cc

@@ -241,7 +241,9 @@ class KeyConvertingIterator : public InternalIterator {
   void SeekToLast() override { iter_->SeekToLast(); }
   void Next() override { iter_->Next(); }
   void Prev() override { iter_->Prev(); }
-  bool IsOutOfBound() override { return iter_->IsOutOfBound(); }
+  IterBoundCheck UpperBoundCheckResult() override {
+    return iter_->UpperBoundCheckResult();
+  }
 
   Slice key() const override {
     assert(Valid());
@@ -4645,13 +4647,13 @@ TEST_P(BlockBasedTableTest, OutOfBoundOnSeek) {
       /*skip_filters=*/false, TableReaderCaller::kUncategorized)));
   iter->SeekToFirst();
   ASSERT_FALSE(iter->Valid());
-  ASSERT_TRUE(iter->IsOutOfBound());
+  ASSERT_TRUE(iter->UpperBoundCheckResult() == IterBoundCheck::kOutOfBound);
   iter.reset(new KeyConvertingIterator(reader->NewIterator(
       read_opt, /*prefix_extractor=*/nullptr, /*arena=*/nullptr,
       /*skip_filters=*/false, TableReaderCaller::kUncategorized)));
   iter->Seek("foo");
   ASSERT_FALSE(iter->Valid());
-  ASSERT_TRUE(iter->IsOutOfBound());
+  ASSERT_TRUE(iter->UpperBoundCheckResult() == IterBoundCheck::kOutOfBound);
 }
 
 // BlockBasedTableIterator should invalidate itself and return
@@ -4686,7 +4688,7 @@ TEST_P(BlockBasedTableTest, OutOfBoundOnNext) {
   ASSERT_EQ("bar", iter->key());
   iter->Next();
   ASSERT_FALSE(iter->Valid());
-  ASSERT_TRUE(iter->IsOutOfBound());
+  ASSERT_TRUE(iter->UpperBoundCheckResult() == IterBoundCheck::kOutOfBound);
   std::string ub2 = "foo_after";
   Slice ub_slice2(ub2);
   read_opt.iterate_upper_bound = &ub_slice2;
@@ -4698,7 +4700,7 @@ TEST_P(BlockBasedTableTest, OutOfBoundOnNext) {
   ASSERT_EQ("foo", iter->key());
   iter->Next();
   ASSERT_FALSE(iter->Valid());
-  ASSERT_FALSE(iter->IsOutOfBound());
+  ASSERT_FALSE(iter->UpperBoundCheckResult() == IterBoundCheck::kOutOfBound);
 }
 
 }  // namespace ROCKSDB_NAMESPACE