ceph/src/rocksdb/db/db_memtable_test.cc

   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 #include <memory>
   7 #include <string>
   8
   9 #include "db/db_test_util.h"
  10 #include "db/memtable.h"
  11 #include "port/stack_trace.h"
  12 #include "rocksdb/memtablerep.h"
  13 #include "rocksdb/slice_transform.h"
  14
  15 namespace rocksdb {
  16
  17 class DBMemTableTest : public DBTestBase {
  18  public:
  19   DBMemTableTest() : DBTestBase("/db_memtable_test") {}
  20 };
  21
  22 class MockMemTableRep : public MemTableRep {
  23  public:
  24   explicit MockMemTableRep(Allocator* allocator, MemTableRep* rep)
  25       : MemTableRep(allocator), rep_(rep), num_insert_with_hint_(0) {}
  26
  27   virtual KeyHandle Allocate(const size_t len, char** buf) override {
  28     return rep_->Allocate(len, buf);
  29   }
  30
  31   virtual void Insert(KeyHandle handle) override { rep_->Insert(handle); }
  32
  33   virtual void InsertWithHint(KeyHandle handle, void** hint) override {
  34     num_insert_with_hint_++;
  35     EXPECT_NE(nullptr, hint);
  36     last_hint_in_ = *hint;
  37     rep_->InsertWithHint(handle, hint);
  38     last_hint_out_ = *hint;
  39   }
  40
  41   virtual bool Contains(const char* key) const override {
  42     return rep_->Contains(key);
  43   }
  44
  45   virtual void Get(const LookupKey& k, void* callback_args,
  46                    bool (*callback_func)(void* arg,
  47                                          const char* entry)) override {
  48     rep_->Get(k, callback_args, callback_func);
  49   }
  50
  51   virtual size_t ApproximateMemoryUsage() override {
  52     return rep_->ApproximateMemoryUsage();
  53   }
  54
  55   virtual Iterator* GetIterator(Arena* arena) override {
  56     return rep_->GetIterator(arena);
  57   }
  58
  59   void* last_hint_in() { return last_hint_in_; }
  60   void* last_hint_out() { return last_hint_out_; }
  61   int num_insert_with_hint() { return num_insert_with_hint_; }
  62
  63  private:
  64   std::unique_ptr<MemTableRep> rep_;
  65   void* last_hint_in_;
  66   void* last_hint_out_;
  67   int num_insert_with_hint_;
  68 };
  69
  70 class MockMemTableRepFactory : public MemTableRepFactory {
  71  public:
  72   virtual MemTableRep* CreateMemTableRep(const MemTableRep::KeyComparator& cmp,
  73                                          Allocator* allocator,
  74                                          const SliceTransform* transform,
  75                                          Logger* logger) override {
  76     SkipListFactory factory;
  77     MemTableRep* skiplist_rep =
  78         factory.CreateMemTableRep(cmp, allocator, transform, logger);
  79     mock_rep_ = new MockMemTableRep(allocator, skiplist_rep);
  80     return mock_rep_;
  81   }
  82
  83   virtual MemTableRep* CreateMemTableRep(const MemTableRep::KeyComparator& cmp,
  84                                          Allocator* allocator,
  85                                          const SliceTransform* transform,
  86                                          Logger* logger,
  87                                          uint32_t column_family_id) override {
  88     last_column_family_id_ = column_family_id;
  89     return CreateMemTableRep(cmp, allocator, transform, logger);
  90   }
  91
  92   virtual const char* Name() const override { return "MockMemTableRepFactory"; }
  93
  94   MockMemTableRep* rep() { return mock_rep_; }
  95
  96   bool IsInsertConcurrentlySupported() const override { return false; }
  97
  98   uint32_t GetLastColumnFamilyId() { return last_column_family_id_; }
  99
 100  private:
 101   MockMemTableRep* mock_rep_;
 102   // workaround since there's no port::kMaxUint32 yet.
 103   uint32_t last_column_family_id_ = static_cast<uint32_t>(-1);
 104 };
 105
 106 class TestPrefixExtractor : public SliceTransform {
 107  public:
 108   virtual const char* Name() const override { return "TestPrefixExtractor"; }
 109
 110   virtual Slice Transform(const Slice& key) const override {
 111     const char* p = separator(key);
 112     if (p == nullptr) {
 113       return Slice();
 114     }
 115     return Slice(key.data(), p - key.data() + 1);
 116   }
 117
 118   virtual bool InDomain(const Slice& key) const override {
 119     return separator(key) != nullptr;
 120   }
 121
 122   virtual bool InRange(const Slice& /*key*/) const override { return false; }
 123
 124  private:
 125   const char* separator(const Slice& key) const {
 126     return reinterpret_cast<const char*>(memchr(key.data(), '_', key.size()));
 127   }
 128 };
 129
 130 // Test that ::Add properly returns false when inserting duplicate keys
 131 TEST_F(DBMemTableTest, DuplicateSeq) {
 132   SequenceNumber seq = 123;
 133   std::string value;
 134   Status s;
 135   MergeContext merge_context;
 136   Options options;
 137   InternalKeyComparator ikey_cmp(options.comparator);
 138   RangeDelAggregator range_del_agg(ikey_cmp, {} /* snapshots */);
 139
 140   // Create a MemTable
 141   InternalKeyComparator cmp(BytewiseComparator());
 142   auto factory = std::make_shared<SkipListFactory>();
 143   options.memtable_factory = factory;
 144   ImmutableCFOptions ioptions(options);
 145   WriteBufferManager wb(options.db_write_buffer_size);
 146   MemTable* mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
 147                                kMaxSequenceNumber, 0 /* column_family_id */);
 148
 149   // Write some keys and make sure it returns false on duplicates
 150   bool res;
 151   res = mem->Add(seq, kTypeValue, "key", "value2");
 152   ASSERT_TRUE(res);
 153   res = mem->Add(seq, kTypeValue, "key", "value2");
 154   ASSERT_FALSE(res);
 155   // Changing the type should still cause the duplicatae key
 156   res = mem->Add(seq, kTypeMerge, "key", "value2");
 157   ASSERT_FALSE(res);
 158   // Changing the seq number will make the key fresh
 159   res = mem->Add(seq + 1, kTypeMerge, "key", "value2");
 160   ASSERT_TRUE(res);
 161   // Test with different types for duplicate keys
 162   res = mem->Add(seq, kTypeDeletion, "key", "");
 163   ASSERT_FALSE(res);
 164   res = mem->Add(seq, kTypeSingleDeletion, "key", "");
 165   ASSERT_FALSE(res);
 166
 167   // Test the duplicate keys under stress
 168   for (int i = 0; i < 10000; i++) {
 169     bool insert_dup = i % 10 == 1;
 170     if (!insert_dup) {
 171       seq++;
 172     }
 173     res = mem->Add(seq, kTypeValue, "foo", "value" + ToString(seq));
 174     if (insert_dup) {
 175       ASSERT_FALSE(res);
 176     } else {
 177       ASSERT_TRUE(res);
 178     }
 179   }
 180   delete mem;
 181
 182   // Test with InsertWithHint
 183   options.memtable_insert_with_hint_prefix_extractor.reset(
 184       new TestPrefixExtractor());  // which uses _ to extract the prefix
 185   ioptions = ImmutableCFOptions(options);
 186   mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
 187                      kMaxSequenceNumber, 0 /* column_family_id */);
 188   // Insert a duplicate key with _ in it
 189   res = mem->Add(seq, kTypeValue, "key_1", "value");
 190   ASSERT_TRUE(res);
 191   res = mem->Add(seq, kTypeValue, "key_1", "value");
 192   ASSERT_FALSE(res);
 193   delete mem;
 194
 195   // Test when InsertConcurrently will be invoked
 196   options.allow_concurrent_memtable_write = true;
 197   ioptions = ImmutableCFOptions(options);
 198   mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
 199                      kMaxSequenceNumber, 0 /* column_family_id */);
 200   MemTablePostProcessInfo post_process_info;
 201   res = mem->Add(seq, kTypeValue, "key", "value", true, &post_process_info);
 202   ASSERT_TRUE(res);
 203   res = mem->Add(seq, kTypeValue, "key", "value", true, &post_process_info);
 204   ASSERT_FALSE(res);
 205   delete mem;
 206 }
 207
 208 TEST_F(DBMemTableTest, InsertWithHint) {
 209   Options options;
 210   options.allow_concurrent_memtable_write = false;
 211   options.create_if_missing = true;
 212   options.memtable_factory.reset(new MockMemTableRepFactory());
 213   options.memtable_insert_with_hint_prefix_extractor.reset(
 214       new TestPrefixExtractor());
 215   options.env = env_;
 216   Reopen(options);
 217   MockMemTableRep* rep =
 218       reinterpret_cast<MockMemTableRepFactory*>(options.memtable_factory.get())
 219           ->rep();
 220   ASSERT_OK(Put("foo_k1", "foo_v1"));
 221   ASSERT_EQ(nullptr, rep->last_hint_in());
 222   void* hint_foo = rep->last_hint_out();
 223   ASSERT_OK(Put("foo_k2", "foo_v2"));
 224   ASSERT_EQ(hint_foo, rep->last_hint_in());
 225   ASSERT_EQ(hint_foo, rep->last_hint_out());
 226   ASSERT_OK(Put("foo_k3", "foo_v3"));
 227   ASSERT_EQ(hint_foo, rep->last_hint_in());
 228   ASSERT_EQ(hint_foo, rep->last_hint_out());
 229   ASSERT_OK(Put("bar_k1", "bar_v1"));
 230   ASSERT_EQ(nullptr, rep->last_hint_in());
 231   void* hint_bar = rep->last_hint_out();
 232   ASSERT_NE(hint_foo, hint_bar);
 233   ASSERT_OK(Put("bar_k2", "bar_v2"));
 234   ASSERT_EQ(hint_bar, rep->last_hint_in());
 235   ASSERT_EQ(hint_bar, rep->last_hint_out());
 236   ASSERT_EQ(5, rep->num_insert_with_hint());
 237   ASSERT_OK(Put("whitelisted", "vvv"));
 238   ASSERT_EQ(5, rep->num_insert_with_hint());
 239   ASSERT_EQ("foo_v1", Get("foo_k1"));
 240   ASSERT_EQ("foo_v2", Get("foo_k2"));
 241   ASSERT_EQ("foo_v3", Get("foo_k3"));
 242   ASSERT_EQ("bar_v1", Get("bar_k1"));
 243   ASSERT_EQ("bar_v2", Get("bar_k2"));
 244   ASSERT_EQ("vvv", Get("whitelisted"));
 245 }
 246
 247 TEST_F(DBMemTableTest, ColumnFamilyId) {
 248   // Verifies MemTableRepFactory is told the right column family id.
 249   Options options;
 250   options.allow_concurrent_memtable_write = false;
 251   options.create_if_missing = true;
 252   options.memtable_factory.reset(new MockMemTableRepFactory());
 253   DestroyAndReopen(options);
 254   CreateAndReopenWithCF({"pikachu"}, options);
 255
 256   for (int cf = 0; cf < 2; ++cf) {
 257     ASSERT_OK(Put(cf, "key", "val"));
 258     ASSERT_OK(Flush(cf));
 259     ASSERT_EQ(
 260         cf, static_cast<MockMemTableRepFactory*>(options.memtable_factory.get())
 261                 ->GetLastColumnFamilyId());
 262   }
 263 }
 264
 265 }  // namespace rocksdb
 266
 267 int main(int argc, char** argv) {
 268   rocksdb::port::InstallStackTraceHandler();
 269   ::testing::InitGoogleTest(&argc, argv);
 270   return RUN_ALL_TESTS();
 271 }