From 8ada876dfe2356a9f078ce17aa51057eeb8f17b9 Mon Sep 17 00:00:00 2001
From: Sagar Vemuri <svemuri@fb.com>
Date: Tue, 20 Feb 2018 19:05:21 -0800
Subject: [PATCH] Add rocksdb.iterator.internal-key property

Summary:
Added a new iterator property: `rocksdb.iterator.internal-key` to get the internal-key (converted to user key) at which the iterator stopped.
Closes https://github.com/facebook/rocksdb/pull/3525

Differential Revision: D7033694

Pulled By: sagar0

fbshipit-source-id: d51e6c00f5e9d766c6276ef79774b81c6c5216f8
---
 db/db_iter.cc              |  3 +++
 db/db_iterator_test.cc     | 49 ++++++++++++++++++++++++++++++++++++++
 include/rocksdb/iterator.h |  3 +++
 3 files changed, 55 insertions(+)

diff --git a/db/db_iter.cc b/db/db_iter.cc
index 9f2d283fbf..29b538de95 100644
--- a/db/db_iter.cc
+++ b/db/db_iter.cc
@@ -217,6 +217,9 @@ class DBIter final: public Iterator {
         *prop = "Iterator is not valid.";
       }
       return Status::OK();
+    } else if (prop_name == "rocksdb.iterator.internal-key") {
+      *prop = saved_key_.GetUserKey().ToString();
+      return Status::OK();
     }
     return Status::InvalidArgument("Undentified property.");
   }
diff --git a/db/db_iterator_test.cc b/db/db_iterator_test.cc
index e59fe3b7e8..a0ecbb6816 100644
--- a/db/db_iterator_test.cc
+++ b/db/db_iterator_test.cc
@@ -81,6 +81,7 @@ TEST_P(DBIteratorTest, IteratorProperty) {
   Options options = CurrentOptions();
   CreateAndReopenWithCF({"pikachu"}, options);
   Put(1, "1", "2");
+  Delete(1, "2");
   ReadOptions ropt;
   ropt.pin_data = false;
   {
@@ -90,9 +91,15 @@ TEST_P(DBIteratorTest, IteratorProperty) {
     ASSERT_NOK(iter->GetProperty("non_existing.value", &prop_value));
     ASSERT_OK(iter->GetProperty("rocksdb.iterator.is-key-pinned", &prop_value));
     ASSERT_EQ("0", prop_value);
+    ASSERT_OK(iter->GetProperty("rocksdb.iterator.internal-key", &prop_value));
+    ASSERT_EQ("1", prop_value);
     iter->Next();
     ASSERT_OK(iter->GetProperty("rocksdb.iterator.is-key-pinned", &prop_value));
     ASSERT_EQ("Iterator is not valid.", prop_value);
+
+    // Get internal key at which the iteration stopped (tombstone in this case).
+    ASSERT_OK(iter->GetProperty("rocksdb.iterator.internal-key", &prop_value));
+    ASSERT_EQ("2", prop_value);
   }
   Close();
 }
@@ -2183,6 +2190,48 @@ TEST_P(DBIteratorTest, SkipStatistics) {
   ASSERT_EQ(skip_count, TestGetTickerCount(options, NUMBER_ITER_SKIP));
 }
 
+TEST_P(DBIteratorTest, SeekAfterHittingManyInternalKeys) {
+  Options options = CurrentOptions();
+  DestroyAndReopen(options);
+  ReadOptions ropts;
+  ropts.max_skippable_internal_keys = 2;
+
+  Put("1", "val_1");
+  // Add more tombstones than max_skippable_internal_keys so that Next() fails.
+  Delete("2");
+  Delete("3");
+  Delete("4");
+  Delete("5");
+  Put("6", "val_6");
+
+  unique_ptr<Iterator> iter(NewIterator(ropts));
+  iter->SeekToFirst();
+
+  ASSERT_TRUE(iter->Valid());
+  ASSERT_EQ(iter->key().ToString(), "1");
+  ASSERT_EQ(iter->value().ToString(), "val_1");
+
+  // This should fail as incomplete due to too many non-visible internal keys on
+  // the way to the next valid user key.
+  iter->Next();
+  ASSERT_TRUE(!iter->Valid());
+  ASSERT_TRUE(iter->status().IsIncomplete());
+
+  // Get the internal key at which Next() failed.
+  std::string prop_value;
+  ASSERT_OK(iter->GetProperty("rocksdb.iterator.internal-key", &prop_value));
+  ASSERT_EQ("4", prop_value);
+
+  // Create a new iterator to seek to the internal key.
+  unique_ptr<Iterator> iter2(NewIterator(ropts));
+  iter2->Seek(prop_value);
+  ASSERT_TRUE(iter2->Valid());
+  ASSERT_OK(iter2->status());
+
+  ASSERT_EQ(iter2->key().ToString(), "6");
+  ASSERT_EQ(iter2->value().ToString(), "val_6");
+}
+
 INSTANTIATE_TEST_CASE_P(DBIteratorTestInstance, DBIteratorTest,
                         testing::Values(true, false));
 
diff --git a/include/rocksdb/iterator.h b/include/rocksdb/iterator.h
index e3e4f7ce7d..14f0ae0ab2 100644
--- a/include/rocksdb/iterator.h
+++ b/include/rocksdb/iterator.h
@@ -97,6 +97,9 @@ class Iterator : public Cleanable {
   // Property "rocksdb.iterator.super-version-number":
   //   LSM version used by the iterator. The same format as DB Property
   //   kCurrentSuperVersionNumber. See its comment for more information.
+  // Property "rocksdb.iterator.internal-key":
+  //   Get the user-key portion of the internal key at which the iteration
+  //   stopped.
   virtual Status GetProperty(std::string prop_name, std::string* prop);
 
  private: