diff --git a/Makefile b/Makefile index cb29a0079..a8612b343 100644 --- a/Makefile +++ b/Makefile @@ -246,6 +246,9 @@ table_reader_bench: table/table_reader_bench.o $(LIBOBJECTS) $(TESTHARNESS) perf_context_test: db/perf_context_test.o $(LIBOBJECTS) $(TESTHARNESS) $(CXX) db/perf_context_test.o $(LIBOBJECTS) $(TESTHARNESS) $(EXEC_LDFLAGS) -o $@ $(LDFLAGS) +prefix_test: db/prefix_test.o $(LIBOBJECTS) $(TESTHARNESS) + $(CXX) db/prefix_test.o $(LIBOBJECTS) $(TESTHARNESS) $(EXEC_LDFLAGS) -o $@ $(LDFLAGS) + ttl_test: utilities/ttl/ttl_test.o $(LIBOBJECTS) $(TESTHARNESS) $(CXX) utilities/ttl/ttl_test.o $(LIBOBJECTS) $(TESTHARNESS) $(EXEC_LDFLAGS) -o $@ $(LDFLAGS) $(COVERAGEFLAGS) diff --git a/db/prefix_test.cc b/db/prefix_test.cc new file mode 100644 index 000000000..0ec84269b --- /dev/null +++ b/db/prefix_test.cc @@ -0,0 +1,225 @@ +#include +#include +#include + +#include +#include "rocksdb/comparator.h" +#include "rocksdb/db.h" +#include "rocksdb/perf_context.h" +#include "util/histogram.h" +#include "util/stop_watch.h" +#include "util/testharness.h" + +DEFINE_bool(use_prefix_hash_memtable, true, ""); +DEFINE_bool(trigger_deadlock, false, + "issue delete in range scan to trigger PrefixHashMap deadlock"); +DEFINE_uint64(bucket_count, 100000, "number of buckets"); +DEFINE_uint64(num_locks, 10001, "number of locks"); +DEFINE_bool(random_prefix, false, "randomize prefix"); +DEFINE_uint64(total_prefixes, 1000, "total number of prefixes"); +DEFINE_uint64(items_per_prefix, 10, "total number of values per prefix"); +DEFINE_int64(write_buffer_size, 1000000000, ""); +DEFINE_int64(max_write_buffer_number, 8, ""); +DEFINE_int64(min_write_buffer_number_to_merge, 7, ""); + +// Path to the database on file system +const std::string kDbName = rocksdb::test::TmpDir() + "/prefix_test"; + +namespace rocksdb { + +struct TestKey { + uint64_t prefix; + uint64_t sorted; + + TestKey(uint64_t prefix, uint64_t sorted) : prefix(prefix), sorted(sorted) {} +}; + +// return a slice backed by test_key +inline Slice TestKeyToSlice(const TestKey& test_key) { + return Slice((const char*)&test_key, sizeof(test_key)); +} + +inline const TestKey* SliceToTestKey(const Slice& slice) { + assert(slice.size() == sizeof(TestKey)); + return (const TestKey*)slice.data(); +} + +class TestKeyComparator : public Comparator { + public: + virtual int Compare(const Slice& a, const Slice& b) const { + const TestKey* key_a = SliceToTestKey(a); + const TestKey* key_b = SliceToTestKey(b); + if (key_a->prefix != key_b->prefix) { + if (key_a->prefix < key_b->prefix) return -1; + if (key_a->prefix > key_b->prefix) return 1; + assert(false); + } else { + if (key_a->sorted < key_b->sorted) return -1; + if (key_a->sorted > key_b->sorted) return 1; + if (key_a->sorted == key_b->sorted) return 0; + assert(false); + } + assert(false); + return 0; + } + + virtual const char* Name() const override { + return "TestKeyComparator"; + } + + virtual void FindShortestSeparator( + std::string* start, + const Slice& limit) const { + } + + virtual void FindShortSuccessor(std::string* key) const {} + + private: + +}; + +class PrefixTest { + public: + std::shared_ptr OpenDb() { + DB* db; + + options.create_if_missing = true; + options.write_buffer_size = FLAGS_write_buffer_size; + options.max_write_buffer_number = FLAGS_max_write_buffer_number; + options.min_write_buffer_number_to_merge = + FLAGS_min_write_buffer_number_to_merge; + + options.comparator = new TestKeyComparator(); + if (FLAGS_use_prefix_hash_memtable) { + auto prefix_extractor = NewFixedPrefixTransform(8); + options.prefix_extractor = prefix_extractor; + options.memtable_factory = + std::make_shared( + prefix_extractor, FLAGS_bucket_count, FLAGS_num_locks); + } + + Status s = DB::Open(options, kDbName, &db); + ASSERT_OK(s); + return std::shared_ptr(db); + } + protected: + Options options; +}; + + +TEST(PrefixTest, PrefixHash) { + + DestroyDB(kDbName, Options()); + auto db = OpenDb(); + WriteOptions write_options; + ReadOptions read_options; + + std::vector prefixes; + for (uint64_t i = 0; i < FLAGS_total_prefixes; ++i) { + prefixes.push_back(i); + } + + if (FLAGS_random_prefix) { + std::random_shuffle(prefixes.begin(), prefixes.end()); + } + + // insert x random prefix, each with y continuous element. + HistogramImpl hist_put_time; + HistogramImpl hist_put_comparison; + + for (auto prefix : prefixes) { + for (uint64_t sorted = 0; sorted < FLAGS_items_per_prefix; sorted++) { + TestKey test_key(prefix, sorted); + + Slice key = TestKeyToSlice(test_key); + std::string value = "v" + std::to_string(sorted); + + perf_context.Reset(); + StopWatchNano timer(Env::Default(), true); + ASSERT_OK(db->Put(write_options, key, value)); + hist_put_time.Add(timer.ElapsedNanos()); + hist_put_comparison.Add(perf_context.user_key_comparison_count); + } + } + + std::cout << "Put key comparison: \n" << hist_put_comparison.ToString() + << "Put time: \n" << hist_put_time.ToString(); + + + // test seek existing keys + HistogramImpl hist_seek_time; + HistogramImpl hist_seek_comparison; + + for (auto prefix : prefixes) { + TestKey test_key(prefix, 0); + Slice key = TestKeyToSlice(test_key); + std::string value = "v" + std::to_string(0); + + Slice key_prefix; + if (FLAGS_use_prefix_hash_memtable) { + key_prefix = options.prefix_extractor->Transform(key); + read_options.prefix = &key_prefix; + } + std::unique_ptr iter(db->NewIterator(read_options)); + + perf_context.Reset(); + StopWatchNano timer(Env::Default(), true); + uint64_t total_keys = 0; + for (iter->Seek(key); iter->Valid(); iter->Next()) { + if (FLAGS_trigger_deadlock) { + std::cout << "Behold the deadlock!\n"; + db->Delete(write_options, iter->key()); + } + auto test_key = SliceToTestKey(iter->key()); + if (test_key->prefix != prefix) break; + total_keys++; + } + hist_seek_time.Add(timer.ElapsedNanos()); + hist_seek_comparison.Add(perf_context.user_key_comparison_count); + ASSERT_EQ(total_keys, FLAGS_items_per_prefix); + } + + std::cout << "Seek key comparison: \n" + << hist_seek_comparison.ToString() + << "Seek time: \n" + << hist_seek_time.ToString(); + + // test non-existing keys + HistogramImpl hist_no_seek_time; + HistogramImpl hist_no_seek_comparison; + + for (auto prefix = FLAGS_total_prefixes; + prefix < FLAGS_total_prefixes + 100; + prefix++) { + TestKey test_key(prefix, 0); + Slice key = TestKeyToSlice(test_key); + + if (FLAGS_use_prefix_hash_memtable) { + Slice key_prefix = options.prefix_extractor->Transform(key); + read_options.prefix = &key_prefix; + } + std::unique_ptr iter(db->NewIterator(read_options)); + + perf_context.Reset(); + StopWatchNano timer(Env::Default(), true); + iter->Seek(key); + hist_no_seek_time.Add(timer.ElapsedNanos()); + hist_no_seek_comparison.Add(perf_context.user_key_comparison_count); + ASSERT_TRUE(!iter->Valid()); + } + + std::cout << "non-existing Seek key comparison: \n" + << hist_no_seek_comparison.ToString() + << "non-existing Seek time: \n" + << hist_no_seek_time.ToString(); +} + +} + +int main(int argc, char** argv) { + google::ParseCommandLineFlags(&argc, &argv, true); + std::cout << kDbName << "\n"; + + rocksdb::test::RunAllTests(); + return 0; +}