Fixed the bug when both whole_key_filtering and prefix_extractor are set.

Summary:
When both whole_key_filtering and prefix_extractor are set, RocksDB will
mistakenly encode prefix + whole key into the database instead of
simply whole key when BlockBasedTable is used.  This patch fixes this bug.

Test Plan: Add a test in table_test

Reviewers: anthony, IslamAbdelRahman, rven, kradhakrishnan, sdong

Reviewed By: sdong

Subscribers: dhruba, leveldb

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

table/block_based_filter_block.cc

@@ -19,19 +19,6 @@
 namespace rocksdb {
 
 namespace {
-bool SamePrefix(const SliceTransform* prefix_extractor,
-                const Slice& key1, const Slice& key2) {
-  if (!prefix_extractor->InDomain(key1) &&
-      !prefix_extractor->InDomain(key2)) {
-    return true;
-  } else if (!prefix_extractor->InDomain(key1) ||
-             !prefix_extractor->InDomain(key2)) {
-    return false;
-  } else {
-    return (prefix_extractor->Transform(key1) ==
-            prefix_extractor->Transform(key2));
-  }
-}
 
 void AppendItem(std::string* props, const std::string& key,
                 const std::string& value) {
@@ -78,7 +65,9 @@ BlockBasedFilterBlockBuilder::BlockBasedFilterBlockBuilder(
     const BlockBasedTableOptions& table_opt)
     : policy_(table_opt.filter_policy.get()),
       prefix_extractor_(prefix_extractor),
-      whole_key_filtering_(table_opt.whole_key_filtering) {
+      whole_key_filtering_(table_opt.whole_key_filtering),
+      prev_prefix_start_(0),
+      prev_prefix_size_(0) {
   assert(policy_);
 }
 
@@ -91,14 +80,13 @@ void BlockBasedFilterBlockBuilder::StartBlock(uint64_t block_offset) {
 }
 
 void BlockBasedFilterBlockBuilder::Add(const Slice& key) {
-  added_to_start_ = 0;
-  if (whole_key_filtering_) {
-    AddKey(key);
-    added_to_start_ = 1;
-  }
   if (prefix_extractor_ && prefix_extractor_->InDomain(key)) {
     AddPrefix(key);
   }
+
+  if (whole_key_filtering_) {
+    AddKey(key);
+  }
 }
 
 // Add key to filter if needed
@@ -111,19 +99,16 @@ inline void BlockBasedFilterBlockBuilder::AddKey(const Slice& key) {
 inline void BlockBasedFilterBlockBuilder::AddPrefix(const Slice& key) {
   // get slice for most recently added entry
   Slice prev;
-  if (start_.size() > added_to_start_) {
-    size_t prev_start = start_[start_.size() - 1 - added_to_start_];
-    const char* base = entries_.data() + prev_start;
-    size_t length = entries_.size() - prev_start;
-    prev = Slice(base, length);
+  if (prev_prefix_size_ > 0) {
+    prev = Slice(entries_.data() + prev_prefix_start_, prev_prefix_size_);
   }
 
-  // this assumes prefix(prefix(key)) == prefix(key), as the last
-  // entry in entries_ may be either a key or prefix, and we use
-  // prefix(last entry) to get the prefix of the last key.
-  if (prev.size() == 0 || !SamePrefix(prefix_extractor_, key, prev)) {
-    Slice prefix = prefix_extractor_->Transform(key);
+  Slice prefix = prefix_extractor_->Transform(key);
+  // insert prefix only when it's different from the previous prefix.
+  if (prev.size() == 0 || prefix != prev) {
     start_.push_back(entries_.size());
+    prev_prefix_start_ = entries_.size();
+    prev_prefix_size_ = prefix.size();
     entries_.append(prefix.data(), prefix.size());
   }
 }
@@ -169,6 +154,8 @@ void BlockBasedFilterBlockBuilder::GenerateFilter() {
   tmp_entries_.clear();
   entries_.clear();
   start_.clear();
+  prev_prefix_start_ = 0;
+  prev_prefix_size_ = 0;
 }
 
 BlockBasedFilterBlockReader::BlockBasedFilterBlockReader(

table/block_based_filter_block.h

@@ -55,9 +55,12 @@ class BlockBasedFilterBlockBuilder : public FilterBlockBuilder {
   const SliceTransform* prefix_extractor_;
   bool whole_key_filtering_;
 
+  size_t prev_prefix_start_;        // the position of the last appended prefix
+                                    // to "entries_".
+  size_t prev_prefix_size_;         // the length of the last appended prefix to
+                                    // "entries_".
   std::string entries_;             // Flattened entry contents
   std::vector<size_t> start_;       // Starting index in entries_ of each entry
-  uint32_t added_to_start_;         // To indicate if key is added
   std::string result_;              // Filter data computed so far
   std::vector<Slice> tmp_entries_;  // policy_->CreateFilter() argument
   std::vector<uint32_t> filter_offsets_;

table/table_test.cc

@@ -2281,6 +2281,88 @@ TEST_F(HarnessTest, FooterTests) {
   }
 }
 
+class PrefixTest : public testing::Test {
+ public:
+  PrefixTest() : testing::Test() {}
+  ~PrefixTest() {}
+};
+
+namespace {
+// A simple PrefixExtractor that only works for test PrefixAndWholeKeyTest
+class TestPrefixExtractor : public rocksdb::SliceTransform {
+ public:
+  ~TestPrefixExtractor() override{};
+  const char* Name() const override { return "TestPrefixExtractor"; }
+
+  rocksdb::Slice Transform(const rocksdb::Slice& src) const override {
+    assert(IsValid(src));
+    return rocksdb::Slice(src.data(), 3);
+  }
+
+  bool InDomain(const rocksdb::Slice& src) const override {
+    assert(IsValid(src));
+    return true;
+  }
+
+  bool InRange(const rocksdb::Slice& dst) const override { return true; }
+
+  bool IsValid(const rocksdb::Slice& src) const {
+    if (src.size() != 4) {
+      return false;
+    }
+    if (src[0] != '[') {
+      return false;
+    }
+    if (src[1] < '0' || src[1] > '9') {
+      return false;
+    }
+    if (src[2] != ']') {
+      return false;
+    }
+    if (src[3] < '0' || src[3] > '9') {
+      return false;
+    }
+    return true;
+  }
+};
+}  // namespace
+
+TEST_F(PrefixTest, PrefixAndWholeKeyTest) {
+  rocksdb::Options options;
+  options.compaction_style = rocksdb::kCompactionStyleUniversal;
+  options.num_levels = 20;
+  options.create_if_missing = true;
+  options.optimize_filters_for_hits = false;
+  options.target_file_size_base = 268435456;
+  options.prefix_extractor = std::make_shared<TestPrefixExtractor>();
+  rocksdb::BlockBasedTableOptions bbto;
+  bbto.filter_policy.reset(rocksdb::NewBloomFilterPolicy(10));
+  bbto.block_size = 262144;
+
+  bbto.whole_key_filtering = true;
+
+  const std::string kDBPath = test::TmpDir() + "/prefix_test";
+  options.table_factory.reset(NewBlockBasedTableFactory(bbto));
+  DestroyDB(kDBPath, options);
+  rocksdb::DB* db;
+  ASSERT_OK(rocksdb::DB::Open(options, kDBPath, &db));
+
+  // Create a bunch of keys with 10 filters.
+  for (int i = 0; i < 10; i++) {
+    std::string prefix = "[" + std::to_string(i) + "]";
+    for (int j = 0; j < 10; j++) {
+      std::string key = prefix + std::to_string(j);
+      db->Put(rocksdb::WriteOptions(), key, "1");
+    }
+  }
+
+  // Trigger compaction.
+  db->CompactRange(CompactRangeOptions(), nullptr, nullptr);
+  delete db;
+  // In the second round, turn whole_key_filtering off and expect
+  // rocksdb still works.
+}
+
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {