From 97edc72d398caae97e86a1f41bb1b0b5057f948a Mon Sep 17 00:00:00 2001
From: Sagar Vemuri <svemuri@fb.com>
Date: Tue, 7 Mar 2017 11:50:02 -0800
Subject: [PATCH] Add a memtable-only iterator

Summary:
This PR is to support a way to iterate over all the keys that are just in memtables.
Closes https://github.com/facebook/rocksdb/pull/1953

Differential Revision: D4663500

Pulled By: sagar0

fbshipit-source-id: 144e177
---
 db/db_impl.cc             |  6 ++--
 db/db_test2.cc            | 64 +++++++++++++++++++++++++++++++++++++++
 include/rocksdb/options.h |  3 +-
 3 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/db/db_impl.cc b/db/db_impl.cc
index 9e8582fa8..a6c00fa3c 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -3919,8 +3919,10 @@ InternalIterator* DBImpl::NewInternalIterator(
   }
   if (s.ok()) {
     // Collect iterators for files in L0 - Ln
-    super_version->current->AddIterators(read_options, env_options_,
-                                         &merge_iter_builder, range_del_agg);
+    if (read_options.read_tier != kMemtableTier) {
+      super_version->current->AddIterators(read_options, env_options_,
+                                           &merge_iter_builder, range_del_agg);
+    }
     internal_iter = merge_iter_builder.Finish();
     IterState* cleanup =
         new IterState(this, &mutex_, super_version,
diff --git a/db/db_test2.cc b/db/db_test2.cc
index 44b7b85b9..e9110772e 100644
--- a/db/db_test2.cc
+++ b/db/db_test2.cc
@@ -2298,6 +2298,70 @@ TEST_F(DBTest2, DirectIO) {
   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
   Reopen(options);
 }
+
+TEST_F(DBTest2, MemtableOnlyIterator) {
+  Options options = CurrentOptions();
+  CreateAndReopenWithCF({"pikachu"}, options);
+
+  ASSERT_OK(Put(1, "foo", "first"));
+  ASSERT_OK(Put(1, "bar", "second"));
+
+  ReadOptions ropt;
+  ropt.read_tier = kMemtableTier;
+  std::string value;
+  Iterator* it = nullptr;
+
+  // Before flushing
+  // point lookups
+  ASSERT_OK(db_->Get(ropt, handles_[1], "foo", &value));
+  ASSERT_EQ("first", value);
+  ASSERT_OK(db_->Get(ropt, handles_[1], "bar", &value));
+  ASSERT_EQ("second", value);
+
+  // Memtable-only iterator (read_tier=kMemtableTier); data not flushed yet.
+  it = db_->NewIterator(ropt, handles_[1]);
+  int count = 0;
+  for (it->SeekToFirst(); it->Valid(); it->Next()) {
+    ASSERT_TRUE(it->Valid());
+    count++;
+  }
+  ASSERT_TRUE(!it->Valid());
+  ASSERT_EQ(2, count);
+  delete it;
+
+  Flush(1);
+
+  // After flushing
+  // point lookups
+  ASSERT_OK(db_->Get(ropt, handles_[1], "foo", &value));
+  ASSERT_EQ("first", value);
+  ASSERT_OK(db_->Get(ropt, handles_[1], "bar", &value));
+  ASSERT_EQ("second", value);
+  // nothing should be returned using memtable-only iterator after flushing.
+  it = db_->NewIterator(ropt, handles_[1]);
+  count = 0;
+  for (it->SeekToFirst(); it->Valid(); it->Next()) {
+    ASSERT_TRUE(it->Valid());
+    count++;
+  }
+  ASSERT_TRUE(!it->Valid());
+  ASSERT_EQ(0, count);
+  delete it;
+
+  // Add a key to memtable
+  ASSERT_OK(Put(1, "foobar", "third"));
+  it = db_->NewIterator(ropt, handles_[1]);
+  count = 0;
+  for (it->SeekToFirst(); it->Valid(); it->Next()) {
+    ASSERT_TRUE(it->Valid());
+    ASSERT_EQ("foobar", it->key().ToString());
+    ASSERT_EQ("third", it->value().ToString());
+    count++;
+  }
+  ASSERT_TRUE(!it->Valid());
+  ASSERT_EQ(1, count);
+  delete it;
+}
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {
diff --git a/include/rocksdb/options.h b/include/rocksdb/options.h
index 9cfd430fd..ffcce4d9e 100644
--- a/include/rocksdb/options.h
+++ b/include/rocksdb/options.h
@@ -888,10 +888,11 @@ struct Options : public DBOptions, public ColumnFamilyOptions {
 enum ReadTier {
   kReadAllTier = 0x0,     // data in memtable, block cache, OS cache or storage
   kBlockCacheTier = 0x1,  // data in memtable or block cache
-  kPersistedTier = 0x2    // persisted data.  When WAL is disabled, this option
+  kPersistedTier = 0x2,   // persisted data.  When WAL is disabled, this option
                           // will skip data in memtable.
                           // Note that this ReadTier currently only supports
                           // Get and MultiGet and does not support iterators.
+  kMemtableTier = 0x3     // data in memtable. used for memtable-only iterators.
 };
 
 // Options that control read operations