1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 
6 #ifndef ROCKSDB_LITE
7 
8 #ifndef GFLAGS
9 #include <cstdio>
main()10 int main() {
11   fprintf(stderr, "Please install gflags to run this test... Skipping...\n");
12   return 0;
13 }
14 #else
15 
16 #include <algorithm>
17 #include <iostream>
18 #include <vector>
19 
20 #include "db/db_impl/db_impl.h"
21 #include "monitoring/histogram.h"
22 #include "rocksdb/comparator.h"
23 #include "rocksdb/db.h"
24 #include "rocksdb/filter_policy.h"
25 #include "rocksdb/memtablerep.h"
26 #include "rocksdb/perf_context.h"
27 #include "rocksdb/slice_transform.h"
28 #include "rocksdb/table.h"
29 #include "test_util/testharness.h"
30 #include "util/coding.h"
31 #include "util/gflags_compat.h"
32 #include "util/random.h"
33 #include "util/stop_watch.h"
34 #include "util/string_util.h"
35 #include "utilities/merge_operators.h"
36 
37 using GFLAGS_NAMESPACE::ParseCommandLineFlags;
38 
39 DEFINE_bool(trigger_deadlock, false,
40             "issue delete in range scan to trigger PrefixHashMap deadlock");
41 DEFINE_int32(bucket_count, 100000, "number of buckets");
42 DEFINE_uint64(num_locks, 10001, "number of locks");
43 DEFINE_bool(random_prefix, false, "randomize prefix");
44 DEFINE_uint64(total_prefixes, 100000, "total number of prefixes");
45 DEFINE_uint64(items_per_prefix, 1, "total number of values per prefix");
46 DEFINE_int64(write_buffer_size, 33554432, "");
47 DEFINE_int32(max_write_buffer_number, 2, "");
48 DEFINE_int32(min_write_buffer_number_to_merge, 1, "");
49 DEFINE_int32(skiplist_height, 4, "");
50 DEFINE_double(memtable_prefix_bloom_size_ratio, 0.1, "");
51 DEFINE_int32(memtable_huge_page_size, 2 * 1024 * 1024, "");
52 DEFINE_int32(value_size, 40, "");
53 DEFINE_bool(enable_print, false, "Print options generated to console.");
54 
55 // Path to the database on file system
56 const std::string kDbName =
57     ROCKSDB_NAMESPACE::test::PerThreadDBPath("prefix_test");
58 
59 namespace ROCKSDB_NAMESPACE {
60 
61 struct TestKey {
62   uint64_t prefix;
63   uint64_t sorted;
64 
TestKeyROCKSDB_NAMESPACE::TestKey65   TestKey(uint64_t _prefix, uint64_t _sorted)
66       : prefix(_prefix), sorted(_sorted) {}
67 };
68 
69 // return a slice backed by test_key
TestKeyToSlice(std::string & s,const TestKey & test_key)70 inline Slice TestKeyToSlice(std::string &s, const TestKey& test_key) {
71   s.clear();
72   PutFixed64(&s, test_key.prefix);
73   PutFixed64(&s, test_key.sorted);
74   return Slice(s.c_str(), s.size());
75 }
76 
SliceToTestKey(const Slice & slice)77 inline const TestKey SliceToTestKey(const Slice& slice) {
78   return TestKey(DecodeFixed64(slice.data()),
79     DecodeFixed64(slice.data() + 8));
80 }
81 
82 class TestKeyComparator : public Comparator {
83  public:
84 
85   // Compare needs to be aware of the possibility of a and/or b is
86   // prefix only
Compare(const Slice & a,const Slice & b) const87   int Compare(const Slice& a, const Slice& b) const override {
88     const TestKey kkey_a = SliceToTestKey(a);
89     const TestKey kkey_b = SliceToTestKey(b);
90     const TestKey *key_a = &kkey_a;
91     const TestKey *key_b = &kkey_b;
92     if (key_a->prefix != key_b->prefix) {
93       if (key_a->prefix < key_b->prefix) return -1;
94       if (key_a->prefix > key_b->prefix) return 1;
95     } else {
96       EXPECT_TRUE(key_a->prefix == key_b->prefix);
97       // note, both a and b could be prefix only
98       if (a.size() != b.size()) {
99         // one of them is prefix
100         EXPECT_TRUE(
101             (a.size() == sizeof(uint64_t) && b.size() == sizeof(TestKey)) ||
102             (b.size() == sizeof(uint64_t) && a.size() == sizeof(TestKey)));
103         if (a.size() < b.size()) return -1;
104         if (a.size() > b.size()) return 1;
105       } else {
106         // both a and b are prefix
107         if (a.size() == sizeof(uint64_t)) {
108           return 0;
109         }
110 
111         // both a and b are whole key
112         EXPECT_TRUE(a.size() == sizeof(TestKey) && b.size() == sizeof(TestKey));
113         if (key_a->sorted < key_b->sorted) return -1;
114         if (key_a->sorted > key_b->sorted) return 1;
115         if (key_a->sorted == key_b->sorted) return 0;
116       }
117     }
118     return 0;
119   }
120 
operator ()(const TestKey & a,const TestKey & b) const121   bool operator()(const TestKey& a, const TestKey& b) const {
122     std::string sa, sb;
123     return Compare(TestKeyToSlice(sa, a), TestKeyToSlice(sb, b)) < 0;
124   }
125 
Name() const126   const char* Name() const override { return "TestKeyComparator"; }
127 
FindShortestSeparator(std::string *,const Slice &) const128   void FindShortestSeparator(std::string* /*start*/,
129                              const Slice& /*limit*/) const override {}
130 
FindShortSuccessor(std::string *) const131   void FindShortSuccessor(std::string* /*key*/) const override {}
132 };
133 
134 namespace {
PutKey(DB * db,WriteOptions write_options,uint64_t prefix,uint64_t suffix,const Slice & value)135 void PutKey(DB* db, WriteOptions write_options, uint64_t prefix,
136             uint64_t suffix, const Slice& value) {
137   TestKey test_key(prefix, suffix);
138   std::string s;
139   Slice key = TestKeyToSlice(s, test_key);
140   ASSERT_OK(db->Put(write_options, key, value));
141 }
142 
PutKey(DB * db,WriteOptions write_options,const TestKey & test_key,const Slice & value)143 void PutKey(DB* db, WriteOptions write_options, const TestKey& test_key,
144             const Slice& value) {
145   std::string s;
146   Slice key = TestKeyToSlice(s, test_key);
147   ASSERT_OK(db->Put(write_options, key, value));
148 }
149 
MergeKey(DB * db,WriteOptions write_options,const TestKey & test_key,const Slice & value)150 void MergeKey(DB* db, WriteOptions write_options, const TestKey& test_key,
151               const Slice& value) {
152   std::string s;
153   Slice key = TestKeyToSlice(s, test_key);
154   ASSERT_OK(db->Merge(write_options, key, value));
155 }
156 
DeleteKey(DB * db,WriteOptions write_options,const TestKey & test_key)157 void DeleteKey(DB* db, WriteOptions write_options, const TestKey& test_key) {
158   std::string s;
159   Slice key = TestKeyToSlice(s, test_key);
160   ASSERT_OK(db->Delete(write_options, key));
161 }
162 
SeekIterator(Iterator * iter,uint64_t prefix,uint64_t suffix)163 void SeekIterator(Iterator* iter, uint64_t prefix, uint64_t suffix) {
164   TestKey test_key(prefix, suffix);
165   std::string s;
166   Slice key = TestKeyToSlice(s, test_key);
167   iter->Seek(key);
168 }
169 
170 const std::string kNotFoundResult = "NOT_FOUND";
171 
Get(DB * db,const ReadOptions & read_options,uint64_t prefix,uint64_t suffix)172 std::string Get(DB* db, const ReadOptions& read_options, uint64_t prefix,
173                 uint64_t suffix) {
174   TestKey test_key(prefix, suffix);
175   std::string s2;
176   Slice key = TestKeyToSlice(s2, test_key);
177 
178   std::string result;
179   Status s = db->Get(read_options, key, &result);
180   if (s.IsNotFound()) {
181     result = kNotFoundResult;
182   } else if (!s.ok()) {
183     result = s.ToString();
184   }
185   return result;
186 }
187 
188 class SamePrefixTransform : public SliceTransform {
189  private:
190   const Slice prefix_;
191   std::string name_;
192 
193  public:
SamePrefixTransform(const Slice & prefix)194   explicit SamePrefixTransform(const Slice& prefix)
195       : prefix_(prefix), name_("rocksdb.SamePrefix." + prefix.ToString()) {}
196 
Name() const197   const char* Name() const override { return name_.c_str(); }
198 
Transform(const Slice & src) const199   Slice Transform(const Slice& src) const override {
200     assert(InDomain(src));
201     return prefix_;
202   }
203 
InDomain(const Slice & src) const204   bool InDomain(const Slice& src) const override {
205     if (src.size() >= prefix_.size()) {
206       return Slice(src.data(), prefix_.size()) == prefix_;
207     }
208     return false;
209   }
210 
InRange(const Slice & dst) const211   bool InRange(const Slice& dst) const override { return dst == prefix_; }
212 
FullLengthEnabled(size_t *) const213   bool FullLengthEnabled(size_t* /*len*/) const override { return false; }
214 };
215 
216 }  // namespace
217 
218 class PrefixTest : public testing::Test {
219  public:
OpenDb()220   std::shared_ptr<DB> OpenDb() {
221     DB* db;
222 
223     options.create_if_missing = true;
224     options.write_buffer_size = FLAGS_write_buffer_size;
225     options.max_write_buffer_number = FLAGS_max_write_buffer_number;
226     options.min_write_buffer_number_to_merge =
227       FLAGS_min_write_buffer_number_to_merge;
228 
229     options.memtable_prefix_bloom_size_ratio =
230         FLAGS_memtable_prefix_bloom_size_ratio;
231     options.memtable_huge_page_size = FLAGS_memtable_huge_page_size;
232 
233     options.prefix_extractor.reset(NewFixedPrefixTransform(8));
234     BlockBasedTableOptions bbto;
235     bbto.filter_policy.reset(NewBloomFilterPolicy(10, false));
236     bbto.whole_key_filtering = false;
237     options.table_factory.reset(NewBlockBasedTableFactory(bbto));
238     options.allow_concurrent_memtable_write = false;
239 
240     Status s = DB::Open(options, kDbName,  &db);
241     EXPECT_OK(s);
242     return std::shared_ptr<DB>(db);
243   }
244 
FirstOption()245   void FirstOption() {
246     option_config_ = kBegin;
247   }
248 
NextOptions(int bucket_count)249   bool NextOptions(int bucket_count) {
250     // skip some options
251     option_config_++;
252     if (option_config_ < kEnd) {
253       options.prefix_extractor.reset(NewFixedPrefixTransform(8));
254       switch(option_config_) {
255         case kHashSkipList:
256           options.memtable_factory.reset(
257               NewHashSkipListRepFactory(bucket_count, FLAGS_skiplist_height));
258           return true;
259         case kHashLinkList:
260           options.memtable_factory.reset(
261               NewHashLinkListRepFactory(bucket_count));
262           return true;
263         case kHashLinkListHugePageTlb:
264           options.memtable_factory.reset(
265               NewHashLinkListRepFactory(bucket_count, 2 * 1024 * 1024));
266           return true;
267         case kHashLinkListTriggerSkipList:
268           options.memtable_factory.reset(
269               NewHashLinkListRepFactory(bucket_count, 0, 3));
270           return true;
271         default:
272           return false;
273       }
274     }
275     return false;
276   }
277 
PrefixTest()278   PrefixTest() : option_config_(kBegin) {
279     options.comparator = new TestKeyComparator();
280   }
~PrefixTest()281   ~PrefixTest() override { delete options.comparator; }
282 
283  protected:
284   enum OptionConfig {
285     kBegin,
286     kHashSkipList,
287     kHashLinkList,
288     kHashLinkListHugePageTlb,
289     kHashLinkListTriggerSkipList,
290     kEnd
291   };
292   int option_config_;
293   Options options;
294 };
295 
TEST(SamePrefixTest,InDomainTest)296 TEST(SamePrefixTest, InDomainTest) {
297   DB* db;
298   Options options;
299   options.create_if_missing = true;
300   options.prefix_extractor.reset(new SamePrefixTransform("HHKB"));
301   BlockBasedTableOptions bbto;
302   bbto.filter_policy.reset(NewBloomFilterPolicy(10, false));
303   bbto.whole_key_filtering = false;
304   options.table_factory.reset(NewBlockBasedTableFactory(bbto));
305   WriteOptions write_options;
306   ReadOptions read_options;
307   {
308     ASSERT_OK(DestroyDB(kDbName, Options()));
309     ASSERT_OK(DB::Open(options, kDbName, &db));
310     ASSERT_OK(db->Put(write_options, "HHKB pro2", "Mar 24, 2006"));
311     ASSERT_OK(db->Put(write_options, "HHKB pro2 Type-S", "June 29, 2011"));
312     ASSERT_OK(db->Put(write_options, "Realforce 87u", "idk"));
313     db->Flush(FlushOptions());
314     std::string result;
315     auto db_iter = db->NewIterator(ReadOptions());
316 
317     db_iter->Seek("Realforce 87u");
318     ASSERT_TRUE(db_iter->Valid());
319     ASSERT_OK(db_iter->status());
320     ASSERT_EQ(db_iter->key(), "Realforce 87u");
321     ASSERT_EQ(db_iter->value(), "idk");
322 
323     delete db_iter;
324     delete db;
325     ASSERT_OK(DestroyDB(kDbName, Options()));
326   }
327 
328   {
329     ASSERT_OK(DB::Open(options, kDbName, &db));
330     ASSERT_OK(db->Put(write_options, "pikachu", "1"));
331     ASSERT_OK(db->Put(write_options, "Meowth", "1"));
332     ASSERT_OK(db->Put(write_options, "Mewtwo", "idk"));
333     db->Flush(FlushOptions());
334     std::string result;
335     auto db_iter = db->NewIterator(ReadOptions());
336 
337     db_iter->Seek("Mewtwo");
338     ASSERT_TRUE(db_iter->Valid());
339     ASSERT_OK(db_iter->status());
340     delete db_iter;
341     delete db;
342     ASSERT_OK(DestroyDB(kDbName, Options()));
343   }
344 }
345 
TEST_F(PrefixTest,TestResult)346 TEST_F(PrefixTest, TestResult) {
347   for (int num_buckets = 1; num_buckets <= 2; num_buckets++) {
348     FirstOption();
349     while (NextOptions(num_buckets)) {
350       std::cout << "*** Mem table: " << options.memtable_factory->Name()
351                 << " number of buckets: " << num_buckets
352                 << std::endl;
353       DestroyDB(kDbName, Options());
354       auto db = OpenDb();
355       WriteOptions write_options;
356       ReadOptions read_options;
357 
358       // 1. Insert one row.
359       Slice v16("v16");
360       PutKey(db.get(), write_options, 1, 6, v16);
361       std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
362       SeekIterator(iter.get(), 1, 6);
363       ASSERT_TRUE(iter->Valid());
364       ASSERT_TRUE(v16 == iter->value());
365       SeekIterator(iter.get(), 1, 5);
366       ASSERT_TRUE(iter->Valid());
367       ASSERT_TRUE(v16 == iter->value());
368       SeekIterator(iter.get(), 1, 5);
369       ASSERT_TRUE(iter->Valid());
370       ASSERT_TRUE(v16 == iter->value());
371       iter->Next();
372       ASSERT_TRUE(!iter->Valid());
373 
374       SeekIterator(iter.get(), 2, 0);
375       ASSERT_TRUE(!iter->Valid());
376 
377       ASSERT_EQ(v16.ToString(), Get(db.get(), read_options, 1, 6));
378       ASSERT_EQ(kNotFoundResult, Get(db.get(), read_options, 1, 5));
379       ASSERT_EQ(kNotFoundResult, Get(db.get(), read_options, 1, 7));
380       ASSERT_EQ(kNotFoundResult, Get(db.get(), read_options, 0, 6));
381       ASSERT_EQ(kNotFoundResult, Get(db.get(), read_options, 2, 6));
382 
383       // 2. Insert an entry for the same prefix as the last entry in the bucket.
384       Slice v17("v17");
385       PutKey(db.get(), write_options, 1, 7, v17);
386       iter.reset(db->NewIterator(read_options));
387       SeekIterator(iter.get(), 1, 7);
388       ASSERT_TRUE(iter->Valid());
389       ASSERT_TRUE(v17 == iter->value());
390 
391       SeekIterator(iter.get(), 1, 6);
392       ASSERT_TRUE(iter->Valid());
393       ASSERT_TRUE(v16 == iter->value());
394       iter->Next();
395       ASSERT_TRUE(iter->Valid());
396       ASSERT_TRUE(v17 == iter->value());
397       iter->Next();
398       ASSERT_TRUE(!iter->Valid());
399 
400       SeekIterator(iter.get(), 2, 0);
401       ASSERT_TRUE(!iter->Valid());
402 
403       // 3. Insert an entry for the same prefix as the head of the bucket.
404       Slice v15("v15");
405       PutKey(db.get(), write_options, 1, 5, v15);
406       iter.reset(db->NewIterator(read_options));
407 
408       SeekIterator(iter.get(), 1, 7);
409       ASSERT_TRUE(iter->Valid());
410       ASSERT_TRUE(v17 == iter->value());
411 
412       SeekIterator(iter.get(), 1, 5);
413       ASSERT_TRUE(iter->Valid());
414       ASSERT_TRUE(v15 == iter->value());
415       iter->Next();
416       ASSERT_TRUE(iter->Valid());
417       ASSERT_TRUE(v16 == iter->value());
418       iter->Next();
419       ASSERT_TRUE(iter->Valid());
420       ASSERT_TRUE(v17 == iter->value());
421 
422       SeekIterator(iter.get(), 1, 5);
423       ASSERT_TRUE(iter->Valid());
424       ASSERT_TRUE(v15 == iter->value());
425 
426       ASSERT_EQ(v15.ToString(), Get(db.get(), read_options, 1, 5));
427       ASSERT_EQ(v16.ToString(), Get(db.get(), read_options, 1, 6));
428       ASSERT_EQ(v17.ToString(), Get(db.get(), read_options, 1, 7));
429 
430       // 4. Insert an entry with a larger prefix
431       Slice v22("v22");
432       PutKey(db.get(), write_options, 2, 2, v22);
433       iter.reset(db->NewIterator(read_options));
434 
435       SeekIterator(iter.get(), 2, 2);
436       ASSERT_TRUE(iter->Valid());
437       ASSERT_TRUE(v22 == iter->value());
438       SeekIterator(iter.get(), 2, 0);
439       ASSERT_TRUE(iter->Valid());
440       ASSERT_TRUE(v22 == iter->value());
441 
442       SeekIterator(iter.get(), 1, 5);
443       ASSERT_TRUE(iter->Valid());
444       ASSERT_TRUE(v15 == iter->value());
445 
446       SeekIterator(iter.get(), 1, 7);
447       ASSERT_TRUE(iter->Valid());
448       ASSERT_TRUE(v17 == iter->value());
449 
450       // 5. Insert an entry with a smaller prefix
451       Slice v02("v02");
452       PutKey(db.get(), write_options, 0, 2, v02);
453       iter.reset(db->NewIterator(read_options));
454 
455       SeekIterator(iter.get(), 0, 2);
456       ASSERT_TRUE(iter->Valid());
457       ASSERT_TRUE(v02 == iter->value());
458       SeekIterator(iter.get(), 0, 0);
459       ASSERT_TRUE(iter->Valid());
460       ASSERT_TRUE(v02 == iter->value());
461 
462       SeekIterator(iter.get(), 2, 0);
463       ASSERT_TRUE(iter->Valid());
464       ASSERT_TRUE(v22 == iter->value());
465 
466       SeekIterator(iter.get(), 1, 5);
467       ASSERT_TRUE(iter->Valid());
468       ASSERT_TRUE(v15 == iter->value());
469 
470       SeekIterator(iter.get(), 1, 7);
471       ASSERT_TRUE(iter->Valid());
472       ASSERT_TRUE(v17 == iter->value());
473 
474       // 6. Insert to the beginning and the end of the first prefix
475       Slice v13("v13");
476       Slice v18("v18");
477       PutKey(db.get(), write_options, 1, 3, v13);
478       PutKey(db.get(), write_options, 1, 8, v18);
479       iter.reset(db->NewIterator(read_options));
480       SeekIterator(iter.get(), 1, 7);
481       ASSERT_TRUE(iter->Valid());
482       ASSERT_TRUE(v17 == iter->value());
483 
484       SeekIterator(iter.get(), 1, 3);
485       ASSERT_TRUE(iter->Valid());
486       ASSERT_TRUE(v13 == iter->value());
487       iter->Next();
488       ASSERT_TRUE(iter->Valid());
489       ASSERT_TRUE(v15 == iter->value());
490       iter->Next();
491       ASSERT_TRUE(iter->Valid());
492       ASSERT_TRUE(v16 == iter->value());
493       iter->Next();
494       ASSERT_TRUE(iter->Valid());
495       ASSERT_TRUE(v17 == iter->value());
496       iter->Next();
497       ASSERT_TRUE(iter->Valid());
498       ASSERT_TRUE(v18 == iter->value());
499 
500       SeekIterator(iter.get(), 0, 0);
501       ASSERT_TRUE(iter->Valid());
502       ASSERT_TRUE(v02 == iter->value());
503 
504       SeekIterator(iter.get(), 2, 0);
505       ASSERT_TRUE(iter->Valid());
506       ASSERT_TRUE(v22 == iter->value());
507 
508       ASSERT_EQ(v22.ToString(), Get(db.get(), read_options, 2, 2));
509       ASSERT_EQ(v02.ToString(), Get(db.get(), read_options, 0, 2));
510       ASSERT_EQ(v13.ToString(), Get(db.get(), read_options, 1, 3));
511       ASSERT_EQ(v15.ToString(), Get(db.get(), read_options, 1, 5));
512       ASSERT_EQ(v16.ToString(), Get(db.get(), read_options, 1, 6));
513       ASSERT_EQ(v17.ToString(), Get(db.get(), read_options, 1, 7));
514       ASSERT_EQ(v18.ToString(), Get(db.get(), read_options, 1, 8));
515     }
516   }
517 }
518 
519 // Show results in prefix
TEST_F(PrefixTest,PrefixValid)520 TEST_F(PrefixTest, PrefixValid) {
521   for (int num_buckets = 1; num_buckets <= 2; num_buckets++) {
522     FirstOption();
523     while (NextOptions(num_buckets)) {
524       std::cout << "*** Mem table: " << options.memtable_factory->Name()
525                 << " number of buckets: " << num_buckets << std::endl;
526       DestroyDB(kDbName, Options());
527       auto db = OpenDb();
528       WriteOptions write_options;
529       ReadOptions read_options;
530 
531       // Insert keys with common prefix and one key with different
532       Slice v16("v16");
533       Slice v17("v17");
534       Slice v18("v18");
535       Slice v19("v19");
536       PutKey(db.get(), write_options, 12345, 6, v16);
537       PutKey(db.get(), write_options, 12345, 7, v17);
538       PutKey(db.get(), write_options, 12345, 8, v18);
539       PutKey(db.get(), write_options, 12345, 9, v19);
540       PutKey(db.get(), write_options, 12346, 8, v16);
541       db->Flush(FlushOptions());
542       TestKey test_key(12346, 8);
543       std::string s;
544       db->Delete(write_options, TestKeyToSlice(s, test_key));
545       db->Flush(FlushOptions());
546       read_options.prefix_same_as_start = true;
547       std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
548       SeekIterator(iter.get(), 12345, 6);
549       ASSERT_TRUE(iter->Valid());
550       ASSERT_TRUE(v16 == iter->value());
551 
552       iter->Next();
553       ASSERT_TRUE(iter->Valid());
554       ASSERT_TRUE(v17 == iter->value());
555 
556       iter->Next();
557       ASSERT_TRUE(iter->Valid());
558       ASSERT_TRUE(v18 == iter->value());
559 
560       iter->Next();
561       ASSERT_TRUE(iter->Valid());
562       ASSERT_TRUE(v19 == iter->value());
563       iter->Next();
564       ASSERT_FALSE(iter->Valid());
565       ASSERT_EQ(kNotFoundResult, Get(db.get(), read_options, 12346, 8));
566 
567       // Verify seeking past the prefix won't return a result.
568       SeekIterator(iter.get(), 12345, 10);
569       ASSERT_TRUE(!iter->Valid());
570     }
571   }
572 }
573 
TEST_F(PrefixTest,DynamicPrefixIterator)574 TEST_F(PrefixTest, DynamicPrefixIterator) {
575   while (NextOptions(FLAGS_bucket_count)) {
576     std::cout << "*** Mem table: " << options.memtable_factory->Name()
577         << std::endl;
578     DestroyDB(kDbName, Options());
579     auto db = OpenDb();
580     WriteOptions write_options;
581     ReadOptions read_options;
582 
583     std::vector<uint64_t> prefixes;
584     for (uint64_t i = 0; i < FLAGS_total_prefixes; ++i) {
585       prefixes.push_back(i);
586     }
587 
588     if (FLAGS_random_prefix) {
589       std::random_shuffle(prefixes.begin(), prefixes.end());
590     }
591 
592     HistogramImpl hist_put_time;
593     HistogramImpl hist_put_comparison;
594 
595     // insert x random prefix, each with y continuous element.
596     for (auto prefix : prefixes) {
597        for (uint64_t sorted = 0; sorted < FLAGS_items_per_prefix; sorted++) {
598         TestKey test_key(prefix, sorted);
599 
600         std::string s;
601         Slice key = TestKeyToSlice(s, test_key);
602         std::string value(FLAGS_value_size, 0);
603 
604         get_perf_context()->Reset();
605         StopWatchNano timer(Env::Default(), true);
606         ASSERT_OK(db->Put(write_options, key, value));
607         hist_put_time.Add(timer.ElapsedNanos());
608         hist_put_comparison.Add(get_perf_context()->user_key_comparison_count);
609       }
610     }
611 
612     std::cout << "Put key comparison: \n" << hist_put_comparison.ToString()
613               << "Put time: \n" << hist_put_time.ToString();
614 
615     // test seek existing keys
616     HistogramImpl hist_seek_time;
617     HistogramImpl hist_seek_comparison;
618 
619     std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
620 
621     for (auto prefix : prefixes) {
622       TestKey test_key(prefix, FLAGS_items_per_prefix / 2);
623       std::string s;
624       Slice key = TestKeyToSlice(s, test_key);
625       std::string value = "v" + ToString(0);
626 
627       get_perf_context()->Reset();
628       StopWatchNano timer(Env::Default(), true);
629       auto key_prefix = options.prefix_extractor->Transform(key);
630       uint64_t total_keys = 0;
631       for (iter->Seek(key);
632            iter->Valid() && iter->key().starts_with(key_prefix);
633            iter->Next()) {
634         if (FLAGS_trigger_deadlock) {
635           std::cout << "Behold the deadlock!\n";
636           db->Delete(write_options, iter->key());
637         }
638         total_keys++;
639       }
640       hist_seek_time.Add(timer.ElapsedNanos());
641       hist_seek_comparison.Add(get_perf_context()->user_key_comparison_count);
642       ASSERT_EQ(total_keys, FLAGS_items_per_prefix - FLAGS_items_per_prefix/2);
643     }
644 
645     std::cout << "Seek key comparison: \n"
646               << hist_seek_comparison.ToString()
647               << "Seek time: \n"
648               << hist_seek_time.ToString();
649 
650     // test non-existing keys
651     HistogramImpl hist_no_seek_time;
652     HistogramImpl hist_no_seek_comparison;
653 
654     for (auto prefix = FLAGS_total_prefixes;
655          prefix < FLAGS_total_prefixes + 10000;
656          prefix++) {
657       TestKey test_key(prefix, 0);
658       std::string s;
659       Slice key = TestKeyToSlice(s, test_key);
660 
661       get_perf_context()->Reset();
662       StopWatchNano timer(Env::Default(), true);
663       iter->Seek(key);
664       hist_no_seek_time.Add(timer.ElapsedNanos());
665       hist_no_seek_comparison.Add(get_perf_context()->user_key_comparison_count);
666       ASSERT_TRUE(!iter->Valid());
667     }
668 
669     std::cout << "non-existing Seek key comparison: \n"
670               << hist_no_seek_comparison.ToString()
671               << "non-existing Seek time: \n"
672               << hist_no_seek_time.ToString();
673   }
674 }
675 
TEST_F(PrefixTest,PrefixSeekModePrev)676 TEST_F(PrefixTest, PrefixSeekModePrev) {
677   // Only for SkipListFactory
678   options.memtable_factory.reset(new SkipListFactory);
679   options.merge_operator = MergeOperators::CreatePutOperator();
680   options.write_buffer_size = 1024 * 1024;
681   Random rnd(1);
682   for (size_t m = 1; m < 100; m++) {
683     std::cout << "[" + std::to_string(m) + "]" + "*** Mem table: "
684               << options.memtable_factory->Name() << std::endl;
685     DestroyDB(kDbName, Options());
686     auto db = OpenDb();
687     WriteOptions write_options;
688     ReadOptions read_options;
689     std::map<TestKey, std::string, TestKeyComparator> entry_maps[3], whole_map;
690     for (uint64_t i = 0; i < 10; i++) {
691       int div = i % 3 + 1;
692       for (uint64_t j = 0; j < 10; j++) {
693         whole_map[TestKey(i, j)] = entry_maps[rnd.Uniform(div)][TestKey(i, j)] =
694             'v' + std::to_string(i) + std::to_string(j);
695       }
696     }
697 
698     std::map<TestKey, std::string, TestKeyComparator> type_map;
699     for (size_t i = 0; i < 3; i++) {
700       for (auto& kv : entry_maps[i]) {
701         if (rnd.OneIn(3)) {
702           PutKey(db.get(), write_options, kv.first, kv.second);
703           type_map[kv.first] = "value";
704         } else {
705           MergeKey(db.get(), write_options, kv.first, kv.second);
706           type_map[kv.first] = "merge";
707         }
708       }
709       if (i < 2) {
710         db->Flush(FlushOptions());
711       }
712     }
713 
714     for (size_t i = 0; i < 2; i++) {
715       for (auto& kv : entry_maps[i]) {
716         if (rnd.OneIn(10)) {
717           whole_map.erase(kv.first);
718           DeleteKey(db.get(), write_options, kv.first);
719           entry_maps[2][kv.first] = "delete";
720         }
721       }
722     }
723 
724     if (FLAGS_enable_print) {
725       for (size_t i = 0; i < 3; i++) {
726         for (auto& kv : entry_maps[i]) {
727           std::cout << "[" << i << "]" << kv.first.prefix << kv.first.sorted
728                     << " " << kv.second + " " + type_map[kv.first] << std::endl;
729         }
730       }
731     }
732 
733     std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
734     for (uint64_t prefix = 0; prefix < 10; prefix++) {
735       uint64_t start_suffix = rnd.Uniform(9);
736       SeekIterator(iter.get(), prefix, start_suffix);
737       auto it = whole_map.find(TestKey(prefix, start_suffix));
738       if (it == whole_map.end()) {
739         continue;
740       }
741       ASSERT_NE(it, whole_map.end());
742       ASSERT_TRUE(iter->Valid());
743       if (FLAGS_enable_print) {
744         std::cout << "round " << prefix
745                   << " iter: " << SliceToTestKey(iter->key()).prefix
746                   << SliceToTestKey(iter->key()).sorted
747                   << " | map: " << it->first.prefix << it->first.sorted << " | "
748                   << iter->value().ToString() << " " << it->second << std::endl;
749       }
750       ASSERT_EQ(iter->value(), it->second);
751       uint64_t stored_prefix = prefix;
752       for (size_t k = 0; k < 9; k++) {
753         if (rnd.OneIn(2) || it == whole_map.begin()) {
754           iter->Next();
755           ++it;
756           if (FLAGS_enable_print) {
757             std::cout << "Next >> ";
758           }
759         } else {
760           iter->Prev();
761           it--;
762           if (FLAGS_enable_print) {
763             std::cout << "Prev >> ";
764           }
765         }
766         if (!iter->Valid() ||
767             SliceToTestKey(iter->key()).prefix != stored_prefix) {
768           break;
769         }
770         stored_prefix = SliceToTestKey(iter->key()).prefix;
771         ASSERT_TRUE(iter->Valid());
772         ASSERT_NE(it, whole_map.end());
773         ASSERT_EQ(iter->value(), it->second);
774         if (FLAGS_enable_print) {
775           std::cout << "iter: " << SliceToTestKey(iter->key()).prefix
776                     << SliceToTestKey(iter->key()).sorted
777                     << " | map: " << it->first.prefix << it->first.sorted
778                     << " | " << iter->value().ToString() << " " << it->second
779                     << std::endl;
780         }
781       }
782     }
783   }
784 }
785 
TEST_F(PrefixTest,PrefixSeekModePrev2)786 TEST_F(PrefixTest, PrefixSeekModePrev2) {
787   // Only for SkipListFactory
788   // test the case
789   //        iter1                iter2
790   // | prefix | suffix |  | prefix | suffix |
791   // |   1    |   1    |  |   1    |   2    |
792   // |   1    |   3    |  |   1    |   4    |
793   // |   2    |   1    |  |   3    |   3    |
794   // |   2    |   2    |  |   3    |   4    |
795   // after seek(15), iter1 will be at 21 and iter2 will be 33.
796   // Then if call Prev() in prefix mode where SeekForPrev(21) gets called,
797   // iter2 should turn to invalid state because of bloom filter.
798   options.memtable_factory.reset(new SkipListFactory);
799   options.write_buffer_size = 1024 * 1024;
800   std::string v13("v13");
801   DestroyDB(kDbName, Options());
802   auto db = OpenDb();
803   WriteOptions write_options;
804   ReadOptions read_options;
805   PutKey(db.get(), write_options, TestKey(1, 2), "v12");
806   PutKey(db.get(), write_options, TestKey(1, 4), "v14");
807   PutKey(db.get(), write_options, TestKey(3, 3), "v33");
808   PutKey(db.get(), write_options, TestKey(3, 4), "v34");
809   db->Flush(FlushOptions());
810   reinterpret_cast<DBImpl*>(db.get())->TEST_WaitForFlushMemTable();
811   PutKey(db.get(), write_options, TestKey(1, 1), "v11");
812   PutKey(db.get(), write_options, TestKey(1, 3), "v13");
813   PutKey(db.get(), write_options, TestKey(2, 1), "v21");
814   PutKey(db.get(), write_options, TestKey(2, 2), "v22");
815   db->Flush(FlushOptions());
816   reinterpret_cast<DBImpl*>(db.get())->TEST_WaitForFlushMemTable();
817   std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
818   SeekIterator(iter.get(), 1, 5);
819   iter->Prev();
820   ASSERT_EQ(iter->value(), v13);
821 }
822 
TEST_F(PrefixTest,PrefixSeekModePrev3)823 TEST_F(PrefixTest, PrefixSeekModePrev3) {
824   // Only for SkipListFactory
825   // test SeekToLast() with iterate_upper_bound_ in prefix_seek_mode
826   options.memtable_factory.reset(new SkipListFactory);
827   options.write_buffer_size = 1024 * 1024;
828   std::string v14("v14");
829   TestKey upper_bound_key = TestKey(1, 5);
830   std::string s;
831   Slice upper_bound = TestKeyToSlice(s, upper_bound_key);
832 
833   {
834     DestroyDB(kDbName, Options());
835     auto db = OpenDb();
836     WriteOptions write_options;
837     ReadOptions read_options;
838     read_options.iterate_upper_bound = &upper_bound;
839     PutKey(db.get(), write_options, TestKey(1, 2), "v12");
840     PutKey(db.get(), write_options, TestKey(1, 4), "v14");
841     db->Flush(FlushOptions());
842     reinterpret_cast<DBImpl*>(db.get())->TEST_WaitForFlushMemTable();
843     PutKey(db.get(), write_options, TestKey(1, 1), "v11");
844     PutKey(db.get(), write_options, TestKey(1, 3), "v13");
845     PutKey(db.get(), write_options, TestKey(2, 1), "v21");
846     PutKey(db.get(), write_options, TestKey(2, 2), "v22");
847     db->Flush(FlushOptions());
848     reinterpret_cast<DBImpl*>(db.get())->TEST_WaitForFlushMemTable();
849     std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
850     iter->SeekToLast();
851     ASSERT_EQ(iter->value(), v14);
852   }
853   {
854     DestroyDB(kDbName, Options());
855     auto db = OpenDb();
856     WriteOptions write_options;
857     ReadOptions read_options;
858     read_options.iterate_upper_bound = &upper_bound;
859     PutKey(db.get(), write_options, TestKey(1, 2), "v12");
860     PutKey(db.get(), write_options, TestKey(1, 4), "v14");
861     PutKey(db.get(), write_options, TestKey(3, 3), "v33");
862     PutKey(db.get(), write_options, TestKey(3, 4), "v34");
863     db->Flush(FlushOptions());
864     reinterpret_cast<DBImpl*>(db.get())->TEST_WaitForFlushMemTable();
865     PutKey(db.get(), write_options, TestKey(1, 1), "v11");
866     PutKey(db.get(), write_options, TestKey(1, 3), "v13");
867     db->Flush(FlushOptions());
868     reinterpret_cast<DBImpl*>(db.get())->TEST_WaitForFlushMemTable();
869     std::unique_ptr<Iterator> iter(db->NewIterator(read_options));
870     iter->SeekToLast();
871     ASSERT_EQ(iter->value(), v14);
872   }
873 }
874 
875 }  // namespace ROCKSDB_NAMESPACE
876 
main(int argc,char ** argv)877 int main(int argc, char** argv) {
878   ::testing::InitGoogleTest(&argc, argv);
879   ParseCommandLineFlags(&argc, &argv, true);
880   return RUN_ALL_TESTS();
881 }
882 
883 #endif  // GFLAGS
884 
885 #else
886 #include <stdio.h>
887 
main(int,char **)888 int main(int /*argc*/, char** /*argv*/) {
889   fprintf(stderr,
890           "SKIPPED as HashSkipList and HashLinkList are not supported in "
891           "ROCKSDB_LITE\n");
892   return 0;
893 }
894 
895 #endif  // !ROCKSDB_LITE
896