[RocksDB] Generalize prefix-aware iterator to be used for more than one Seek

Summary: Added a prefix_seek flag in ReadOptions to indicate that Seek is prefix aware(might not return data with different prefix), and also not bound to a specific prefix. Multiple Seeks and range scans can be invoked on the same iterator. If a specific prefix is specified, this flag will be ignored. Just a quick prototype that works for PrefixHashRep, the new lockless memtable could be easily extended with this support too.

Test Plan: test it on Leaf

Reviewers: dhruba, kailiu, sdong, igor

Reviewed By: igor

CC: leveldb

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

db/db_impl.cc

@@ -2454,7 +2454,7 @@ Iterator* DBImpl::NewInternalIterator(const ReadOptions& options,
   // Collect together all needed child iterators for mem
   std::vector<Iterator*> list;
   mem_->Ref();
-  list.push_back(mem_->NewIterator(options.prefix));
+  list.push_back(mem_->NewIterator(options));
 
   cleanup->mem.push_back(mem_);
 
@@ -2464,7 +2464,7 @@ Iterator* DBImpl::NewInternalIterator(const ReadOptions& options,
   for (unsigned int i = 0; i < immutables.size(); i++) {
     MemTable* m = immutables[i];
     m->Ref();
-    list.push_back(m->NewIterator(options.prefix));
+    list.push_back(m->NewIterator(options));
     cleanup->mem.push_back(m);
   }
 

db/memtable.cc

@@ -84,11 +84,16 @@ static const char* EncodeKey(std::string* scratch, const Slice& target) {
 
 class MemTableIterator: public Iterator {
  public:
-  explicit MemTableIterator(MemTableRep* table)
-    : iter_(table->GetIterator()) { }
-
-  MemTableIterator(MemTableRep* table, const Slice* prefix)
-    : iter_(table->GetPrefixIterator(*prefix)) { }
+  MemTableIterator(MemTableRep* table, const ReadOptions& options)
+    : iter_() {
+    if (options.prefix) {
+      iter_ = table->GetPrefixIterator(*options.prefix);
+    } else if (options.prefix_seek) {
+      iter_ = table->GetDynamicPrefixIterator();
+    } else {
+      iter_ = table->GetIterator();
+    }
+  }
 
   virtual bool Valid() const { return iter_->Valid(); }
   virtual void Seek(const Slice& k) { iter_->Seek(EncodeKey(&tmp_, k)); }
@@ -115,12 +120,8 @@ class MemTableIterator: public Iterator {
   void operator=(const MemTableIterator&);
 };
 
-Iterator* MemTable::NewIterator(const Slice* prefix) {
-  if (prefix) {
-    return new MemTableIterator(table_.get(), prefix);
-  } else {
-    return new MemTableIterator(table_.get());
-  }
+Iterator* MemTable::NewIterator(const ReadOptions& options) {
+  return new MemTableIterator(table_.get(), options);
 }
 
 port::RWMutex* MemTable::GetLock(const Slice& key) {

db/memtable.h

@@ -65,10 +65,14 @@ class MemTable {
   // iterator are internal keys encoded by AppendInternalKey in the
   // db/dbformat.{h,cc} module.
   //
-  // If a prefix is supplied, it is passed to the underlying MemTableRep as a
-  // hint that the iterator only need to support access to keys with that
-  // prefix.
-  Iterator* NewIterator(const Slice* prefix = nullptr);
+  // If options.prefix is supplied, it is passed to the underlying MemTableRep
+  // as a hint that the iterator only need to support access to keys with that
+  // specific prefix.
+  // If options.prefix is not supplied and options.prefix_seek is set, the
+  // iterator is not bound to a specific prefix. However, the semantics of
+  // Seek is changed - the result might only include keys with the same prefix
+  // as the seek-key.
+  Iterator* NewIterator(const ReadOptions& options = ReadOptions());
 
   // Add an entry into memtable that maps key to value at the
   // specified sequence number and with the specified type.

db/prefix_test.cc

@@ -75,8 +75,6 @@ class TestKeyComparator : public Comparator {
 
   virtual void FindShortSuccessor(std::string* key) const {}
 
- private:
-
 };
 
 class PrefixTest {
@@ -116,6 +114,93 @@ class PrefixTest {
   Options options;
 };
 
+TEST(PrefixTest, DynamicPrefixIterator) {
+
+  DestroyDB(kDbName, Options());
+  auto db = OpenDb();
+  WriteOptions write_options;
+  ReadOptions read_options;
+
+  std::vector<uint64_t> prefixes;
+  for (uint64_t i = 0; i < FLAGS_total_prefixes; ++i) {
+    prefixes.push_back(i);
+  }
+
+  if (FLAGS_random_prefix) {
+    std::random_shuffle(prefixes.begin(), prefixes.end());
+  }
+
+  // insert x random prefix, each with y continuous element.
+  for (auto prefix : prefixes) {
+     for (uint64_t sorted = 0; sorted < FLAGS_items_per_prefix; sorted++) {
+      TestKey test_key(prefix, sorted);
+
+      Slice key = TestKeyToSlice(test_key);
+      std::string value = "v" + std::to_string(sorted);
+
+      ASSERT_OK(db->Put(write_options, key, value));
+    }
+  }
+
+  // test seek existing keys
+  HistogramImpl hist_seek_time;
+  HistogramImpl hist_seek_comparison;
+
+  if (FLAGS_use_prefix_hash_memtable) {
+    read_options.prefix_seek = true;
+  }
+  std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
+
+  for (auto prefix : prefixes) {
+    TestKey test_key(prefix, FLAGS_items_per_prefix / 2);
+    Slice key = TestKeyToSlice(test_key);
+    std::string value = "v" + std::to_string(0);
+
+    perf_context.Reset();
+    StopWatchNano timer(Env::Default(), true);
+    uint64_t total_keys = 0;
+    for (iter->Seek(key); iter->Valid(); iter->Next()) {
+      if (FLAGS_trigger_deadlock) {
+        std::cout << "Behold the deadlock!\n";
+        db->Delete(write_options, iter->key());
+      }
+      auto test_key = SliceToTestKey(iter->key());
+      if (test_key->prefix != prefix) break;
+      total_keys++;
+    }
+    hist_seek_time.Add(timer.ElapsedNanos());
+    hist_seek_comparison.Add(perf_context.user_key_comparison_count);
+    ASSERT_EQ(total_keys, FLAGS_items_per_prefix - FLAGS_items_per_prefix/2);
+  }
+
+  std::cout << "Seek key comparison: \n"
+            << hist_seek_comparison.ToString()
+            << "Seek time: \n"
+            << hist_seek_time.ToString();
+
+  // test non-existing keys
+  HistogramImpl hist_no_seek_time;
+  HistogramImpl hist_no_seek_comparison;
+
+  for (auto prefix = FLAGS_total_prefixes;
+       prefix < FLAGS_total_prefixes + 100;
+       prefix++) {
+    TestKey test_key(prefix, 0);
+    Slice key = TestKeyToSlice(test_key);
+
+    perf_context.Reset();
+    StopWatchNano timer(Env::Default(), true);
+    iter->Seek(key);
+    hist_no_seek_time.Add(timer.ElapsedNanos());
+    hist_no_seek_comparison.Add(perf_context.user_key_comparison_count);
+    ASSERT_TRUE(!iter->Valid());
+  }
+
+  std::cout << "non-existing Seek key comparison: \n"
+            << hist_no_seek_comparison.ToString()
+            << "non-existing Seek time: \n"
+            << hist_no_seek_time.ToString();
+}
 
 TEST(PrefixTest, PrefixHash) {