diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-21 11:54:28 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-21 11:54:28 +0000 |
commit | e6918187568dbd01842d8d1d2c808ce16a894239 (patch) | |
tree | 64f88b554b444a49f656b6c656111a145cbbaa28 /src/rocksdb/db/db_memtable_test.cc | |
parent | Initial commit. (diff) | |
download | ceph-upstream/18.2.2.tar.xz ceph-upstream/18.2.2.zip |
Adding upstream version 18.2.2.upstream/18.2.2
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to '')
-rw-r--r-- | src/rocksdb/db/db_memtable_test.cc | 344 |
1 files changed, 344 insertions, 0 deletions
diff --git a/src/rocksdb/db/db_memtable_test.cc b/src/rocksdb/db/db_memtable_test.cc new file mode 100644 index 000000000..cae592db3 --- /dev/null +++ b/src/rocksdb/db/db_memtable_test.cc @@ -0,0 +1,344 @@ +// Copyright (c) 2011-present, Facebook, Inc. All rights reserved. +// This source code is licensed under both the GPLv2 (found in the +// COPYING file in the root directory) and Apache 2.0 License +// (found in the LICENSE.Apache file in the root directory). + +#include <memory> +#include <string> + +#include "db/db_test_util.h" +#include "db/memtable.h" +#include "db/range_del_aggregator.h" +#include "port/stack_trace.h" +#include "rocksdb/memtablerep.h" +#include "rocksdb/slice_transform.h" + +namespace ROCKSDB_NAMESPACE { + +class DBMemTableTest : public DBTestBase { + public: + DBMemTableTest() : DBTestBase("db_memtable_test", /*env_do_fsync=*/true) {} +}; + +class MockMemTableRep : public MemTableRep { + public: + explicit MockMemTableRep(Allocator* allocator, MemTableRep* rep) + : MemTableRep(allocator), rep_(rep), num_insert_with_hint_(0) {} + + KeyHandle Allocate(const size_t len, char** buf) override { + return rep_->Allocate(len, buf); + } + + void Insert(KeyHandle handle) override { rep_->Insert(handle); } + + void InsertWithHint(KeyHandle handle, void** hint) override { + num_insert_with_hint_++; + EXPECT_NE(nullptr, hint); + last_hint_in_ = *hint; + rep_->InsertWithHint(handle, hint); + last_hint_out_ = *hint; + } + + bool Contains(const char* key) const override { return rep_->Contains(key); } + + void Get(const LookupKey& k, void* callback_args, + bool (*callback_func)(void* arg, const char* entry)) override { + rep_->Get(k, callback_args, callback_func); + } + + size_t ApproximateMemoryUsage() override { + return rep_->ApproximateMemoryUsage(); + } + + Iterator* GetIterator(Arena* arena) override { + return rep_->GetIterator(arena); + } + + void* last_hint_in() { return last_hint_in_; } + void* last_hint_out() { return last_hint_out_; } + int num_insert_with_hint() { return num_insert_with_hint_; } + + private: + std::unique_ptr<MemTableRep> rep_; + void* last_hint_in_; + void* last_hint_out_; + int num_insert_with_hint_; +}; + +class MockMemTableRepFactory : public MemTableRepFactory { + public: + MemTableRep* CreateMemTableRep(const MemTableRep::KeyComparator& cmp, + Allocator* allocator, + const SliceTransform* transform, + Logger* logger) override { + SkipListFactory factory; + MemTableRep* skiplist_rep = + factory.CreateMemTableRep(cmp, allocator, transform, logger); + mock_rep_ = new MockMemTableRep(allocator, skiplist_rep); + return mock_rep_; + } + + MemTableRep* CreateMemTableRep(const MemTableRep::KeyComparator& cmp, + Allocator* allocator, + const SliceTransform* transform, + Logger* logger, + uint32_t column_family_id) override { + last_column_family_id_ = column_family_id; + return CreateMemTableRep(cmp, allocator, transform, logger); + } + + const char* Name() const override { return "MockMemTableRepFactory"; } + + MockMemTableRep* rep() { return mock_rep_; } + + bool IsInsertConcurrentlySupported() const override { return false; } + + uint32_t GetLastColumnFamilyId() { return last_column_family_id_; } + + private: + MockMemTableRep* mock_rep_; + // workaround since there's no std::numeric_limits<uint32_t>::max() yet. + uint32_t last_column_family_id_ = static_cast<uint32_t>(-1); +}; + +class TestPrefixExtractor : public SliceTransform { + public: + const char* Name() const override { return "TestPrefixExtractor"; } + + Slice Transform(const Slice& key) const override { + const char* p = separator(key); + if (p == nullptr) { + return Slice(); + } + return Slice(key.data(), p - key.data() + 1); + } + + bool InDomain(const Slice& key) const override { + return separator(key) != nullptr; + } + + bool InRange(const Slice& /*key*/) const override { return false; } + + private: + const char* separator(const Slice& key) const { + return reinterpret_cast<const char*>(memchr(key.data(), '_', key.size())); + } +}; + +// Test that ::Add properly returns false when inserting duplicate keys +TEST_F(DBMemTableTest, DuplicateSeq) { + SequenceNumber seq = 123; + std::string value; + MergeContext merge_context; + Options options; + InternalKeyComparator ikey_cmp(options.comparator); + ReadRangeDelAggregator range_del_agg(&ikey_cmp, + kMaxSequenceNumber /* upper_bound */); + + // Create a MemTable + InternalKeyComparator cmp(BytewiseComparator()); + auto factory = std::make_shared<SkipListFactory>(); + options.memtable_factory = factory; + ImmutableOptions ioptions(options); + WriteBufferManager wb(options.db_write_buffer_size); + MemTable* mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb, + kMaxSequenceNumber, 0 /* column_family_id */); + + // Write some keys and make sure it returns false on duplicates + ASSERT_OK( + mem->Add(seq, kTypeValue, "key", "value2", nullptr /* kv_prot_info */)); + ASSERT_TRUE( + mem->Add(seq, kTypeValue, "key", "value2", nullptr /* kv_prot_info */) + .IsTryAgain()); + // Changing the type should still cause the duplicatae key + ASSERT_TRUE( + mem->Add(seq, kTypeMerge, "key", "value2", nullptr /* kv_prot_info */) + .IsTryAgain()); + // Changing the seq number will make the key fresh + ASSERT_OK(mem->Add(seq + 1, kTypeMerge, "key", "value2", + nullptr /* kv_prot_info */)); + // Test with different types for duplicate keys + ASSERT_TRUE( + mem->Add(seq, kTypeDeletion, "key", "", nullptr /* kv_prot_info */) + .IsTryAgain()); + ASSERT_TRUE( + mem->Add(seq, kTypeSingleDeletion, "key", "", nullptr /* kv_prot_info */) + .IsTryAgain()); + + // Test the duplicate keys under stress + for (int i = 0; i < 10000; i++) { + bool insert_dup = i % 10 == 1; + if (!insert_dup) { + seq++; + } + Status s = mem->Add(seq, kTypeValue, "foo", "value" + std::to_string(seq), + nullptr /* kv_prot_info */); + if (insert_dup) { + ASSERT_TRUE(s.IsTryAgain()); + } else { + ASSERT_OK(s); + } + } + delete mem; + + // Test with InsertWithHint + options.memtable_insert_with_hint_prefix_extractor.reset( + new TestPrefixExtractor()); // which uses _ to extract the prefix + ioptions = ImmutableOptions(options); + mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb, + kMaxSequenceNumber, 0 /* column_family_id */); + // Insert a duplicate key with _ in it + ASSERT_OK( + mem->Add(seq, kTypeValue, "key_1", "value", nullptr /* kv_prot_info */)); + ASSERT_TRUE( + mem->Add(seq, kTypeValue, "key_1", "value", nullptr /* kv_prot_info */) + .IsTryAgain()); + delete mem; + + // Test when InsertConcurrently will be invoked + options.allow_concurrent_memtable_write = true; + ioptions = ImmutableOptions(options); + mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb, + kMaxSequenceNumber, 0 /* column_family_id */); + MemTablePostProcessInfo post_process_info; + ASSERT_OK(mem->Add(seq, kTypeValue, "key", "value", + nullptr /* kv_prot_info */, true, &post_process_info)); + ASSERT_TRUE(mem->Add(seq, kTypeValue, "key", "value", + nullptr /* kv_prot_info */, true, &post_process_info) + .IsTryAgain()); + delete mem; +} + +// A simple test to verify that the concurrent merge writes is functional +TEST_F(DBMemTableTest, ConcurrentMergeWrite) { + int num_ops = 1000; + std::string value; + MergeContext merge_context; + Options options; + // A merge operator that is not sensitive to concurrent writes since in this + // test we don't order the writes. + options.merge_operator = MergeOperators::CreateUInt64AddOperator(); + + // Create a MemTable + InternalKeyComparator cmp(BytewiseComparator()); + auto factory = std::make_shared<SkipListFactory>(); + options.memtable_factory = factory; + options.allow_concurrent_memtable_write = true; + ImmutableOptions ioptions(options); + WriteBufferManager wb(options.db_write_buffer_size); + MemTable* mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb, + kMaxSequenceNumber, 0 /* column_family_id */); + + // Put 0 as the base + PutFixed64(&value, static_cast<uint64_t>(0)); + ASSERT_OK(mem->Add(0, kTypeValue, "key", value, nullptr /* kv_prot_info */)); + value.clear(); + + // Write Merge concurrently + ROCKSDB_NAMESPACE::port::Thread write_thread1([&]() { + MemTablePostProcessInfo post_process_info1; + std::string v1; + for (int seq = 1; seq < num_ops / 2; seq++) { + PutFixed64(&v1, seq); + ASSERT_OK(mem->Add(seq, kTypeMerge, "key", v1, nullptr /* kv_prot_info */, + true, &post_process_info1)); + v1.clear(); + } + }); + ROCKSDB_NAMESPACE::port::Thread write_thread2([&]() { + MemTablePostProcessInfo post_process_info2; + std::string v2; + for (int seq = num_ops / 2; seq < num_ops; seq++) { + PutFixed64(&v2, seq); + ASSERT_OK(mem->Add(seq, kTypeMerge, "key", v2, nullptr /* kv_prot_info */, + true, &post_process_info2)); + v2.clear(); + } + }); + write_thread1.join(); + write_thread2.join(); + + Status status; + ReadOptions roptions; + SequenceNumber max_covering_tombstone_seq = 0; + LookupKey lkey("key", kMaxSequenceNumber); + bool res = mem->Get(lkey, &value, /*columns=*/nullptr, /*timestamp=*/nullptr, + &status, &merge_context, &max_covering_tombstone_seq, + roptions, false /* immutable_memtable */); + ASSERT_OK(status); + ASSERT_TRUE(res); + uint64_t ivalue = DecodeFixed64(Slice(value).data()); + uint64_t sum = 0; + for (int seq = 0; seq < num_ops; seq++) { + sum += seq; + } + ASSERT_EQ(ivalue, sum); + + delete mem; +} + +TEST_F(DBMemTableTest, InsertWithHint) { + Options options; + options.allow_concurrent_memtable_write = false; + options.create_if_missing = true; + options.memtable_factory.reset(new MockMemTableRepFactory()); + options.memtable_insert_with_hint_prefix_extractor.reset( + new TestPrefixExtractor()); + options.env = env_; + Reopen(options); + MockMemTableRep* rep = + reinterpret_cast<MockMemTableRepFactory*>(options.memtable_factory.get()) + ->rep(); + ASSERT_OK(Put("foo_k1", "foo_v1")); + ASSERT_EQ(nullptr, rep->last_hint_in()); + void* hint_foo = rep->last_hint_out(); + ASSERT_OK(Put("foo_k2", "foo_v2")); + ASSERT_EQ(hint_foo, rep->last_hint_in()); + ASSERT_EQ(hint_foo, rep->last_hint_out()); + ASSERT_OK(Put("foo_k3", "foo_v3")); + ASSERT_EQ(hint_foo, rep->last_hint_in()); + ASSERT_EQ(hint_foo, rep->last_hint_out()); + ASSERT_OK(Put("bar_k1", "bar_v1")); + ASSERT_EQ(nullptr, rep->last_hint_in()); + void* hint_bar = rep->last_hint_out(); + ASSERT_NE(hint_foo, hint_bar); + ASSERT_OK(Put("bar_k2", "bar_v2")); + ASSERT_EQ(hint_bar, rep->last_hint_in()); + ASSERT_EQ(hint_bar, rep->last_hint_out()); + ASSERT_EQ(5, rep->num_insert_with_hint()); + ASSERT_OK(Put("NotInPrefixDomain", "vvv")); + ASSERT_EQ(5, rep->num_insert_with_hint()); + ASSERT_EQ("foo_v1", Get("foo_k1")); + ASSERT_EQ("foo_v2", Get("foo_k2")); + ASSERT_EQ("foo_v3", Get("foo_k3")); + ASSERT_EQ("bar_v1", Get("bar_k1")); + ASSERT_EQ("bar_v2", Get("bar_k2")); + ASSERT_EQ("vvv", Get("NotInPrefixDomain")); +} + +TEST_F(DBMemTableTest, ColumnFamilyId) { + // Verifies MemTableRepFactory is told the right column family id. + Options options; + options.env = CurrentOptions().env; + options.allow_concurrent_memtable_write = false; + options.create_if_missing = true; + options.memtable_factory.reset(new MockMemTableRepFactory()); + DestroyAndReopen(options); + CreateAndReopenWithCF({"pikachu"}, options); + + for (uint32_t cf = 0; cf < 2; ++cf) { + ASSERT_OK(Put(cf, "key", "val")); + ASSERT_OK(Flush(cf)); + ASSERT_EQ( + cf, static_cast<MockMemTableRepFactory*>(options.memtable_factory.get()) + ->GetLastColumnFamilyId()); + } +} + +} // namespace ROCKSDB_NAMESPACE + +int main(int argc, char** argv) { + ROCKSDB_NAMESPACE::port::InstallStackTraceHandler(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} |