Separate some IndexBlockIter logic from BlockIter (#4136)

Summary:
Some logic only related to IndexBlockIter is separated from BlockIter to IndexBlockIter. This is done by writing an exclusive Seek() and SeekForPrev() for DataBlockIter, and all metadata block iter and tombstone block iter now use data block iter. Dealing with the BinarySeek() sharing problem by passing in the comparator to use.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4136

Reviewed By: maysamyabandeh

Differential Revision: D8859673

Pulled By: siying

fbshipit-source-id: 703e5e6824b82b7cbf4721f3594b94127797ca9e

table/block.cc

@@ -139,17 +139,14 @@ void BlockIter::Prev() {
   prev_entries_idx_ = static_cast<int32_t>(prev_entries_.size()) - 1;
 }
 
-void BlockIter::Seek(const Slice& target) {
+void DataBlockIter::Seek(const Slice& target) {
   Slice seek_key = target;
-  if (!key_includes_seq_) {
-    seek_key = ExtractUserKey(target);
-  }
   PERF_TIMER_GUARD(block_seek_nanos);
   if (data_ == nullptr) {  // Not init yet
     return;
   }
   uint32_t index = 0;
-  bool ok = BinarySeek(seek_key, 0, num_restarts_ - 1, &index);
+  bool ok = BinarySeek(seek_key, 0, num_restarts_ - 1, &index, comparator_);
 
   if (!ok) {
     return;
@@ -178,7 +175,8 @@ void IndexBlockIter::Seek(const Slice& target) {
   if (prefix_index_) {
     ok = PrefixSeek(target, &index);
   } else {
-    ok = BinarySeek(seek_key, 0, num_restarts_ - 1, &index);
+    ok = BinarySeek(seek_key, 0, num_restarts_ - 1, &index,
+                    key_includes_seq_ ? comparator_ : user_comparator_);
   }
 
   if (!ok) {
@@ -194,17 +192,14 @@ void IndexBlockIter::Seek(const Slice& target) {
   }
 }
 
-void BlockIter::SeekForPrev(const Slice& target) {
+void DataBlockIter::SeekForPrev(const Slice& target) {
   PERF_TIMER_GUARD(block_seek_nanos);
   Slice seek_key = target;
-  if (!key_includes_seq_) {
-    seek_key = ExtractUserKey(target);
-  }
   if (data_ == nullptr) {  // Not init yet
     return;
   }
   uint32_t index = 0;
-  bool ok = BinarySeek(seek_key, 0, num_restarts_ - 1, &index);
+  bool ok = BinarySeek(seek_key, 0, num_restarts_ - 1, &index, comparator_);
 
   if (!ok) {
     return;
@@ -321,7 +316,7 @@ bool BlockIter::ParseNextKey() {
 // which means the key of next restart point is larger than target, or
 // the first restart point with a key = target
 bool BlockIter::BinarySeek(const Slice& target, uint32_t left, uint32_t right,
-                           uint32_t* index) {
+                           uint32_t* index, const Comparator* comp) {
   assert(left <= right);
 
   while (left < right) {
@@ -335,7 +330,7 @@ bool BlockIter::BinarySeek(const Slice& target, uint32_t left, uint32_t right,
       return false;
     }
     Slice mid_key(key_ptr, non_shared);
-    int cmp = Compare(mid_key, target);
+    int cmp = comp->Compare(mid_key, target);
     if (cmp < 0) {
       // Key at "mid" is smaller than "target". Therefore all
       // blocks before "mid" are uninteresting.
@@ -354,7 +349,7 @@ bool BlockIter::BinarySeek(const Slice& target, uint32_t left, uint32_t right,
 }
 
 // Compare target key and the block key of the block of `block_index`.
-// Return -1 if error.
+// Return -1 if eror.
 int IndexBlockIter::CompareBlockKey(uint32_t block_index, const Slice& target) {
   uint32_t region_offset = GetRestartPoint(block_index);
   uint32_t shared, non_shared, value_length;
@@ -471,35 +466,6 @@ Block::Block(BlockContents&& contents, SequenceNumber _global_seqno,
   }
 }
 
-template <>
-BlockIter* Block::NewIterator(const Comparator* cmp, const Comparator* ucmp,
-                              BlockIter* iter, Statistics* /*stats*/,
-                              bool /*total_order_seek*/,
-                              bool /*key_includes_seq*/,
-                              BlockPrefixIndex* /*prefix_index*/) {
-  BlockIter* ret_iter;
-  if (iter != nullptr) {
-    ret_iter = iter;
-  } else {
-    ret_iter = new BlockIter;
-  }
-  if (size_ < 2*sizeof(uint32_t)) {
-    ret_iter->Invalidate(Status::Corruption("bad block contents"));
-    return ret_iter;
-  }
-  if (num_restarts_ == 0) {
-    // Empty block.
-    ret_iter->Invalidate(Status::OK());
-    return ret_iter;
-  } else {
-    const bool kKeyIncludesSeq = true;
-    ret_iter->InitializeBase(cmp, ucmp, data_, restart_offset_, num_restarts_,
-                             global_seqno_, kKeyIncludesSeq, cachable());
-  }
-
-  return ret_iter;
-}
-
 template <>
 DataBlockIter* Block::NewIterator(const Comparator* cmp, const Comparator* ucmp,
                                   DataBlockIter* iter, Statistics* stats,

table/block.h

@@ -206,34 +206,7 @@ class Block {
 
 class BlockIter : public InternalIterator {
  public:
-  // Object created using this constructor will behave like an iterator
+  void InitializeBase(const Comparator* comparator, const char* data,
-  // against an empty block. The state after the creation: Valid()=false
-  // and status() is OK.
-  BlockIter()
-      : comparator_(nullptr),
-        user_comparator_(nullptr),
-        data_(nullptr),
-        num_restarts_(0),
-        restart_index_(0),
-        restarts_(0),
-        current_(0),
-        status_(Status::OK()),
-        key_pinned_(false),
-        block_contents_pinned_(false),
-        key_includes_seq_(true),
-        global_seqno_(kDisableGlobalSequenceNumber) {}
-
-  BlockIter(const Comparator* comparator, const Comparator* user_comparator,
-            const char* data, uint32_t restarts, uint32_t num_restarts,
-            SequenceNumber global_seqno, bool key_includes_seq,
-            bool block_contents_pinned)
-      : BlockIter() {
-    InitializeBase(comparator, user_comparator, data, restarts, num_restarts,
-                   global_seqno, key_includes_seq, block_contents_pinned);
-  }
-
-  void InitializeBase(const Comparator* comparator,
-                      const Comparator* user_comparator, const char* data,
                       uint32_t restarts, uint32_t num_restarts,
                       SequenceNumber global_seqno, bool key_includes_seq,
                       bool block_contents_pinned) {
@@ -241,7 +214,6 @@ class BlockIter : public InternalIterator {
     assert(num_restarts > 0);           // Ensure the param is valid
 
     comparator_ = comparator;
-    user_comparator_ = user_comparator;
     data_ = data;
     restarts_ = restarts;
     num_restarts_ = num_restarts;
@@ -287,10 +259,6 @@ class BlockIter : public InternalIterator {
 
   virtual void Prev() override;
 
-  virtual void Seek(const Slice& target) override;
-
-  virtual void SeekForPrev(const Slice& target) override;
-
   virtual void SeekToFirst() override;
 
   virtual void SeekToLast() override;
@@ -324,8 +292,6 @@ class BlockIter : public InternalIterator {
   // Note: The type could be changed to InternalKeyComparator but we see a weird
   // performance drop by that.
   const Comparator* comparator_;
-  // Same as comparator_ if comparator_ is not InernalKeyComparator
-  const Comparator* user_comparator_;
   const char* data_;       // underlying block contents
   uint32_t num_restarts_;  // Number of uint32_t entries in restart array
 
@@ -368,22 +334,6 @@ class BlockIter : public InternalIterator {
   int32_t prev_entries_idx_ = -1;
 
  public:
-  inline int Compare(const Slice& a, const Slice& b) const {
-    if (key_includes_seq_) {
-      return comparator_->Compare(a, b);
-    } else {
-      return user_comparator_->Compare(a, b);
-    }
-  }
-
-  inline int Compare(const IterKey& ikey, const Slice& b) const {
-    if (key_includes_seq_) {
-      return comparator_->Compare(ikey.GetInternalKey(), b);
-    } else {
-      return user_comparator_->Compare(ikey.GetUserKey(), b);
-    }
-  }
-
   // Return the offset in data_ just past the end of the current entry.
   inline uint32_t NextEntryOffset() const {
     // NOTE: We don't support blocks bigger than 2GB
@@ -408,9 +358,8 @@ class BlockIter : public InternalIterator {
   void CorruptionError();
 
   bool ParseNextKey();
-
   bool BinarySeek(const Slice& target, uint32_t left, uint32_t right,
-                  uint32_t* index);
+                  uint32_t* index, const Comparator* comp);
 };
 
 class DataBlockIter final : public BlockIter {
@@ -426,14 +375,14 @@ class DataBlockIter final : public BlockIter {
                global_seqno, read_amp_bitmap, block_contents_pinned);
   }
   void Initialize(const Comparator* comparator,
-                  const Comparator* user_comparator, const char* data,
+                  const Comparator* /*user_comparator*/, const char* data,
                   uint32_t restarts, uint32_t num_restarts,
                   SequenceNumber global_seqno,
                   BlockReadAmpBitmap* read_amp_bitmap,
                   bool block_contents_pinned) {
     const bool kKeyIncludesSeq = true;
-    InitializeBase(comparator, user_comparator, data, restarts, num_restarts,
+    InitializeBase(comparator, data, restarts, num_restarts, global_seqno,
-                   global_seqno, kKeyIncludesSeq, block_contents_pinned);
+                   kKeyIncludesSeq, block_contents_pinned);
     read_amp_bitmap_ = read_amp_bitmap;
     last_bitmap_offset_ = current_ + 1;
   }
@@ -449,11 +398,19 @@ class DataBlockIter final : public BlockIter {
     return value_;
   }
 
+  virtual void Seek(const Slice& target) override;
+
+  virtual void SeekForPrev(const Slice& target) override;
+
  private:
   // read-amp bitmap
   BlockReadAmpBitmap* read_amp_bitmap_;
   // last `current_` value we report to read-amp bitmp
   mutable uint32_t last_bitmap_offset_;
+
+  inline int Compare(const IterKey& ikey, const Slice& b) const {
+    return comparator_->Compare(ikey.GetInternalKey(), b);
+  }
 };
 
 class IndexBlockIter final : public BlockIter {
@@ -479,7 +436,8 @@ class IndexBlockIter final : public BlockIter {
                   uint32_t restarts, uint32_t num_restarts,
                   BlockPrefixIndex* prefix_index, bool key_includes_seq,
                   bool block_contents_pinned) {
-    InitializeBase(comparator, user_comparator, data, restarts, num_restarts,
+    user_comparator_ = user_comparator;
+    InitializeBase(comparator, data, restarts, num_restarts,
                    kDisableGlobalSequenceNumber, key_includes_seq,
                    block_contents_pinned);
     prefix_index_ = prefix_index;
@@ -487,6 +445,17 @@ class IndexBlockIter final : public BlockIter {
 
   virtual void Seek(const Slice& target) override;
 
+  virtual void SeekForPrev(const Slice&) override {
+    assert(false);
+    current_ = restarts_;
+    restart_index_ = num_restarts_;
+    status_ = Status::InvalidArgument(
+        "RocksDB internal error: should never all SeekForPrev() for index "
+        "blocks");
+    key_.Clear();
+    value_.clear();
+  }
+
  private:
   bool PrefixSeek(const Slice& target, uint32_t* index);
   bool BinaryBlockIndexSeek(const Slice& target, uint32_t* block_ids,
@@ -494,6 +463,24 @@ class IndexBlockIter final : public BlockIter {
                             uint32_t* index);
   int CompareBlockKey(uint32_t block_index, const Slice& target);
 
+  inline int Compare(const Slice& a, const Slice& b) const {
+    if (key_includes_seq_) {
+      return comparator_->Compare(a, b);
+    } else {
+      return user_comparator_->Compare(a, b);
+    }
+  }
+
+  inline int Compare(const IterKey& ikey, const Slice& b) const {
+    if (key_includes_seq_) {
+      return comparator_->Compare(ikey.GetInternalKey(), b);
+    } else {
+      return user_comparator_->Compare(ikey.GetUserKey(), b);
+    }
+  }
+
+  // Same as comparator_ if comparator_ is not InernalKeyComparator
+  const Comparator* user_comparator_;
   BlockPrefixIndex* prefix_index_;
 };
 

table/block_based_table_reader.cc

@@ -1119,8 +1119,8 @@ Status BlockBasedTable::ReadMetaBlock(Rep* rep,
 
   *meta_block = std::move(meta);
   // meta block uses bytewise comparator.
-  iter->reset(meta_block->get()->NewIterator<BlockIter>(BytewiseComparator(),
+  iter->reset(meta_block->get()->NewIterator<DataBlockIter>(
-                                                        BytewiseComparator()));
+      BytewiseComparator(), BytewiseComparator()));
   return Status::OK();
 }
 
@@ -1809,7 +1809,7 @@ BlockBasedTable::PartitionedIndexIteratorState::NewSecondaryIterator(
         nullptr, kNullStats, true, index_key_includes_seq_);
   }
   // Create an empty iterator
-  return new BlockIter();
+  return new DataBlockIter();
 }
 
 // This will be broken if the user specifies an unusual implementation
@@ -2211,7 +2211,7 @@ InternalIterator* BlockBasedTable::NewRangeTombstoneIterator(
     Cache* block_cache = rep_->table_options.block_cache.get();
     assert(block_cache != nullptr);
     if (block_cache->Ref(rep_->range_del_entry.cache_handle)) {
-      auto iter = rep_->range_del_entry.value->NewIterator<BlockIter>(
+      auto iter = rep_->range_del_entry.value->NewIterator<DataBlockIter>(
           &rep_->internal_comparator,
           rep_->internal_comparator.user_comparator());
       iter->RegisterCleanup(&ReleaseCachedEntry, block_cache,
@@ -2223,7 +2223,7 @@ InternalIterator* BlockBasedTable::NewRangeTombstoneIterator(
   rep_->range_del_handle.EncodeTo(&str);
   // The meta-block exists but isn't in uncompressed block cache (maybe
   // because it is disabled), so go through the full lookup process.
-  return NewDataBlockIterator<BlockIter>(rep_, read_options, Slice(str));
+  return NewDataBlockIterator<DataBlockIter>(rep_, read_options, Slice(str));
 }
 
 bool BlockBasedTable::FullFilterKeyMayMatch(

table/block_test.cc

@@ -100,7 +100,7 @@ TEST_F(BlockTest, SimpleTest) {
   // read contents of block sequentially
   int count = 0;
   InternalIterator *iter =
-      reader.NewIterator<BlockIter>(options.comparator, options.comparator);
+      reader.NewIterator<DataBlockIter>(options.comparator, options.comparator);
   for (iter->SeekToFirst();iter->Valid(); count++, iter->Next()) {
 
     // read kv from block
@@ -114,7 +114,8 @@ TEST_F(BlockTest, SimpleTest) {
   delete iter;
 
   // read block contents randomly
-  iter = reader.NewIterator<BlockIter>(options.comparator, options.comparator);
+  iter =
+      reader.NewIterator<DataBlockIter>(options.comparator, options.comparator);
   for (int i = 0; i < num_records; i++) {
 
     // find a random key in the lookaside array
@@ -163,8 +164,9 @@ void CheckBlockContents(BlockContents contents, const int max_key,
   std::unique_ptr<const SliceTransform> prefix_extractor(
       NewFixedPrefixTransform(prefix_size));
 
-  std::unique_ptr<InternalIterator> regular_iter(reader2.NewIterator<BlockIter>(
+  std::unique_ptr<InternalIterator> regular_iter(
-      BytewiseComparator(), BytewiseComparator()));
+      reader2.NewIterator<DataBlockIter>(BytewiseComparator(),
+                                         BytewiseComparator()));
 
   // Seek existent keys
   for (size_t i = 0; i < keys.size(); i++) {

table/meta_blocks.cc

@@ -203,8 +203,8 @@ Status ReadProperties(const Slice& handle_value, RandomAccessFileReader* file,
 
   Block properties_block(std::move(block_contents),
                          kDisableGlobalSequenceNumber);
-  BlockIter iter;
+  DataBlockIter iter;
-  properties_block.NewIterator<BlockIter>(BytewiseComparator(),
+  properties_block.NewIterator<DataBlockIter>(BytewiseComparator(),
                                               BytewiseComparator(), &iter);
 
   auto new_table_properties = new TableProperties();
@@ -335,7 +335,7 @@ Status ReadTableProperties(RandomAccessFileReader* file, uint64_t file_size,
   Block metaindex_block(std::move(metaindex_contents),
                         kDisableGlobalSequenceNumber);
   std::unique_ptr<InternalIterator> meta_iter(
-      metaindex_block.NewIterator<BlockIter>(BytewiseComparator(),
+      metaindex_block.NewIterator<DataBlockIter>(BytewiseComparator(),
                                                  BytewiseComparator()));
 
   // -- Read property block
@@ -405,8 +405,8 @@ Status FindMetaBlock(RandomAccessFileReader* file, uint64_t file_size,
                         kDisableGlobalSequenceNumber);
 
   std::unique_ptr<InternalIterator> meta_iter;
-  meta_iter.reset(metaindex_block.NewIterator<BlockIter>(BytewiseComparator(),
+  meta_iter.reset(metaindex_block.NewIterator<DataBlockIter>(
-                                                         BytewiseComparator()));
+      BytewiseComparator(), BytewiseComparator()));
 
   return FindMetaBlock(meta_iter.get(), meta_block_name, block_handle);
 }
@@ -452,8 +452,8 @@ Status ReadMetaBlock(RandomAccessFileReader* file,
                         kDisableGlobalSequenceNumber);
 
   std::unique_ptr<InternalIterator> meta_iter;
-  meta_iter.reset(metaindex_block.NewIterator<BlockIter>(BytewiseComparator(),
+  meta_iter.reset(metaindex_block.NewIterator<DataBlockIter>(
-                                                         BytewiseComparator()));
+      BytewiseComparator(), BytewiseComparator()));
 
   BlockHandle block_handle;
   status = FindMetaBlock(meta_iter.get(), meta_block_name, &block_handle);

table/table_test.cc

@@ -238,7 +238,7 @@ class BlockConstructor: public Constructor {
   }
   virtual InternalIterator* NewIterator(
       const SliceTransform* /*prefix_extractor*/) const override {
-    return block_->NewIterator<BlockIter>(comparator_, comparator_);
+    return block_->NewIterator<DataBlockIter>(comparator_, comparator_);
   }
 
  private:
@@ -3474,7 +3474,7 @@ TEST_P(BlockBasedTableTest, PropertiesBlockRestartPointTest) {
                           kDisableGlobalSequenceNumber);
 
     std::unique_ptr<InternalIterator> meta_iter(
-        metaindex_block.NewIterator<BlockIter>(BytewiseComparator(),
+        metaindex_block.NewIterator<DataBlockIter>(BytewiseComparator(),
                                                    BytewiseComparator()));
     bool found_properties_block = true;
     ASSERT_OK(SeekToPropertiesBlock(meta_iter.get(), &found_properties_block));