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 #include "db/db_impl/db_impl.h"
9 #include "rocksdb/db.h"
10 #include "rocksdb/env.h"
11 #include "table/cuckoo/cuckoo_table_factory.h"
12 #include "table/cuckoo/cuckoo_table_reader.h"
13 #include "table/meta_blocks.h"
14 #include "test_util/testharness.h"
15 #include "test_util/testutil.h"
16 #include "util/string_util.h"
17
18 namespace ROCKSDB_NAMESPACE {
19
20 class CuckooTableDBTest : public testing::Test {
21 private:
22 std::string dbname_;
23 Env* env_;
24 DB* db_;
25
26 public:
CuckooTableDBTest()27 CuckooTableDBTest() : env_(Env::Default()) {
28 dbname_ = test::PerThreadDBPath("cuckoo_table_db_test");
29 EXPECT_OK(DestroyDB(dbname_, Options()));
30 db_ = nullptr;
31 Reopen();
32 }
33
~CuckooTableDBTest()34 ~CuckooTableDBTest() override {
35 delete db_;
36 EXPECT_OK(DestroyDB(dbname_, Options()));
37 }
38
CurrentOptions()39 Options CurrentOptions() {
40 Options options;
41 options.table_factory.reset(NewCuckooTableFactory());
42 options.memtable_factory.reset(NewHashLinkListRepFactory(4, 0, 3, true));
43 options.allow_mmap_reads = true;
44 options.create_if_missing = true;
45 options.allow_concurrent_memtable_write = false;
46 return options;
47 }
48
dbfull()49 DBImpl* dbfull() {
50 return reinterpret_cast<DBImpl*>(db_);
51 }
52
53 // The following util methods are copied from plain_table_db_test.
Reopen(Options * options=nullptr)54 void Reopen(Options* options = nullptr) {
55 delete db_;
56 db_ = nullptr;
57 Options opts;
58 if (options != nullptr) {
59 opts = *options;
60 } else {
61 opts = CurrentOptions();
62 opts.create_if_missing = true;
63 }
64 ASSERT_OK(DB::Open(opts, dbname_, &db_));
65 }
66
Put(const Slice & k,const Slice & v)67 Status Put(const Slice& k, const Slice& v) {
68 return db_->Put(WriteOptions(), k, v);
69 }
70
Delete(const std::string & k)71 Status Delete(const std::string& k) {
72 return db_->Delete(WriteOptions(), k);
73 }
74
Get(const std::string & k)75 std::string Get(const std::string& k) {
76 ReadOptions options;
77 std::string result;
78 Status s = db_->Get(options, k, &result);
79 if (s.IsNotFound()) {
80 result = "NOT_FOUND";
81 } else if (!s.ok()) {
82 result = s.ToString();
83 }
84 return result;
85 }
86
NumTableFilesAtLevel(int level)87 int NumTableFilesAtLevel(int level) {
88 std::string property;
89 EXPECT_TRUE(db_->GetProperty(
90 "rocksdb.num-files-at-level" + NumberToString(level), &property));
91 return atoi(property.c_str());
92 }
93
94 // Return spread of files per level
FilesPerLevel()95 std::string FilesPerLevel() {
96 std::string result;
97 size_t last_non_zero_offset = 0;
98 for (int level = 0; level < db_->NumberLevels(); level++) {
99 int f = NumTableFilesAtLevel(level);
100 char buf[100];
101 snprintf(buf, sizeof(buf), "%s%d", (level ? "," : ""), f);
102 result += buf;
103 if (f > 0) {
104 last_non_zero_offset = result.size();
105 }
106 }
107 result.resize(last_non_zero_offset);
108 return result;
109 }
110 };
111
TEST_F(CuckooTableDBTest,Flush)112 TEST_F(CuckooTableDBTest, Flush) {
113 // Try with empty DB first.
114 ASSERT_TRUE(dbfull() != nullptr);
115 ASSERT_EQ("NOT_FOUND", Get("key2"));
116
117 // Add some values to db.
118 Options options = CurrentOptions();
119 Reopen(&options);
120
121 ASSERT_OK(Put("key1", "v1"));
122 ASSERT_OK(Put("key2", "v2"));
123 ASSERT_OK(Put("key3", "v3"));
124 dbfull()->TEST_FlushMemTable();
125
126 TablePropertiesCollection ptc;
127 reinterpret_cast<DB*>(dbfull())->GetPropertiesOfAllTables(&ptc);
128 ASSERT_EQ(1U, ptc.size());
129 ASSERT_EQ(3U, ptc.begin()->second->num_entries);
130 ASSERT_EQ("1", FilesPerLevel());
131
132 ASSERT_EQ("v1", Get("key1"));
133 ASSERT_EQ("v2", Get("key2"));
134 ASSERT_EQ("v3", Get("key3"));
135 ASSERT_EQ("NOT_FOUND", Get("key4"));
136
137 // Now add more keys and flush.
138 ASSERT_OK(Put("key4", "v4"));
139 ASSERT_OK(Put("key5", "v5"));
140 ASSERT_OK(Put("key6", "v6"));
141 dbfull()->TEST_FlushMemTable();
142
143 reinterpret_cast<DB*>(dbfull())->GetPropertiesOfAllTables(&ptc);
144 ASSERT_EQ(2U, ptc.size());
145 auto row = ptc.begin();
146 ASSERT_EQ(3U, row->second->num_entries);
147 ASSERT_EQ(3U, (++row)->second->num_entries);
148 ASSERT_EQ("2", FilesPerLevel());
149 ASSERT_EQ("v1", Get("key1"));
150 ASSERT_EQ("v2", Get("key2"));
151 ASSERT_EQ("v3", Get("key3"));
152 ASSERT_EQ("v4", Get("key4"));
153 ASSERT_EQ("v5", Get("key5"));
154 ASSERT_EQ("v6", Get("key6"));
155
156 ASSERT_OK(Delete("key6"));
157 ASSERT_OK(Delete("key5"));
158 ASSERT_OK(Delete("key4"));
159 dbfull()->TEST_FlushMemTable();
160 reinterpret_cast<DB*>(dbfull())->GetPropertiesOfAllTables(&ptc);
161 ASSERT_EQ(3U, ptc.size());
162 row = ptc.begin();
163 ASSERT_EQ(3U, row->second->num_entries);
164 ASSERT_EQ(3U, (++row)->second->num_entries);
165 ASSERT_EQ(3U, (++row)->second->num_entries);
166 ASSERT_EQ("3", FilesPerLevel());
167 ASSERT_EQ("v1", Get("key1"));
168 ASSERT_EQ("v2", Get("key2"));
169 ASSERT_EQ("v3", Get("key3"));
170 ASSERT_EQ("NOT_FOUND", Get("key4"));
171 ASSERT_EQ("NOT_FOUND", Get("key5"));
172 ASSERT_EQ("NOT_FOUND", Get("key6"));
173 }
174
TEST_F(CuckooTableDBTest,FlushWithDuplicateKeys)175 TEST_F(CuckooTableDBTest, FlushWithDuplicateKeys) {
176 Options options = CurrentOptions();
177 Reopen(&options);
178 ASSERT_OK(Put("key1", "v1"));
179 ASSERT_OK(Put("key2", "v2"));
180 ASSERT_OK(Put("key1", "v3")); // Duplicate
181 dbfull()->TEST_FlushMemTable();
182
183 TablePropertiesCollection ptc;
184 reinterpret_cast<DB*>(dbfull())->GetPropertiesOfAllTables(&ptc);
185 ASSERT_EQ(1U, ptc.size());
186 ASSERT_EQ(2U, ptc.begin()->second->num_entries);
187 ASSERT_EQ("1", FilesPerLevel());
188 ASSERT_EQ("v3", Get("key1"));
189 ASSERT_EQ("v2", Get("key2"));
190 }
191
192 namespace {
Key(int i)193 static std::string Key(int i) {
194 char buf[100];
195 snprintf(buf, sizeof(buf), "key_______%06d", i);
196 return std::string(buf);
197 }
Uint64Key(uint64_t i)198 static std::string Uint64Key(uint64_t i) {
199 std::string str;
200 str.resize(8);
201 memcpy(&str[0], static_cast<void*>(&i), 8);
202 return str;
203 }
204 } // namespace.
205
TEST_F(CuckooTableDBTest,Uint64Comparator)206 TEST_F(CuckooTableDBTest, Uint64Comparator) {
207 Options options = CurrentOptions();
208 options.comparator = test::Uint64Comparator();
209 Reopen(&options);
210
211 ASSERT_OK(Put(Uint64Key(1), "v1"));
212 ASSERT_OK(Put(Uint64Key(2), "v2"));
213 ASSERT_OK(Put(Uint64Key(3), "v3"));
214 dbfull()->TEST_FlushMemTable();
215
216 ASSERT_EQ("v1", Get(Uint64Key(1)));
217 ASSERT_EQ("v2", Get(Uint64Key(2)));
218 ASSERT_EQ("v3", Get(Uint64Key(3)));
219 ASSERT_EQ("NOT_FOUND", Get(Uint64Key(4)));
220
221 // Add more keys.
222 ASSERT_OK(Delete(Uint64Key(2))); // Delete.
223 dbfull()->TEST_FlushMemTable();
224 ASSERT_OK(Put(Uint64Key(3), "v0")); // Update.
225 ASSERT_OK(Put(Uint64Key(4), "v4"));
226 dbfull()->TEST_FlushMemTable();
227 ASSERT_EQ("v1", Get(Uint64Key(1)));
228 ASSERT_EQ("NOT_FOUND", Get(Uint64Key(2)));
229 ASSERT_EQ("v0", Get(Uint64Key(3)));
230 ASSERT_EQ("v4", Get(Uint64Key(4)));
231 }
232
TEST_F(CuckooTableDBTest,CompactionIntoMultipleFiles)233 TEST_F(CuckooTableDBTest, CompactionIntoMultipleFiles) {
234 // Create a big L0 file and check it compacts into multiple files in L1.
235 Options options = CurrentOptions();
236 options.write_buffer_size = 270 << 10;
237 // Two SST files should be created, each containing 14 keys.
238 // Number of buckets will be 16. Total size ~156 KB.
239 options.target_file_size_base = 160 << 10;
240 Reopen(&options);
241
242 // Write 28 values, each 10016 B ~ 10KB
243 for (int idx = 0; idx < 28; ++idx) {
244 ASSERT_OK(Put(Key(idx), std::string(10000, 'a' + char(idx))));
245 }
246 dbfull()->TEST_WaitForFlushMemTable();
247 ASSERT_EQ("1", FilesPerLevel());
248
249 dbfull()->TEST_CompactRange(0, nullptr, nullptr, nullptr,
250 true /* disallow trivial move */);
251 ASSERT_EQ("0,2", FilesPerLevel());
252 for (int idx = 0; idx < 28; ++idx) {
253 ASSERT_EQ(std::string(10000, 'a' + char(idx)), Get(Key(idx)));
254 }
255 }
256
TEST_F(CuckooTableDBTest,SameKeyInsertedInTwoDifferentFilesAndCompacted)257 TEST_F(CuckooTableDBTest, SameKeyInsertedInTwoDifferentFilesAndCompacted) {
258 // Insert same key twice so that they go to different SST files. Then wait for
259 // compaction and check if the latest value is stored and old value removed.
260 Options options = CurrentOptions();
261 options.write_buffer_size = 100 << 10; // 100KB
262 options.level0_file_num_compaction_trigger = 2;
263 Reopen(&options);
264
265 // Write 11 values, each 10016 B
266 for (int idx = 0; idx < 11; ++idx) {
267 ASSERT_OK(Put(Key(idx), std::string(10000, 'a')));
268 }
269 dbfull()->TEST_WaitForFlushMemTable();
270 ASSERT_EQ("1", FilesPerLevel());
271
272 // Generate one more file in level-0, and should trigger level-0 compaction
273 for (int idx = 0; idx < 11; ++idx) {
274 ASSERT_OK(Put(Key(idx), std::string(10000, 'a' + char(idx))));
275 }
276 dbfull()->TEST_WaitForFlushMemTable();
277 dbfull()->TEST_CompactRange(0, nullptr, nullptr);
278
279 ASSERT_EQ("0,1", FilesPerLevel());
280 for (int idx = 0; idx < 11; ++idx) {
281 ASSERT_EQ(std::string(10000, 'a' + char(idx)), Get(Key(idx)));
282 }
283 }
284
TEST_F(CuckooTableDBTest,AdaptiveTable)285 TEST_F(CuckooTableDBTest, AdaptiveTable) {
286 Options options = CurrentOptions();
287
288 // Ensure options compatible with PlainTable
289 options.prefix_extractor.reset(NewCappedPrefixTransform(8));
290
291 // Write some keys using cuckoo table.
292 options.table_factory.reset(NewCuckooTableFactory());
293 Reopen(&options);
294
295 ASSERT_OK(Put("key1", "v1"));
296 ASSERT_OK(Put("key2", "v2"));
297 ASSERT_OK(Put("key3", "v3"));
298 dbfull()->TEST_FlushMemTable();
299
300 // Write some keys using plain table.
301 std::shared_ptr<TableFactory> block_based_factory(
302 NewBlockBasedTableFactory());
303 std::shared_ptr<TableFactory> plain_table_factory(
304 NewPlainTableFactory());
305 std::shared_ptr<TableFactory> cuckoo_table_factory(
306 NewCuckooTableFactory());
307 options.create_if_missing = false;
308 options.table_factory.reset(NewAdaptiveTableFactory(
309 plain_table_factory, block_based_factory, plain_table_factory,
310 cuckoo_table_factory));
311 Reopen(&options);
312 ASSERT_OK(Put("key4", "v4"));
313 ASSERT_OK(Put("key1", "v5"));
314 dbfull()->TEST_FlushMemTable();
315
316 // Write some keys using block based table.
317 options.table_factory.reset(NewAdaptiveTableFactory(
318 block_based_factory, block_based_factory, plain_table_factory,
319 cuckoo_table_factory));
320 Reopen(&options);
321 ASSERT_OK(Put("key5", "v6"));
322 ASSERT_OK(Put("key2", "v7"));
323 dbfull()->TEST_FlushMemTable();
324
325 ASSERT_EQ("v5", Get("key1"));
326 ASSERT_EQ("v7", Get("key2"));
327 ASSERT_EQ("v3", Get("key3"));
328 ASSERT_EQ("v4", Get("key4"));
329 ASSERT_EQ("v6", Get("key5"));
330 }
331 } // namespace ROCKSDB_NAMESPACE
332
main(int argc,char ** argv)333 int main(int argc, char** argv) {
334 if (ROCKSDB_NAMESPACE::port::kLittleEndian) {
335 ::testing::InitGoogleTest(&argc, argv);
336 return RUN_ALL_TESTS();
337 } else {
338 fprintf(stderr, "SKIPPED as Cuckoo table doesn't support Big Endian\n");
339 return 0;
340 }
341 }
342
343 #else
344 #include <stdio.h>
345
main(int,char **)346 int main(int /*argc*/, char** /*argv*/) {
347 fprintf(stderr, "SKIPPED as Cuckoo table is not supported in ROCKSDB_LITE\n");
348 return 0;
349 }
350
351 #endif // ROCKSDB_LITE
352