From 12ba00ea65b073f388c18e84c29d7af6d201d467 Mon Sep 17 00:00:00 2001
From: Aaron Gao <gzh@fb.com>
Date: Wed, 8 Mar 2017 17:12:04 -0800
Subject: [PATCH] Reset DBIter::saved_key_ with proper user key anywhere before
 pass to DBIter::FindNextUserEntry

Summary:
fix db_iter bug introduced by [facebook#1413](https://github.com/facebook/rocksdb/pull/1413)
Closes https://github.com/facebook/rocksdb/pull/1962

Differential Revision: D4672369

Pulled By: lightmark

fbshipit-source-id: 6a22953
---
 db/db_iter.cc      |  7 +++++-
 db/db_iter_test.cc | 60 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/db/db_iter.cc b/db/db_iter.cc
index 54d8351a8..e46f231fa 100644
--- a/db/db_iter.cc
+++ b/db/db_iter.cc
@@ -341,6 +341,7 @@ void DBIter::Next() {
 //
 // NOTE: In between, saved_key_ can point to a user key that has
 //       a delete marker or a sequence number higher than sequence_
+//       saved_key_ MUST have a proper user_key before calling this function
 //
 // The prefix_check parameter controls whether we check the iterated
 // keys against the prefix of the seeked key. Set to false when
@@ -946,7 +947,6 @@ void DBIter::Seek(const Slice& target) {
   StopWatch sw(env_, statistics_, DB_SEEK);
   ReleaseTempPinnedData();
   saved_key_.Clear();
-  // now saved_key is used to store internal key.
   saved_key_.SetInternalKey(target, sequence_);
 
   {
@@ -961,6 +961,9 @@ void DBIter::Seek(const Slice& target) {
     }
     direction_ = kForward;
     ClearSavedValue();
+    // convert the InternalKey to UserKey in saved_key_ before
+    // passed to FindNextUserEntry
+    saved_key_.Reserve(saved_key_.Size() - 8);
     FindNextUserEntry(false /* not skipping */, prefix_same_as_start_);
     if (!valid_) {
       prefix_start_key_.clear();
@@ -1039,6 +1042,8 @@ void DBIter::SeekToFirst() {
 
   RecordTick(statistics_, NUMBER_DB_SEEK);
   if (iter_->Valid()) {
+    saved_key_.SetKey(ExtractUserKey(iter_->key()),
+                      !iter_->IsKeyPinned() || !pin_thru_lifetime_ /* copy */);
     FindNextUserEntry(false /* not skipping */, false /* no prefix check */);
     if (statistics_ != nullptr) {
       if (valid_) {
diff --git a/db/db_iter_test.cc b/db/db_iter_test.cc
index c8abd0d25..f6492bd3a 100644
--- a/db/db_iter_test.cc
+++ b/db/db_iter_test.cc
@@ -1897,6 +1897,66 @@ TEST_F(DBIteratorTest, DBIterator12) {
   ASSERT_FALSE(db_iter->Valid());
 }
 
+TEST_F(DBIteratorTest, DBIterator13) {
+  Options options;
+  options.merge_operator = nullptr;
+
+  std::string key;
+  key.resize(9);
+  key.assign(9, static_cast<char>(0));
+  key[0] = 'b';
+
+  TestIterator* internal_iter = new TestIterator(BytewiseComparator());
+  internal_iter->AddPut(key, "0");
+  internal_iter->AddPut(key, "1");
+  internal_iter->AddPut(key, "2");
+  internal_iter->AddPut(key, "3");
+  internal_iter->AddPut(key, "4");
+  internal_iter->AddPut(key, "5");
+  internal_iter->AddPut(key, "6");
+  internal_iter->AddPut(key, "7");
+  internal_iter->AddPut(key, "8");
+  internal_iter->Finish();
+
+  std::unique_ptr<Iterator> db_iter(
+      NewDBIterator(env_, ImmutableCFOptions(options), BytewiseComparator(),
+                    internal_iter, 2, 3, 0));
+  db_iter->Seek("b");
+  ASSERT_TRUE(db_iter->Valid());
+  ASSERT_EQ(db_iter->key().ToString(), key);
+  ASSERT_EQ(db_iter->value().ToString(), "2");
+}
+
+TEST_F(DBIteratorTest, DBIterator14) {
+  Options options;
+  options.merge_operator = nullptr;
+
+  std::string key("b");
+  TestIterator* internal_iter = new TestIterator(BytewiseComparator());
+  internal_iter->AddPut("b", "0");
+  internal_iter->AddPut("b", "1");
+  internal_iter->AddPut("b", "2");
+  internal_iter->AddPut("b", "3");
+  internal_iter->AddPut("a", "4");
+  internal_iter->AddPut("a", "5");
+  internal_iter->AddPut("a", "6");
+  internal_iter->AddPut("c", "7");
+  internal_iter->AddPut("c", "8");
+  internal_iter->AddPut("c", "9");
+  internal_iter->Finish();
+
+  std::unique_ptr<Iterator> db_iter(
+      NewDBIterator(env_, ImmutableCFOptions(options), BytewiseComparator(),
+                    internal_iter, 4, 1, 0));
+  db_iter->Seek("b");
+  ASSERT_TRUE(db_iter->Valid());
+  ASSERT_EQ(db_iter->key().ToString(), "b");
+  ASSERT_EQ(db_iter->value().ToString(), "3");
+  db_iter->SeekToFirst();
+  ASSERT_EQ(db_iter->key().ToString(), "a");
+  ASSERT_EQ(db_iter->value().ToString(), "4");
+}
+
 class DBIterWithMergeIterTest : public testing::Test {
  public:
   DBIterWithMergeIterTest()