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 // Copyright (c) 2011 The LevelDB Authors. All rights reserved.
7 // Use of this source code is governed by a BSD-style license that can be
8 // found in the LICENSE file. See the AUTHORS file for names of contributors.
9 
10 #include "db/db_test_util.h"
11 #include "db/forward_iterator.h"
12 #include "rocksdb/env_encryption.h"
13 #include "rocksdb/utilities/object_registry.h"
14 
15 namespace ROCKSDB_NAMESPACE {
16 
17 // Special Env used to delay background operations
18 
SpecialEnv(Env * base)19 SpecialEnv::SpecialEnv(Env* base)
20     : EnvWrapper(base),
21       rnd_(301),
22       sleep_counter_(this),
23       addon_time_(0),
24       time_elapse_only_sleep_(false),
25       no_slowdown_(false) {
26   delay_sstable_sync_.store(false, std::memory_order_release);
27   drop_writes_.store(false, std::memory_order_release);
28   no_space_.store(false, std::memory_order_release);
29   non_writable_.store(false, std::memory_order_release);
30   count_random_reads_ = false;
31   count_sequential_reads_ = false;
32   manifest_sync_error_.store(false, std::memory_order_release);
33   manifest_write_error_.store(false, std::memory_order_release);
34   log_write_error_.store(false, std::memory_order_release);
35   random_file_open_counter_.store(0, std::memory_order_relaxed);
36   delete_count_.store(0, std::memory_order_relaxed);
37   num_open_wal_file_.store(0);
38   log_write_slowdown_ = 0;
39   bytes_written_ = 0;
40   sync_counter_ = 0;
41   non_writeable_rate_ = 0;
42   new_writable_count_ = 0;
43   non_writable_count_ = 0;
44   table_write_callback_ = nullptr;
45 }
46 #ifndef ROCKSDB_LITE
47 ROT13BlockCipher rot13Cipher_(16);
48 #endif  // ROCKSDB_LITE
49 
DBTestBase(const std::string path)50 DBTestBase::DBTestBase(const std::string path)
51     : mem_env_(nullptr), encrypted_env_(nullptr), option_config_(kDefault) {
52   Env* base_env = Env::Default();
53 #ifndef ROCKSDB_LITE
54   const char* test_env_uri = getenv("TEST_ENV_URI");
55   if (test_env_uri) {
56     Env* test_env = nullptr;
57     Status s = Env::LoadEnv(test_env_uri, &test_env, &env_guard_);
58     base_env = test_env;
59     EXPECT_OK(s);
60     EXPECT_NE(Env::Default(), base_env);
61   }
62 #endif  // !ROCKSDB_LITE
63   EXPECT_NE(nullptr, base_env);
64   if (getenv("MEM_ENV")) {
65     mem_env_ = new MockEnv(base_env);
66   }
67 #ifndef ROCKSDB_LITE
68   if (getenv("ENCRYPTED_ENV")) {
69     encrypted_env_ = NewEncryptedEnv(mem_env_ ? mem_env_ : base_env,
70                                      new CTREncryptionProvider(rot13Cipher_));
71   }
72 #endif  // !ROCKSDB_LITE
73   env_ = new SpecialEnv(encrypted_env_ ? encrypted_env_
74                                        : (mem_env_ ? mem_env_ : base_env));
75   env_->SetBackgroundThreads(1, Env::LOW);
76   env_->SetBackgroundThreads(1, Env::HIGH);
77   dbname_ = test::PerThreadDBPath(env_, path);
78   alternative_wal_dir_ = dbname_ + "/wal";
79   alternative_db_log_dir_ = dbname_ + "/db_log_dir";
80   auto options = CurrentOptions();
81   options.env = env_;
82   auto delete_options = options;
83   delete_options.wal_dir = alternative_wal_dir_;
84   EXPECT_OK(DestroyDB(dbname_, delete_options));
85   // Destroy it for not alternative WAL dir is used.
86   EXPECT_OK(DestroyDB(dbname_, options));
87   db_ = nullptr;
88   Reopen(options);
89   Random::GetTLSInstance()->Reset(0xdeadbeef);
90 }
91 
~DBTestBase()92 DBTestBase::~DBTestBase() {
93   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
94   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->LoadDependency({});
95   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
96   Close();
97   Options options;
98   options.db_paths.emplace_back(dbname_, 0);
99   options.db_paths.emplace_back(dbname_ + "_2", 0);
100   options.db_paths.emplace_back(dbname_ + "_3", 0);
101   options.db_paths.emplace_back(dbname_ + "_4", 0);
102   options.env = env_;
103 
104   if (getenv("KEEP_DB")) {
105     printf("DB is still at %s\n", dbname_.c_str());
106   } else {
107     EXPECT_OK(DestroyDB(dbname_, options));
108   }
109   delete env_;
110 }
111 
ShouldSkipOptions(int option_config,int skip_mask)112 bool DBTestBase::ShouldSkipOptions(int option_config, int skip_mask) {
113 #ifdef ROCKSDB_LITE
114     // These options are not supported in ROCKSDB_LITE
115     if (option_config == kHashSkipList ||
116         option_config == kPlainTableFirstBytePrefix ||
117         option_config == kPlainTableCappedPrefix ||
118         option_config == kPlainTableCappedPrefixNonMmap ||
119         option_config == kPlainTableAllBytesPrefix ||
120         option_config == kVectorRep || option_config == kHashLinkList ||
121         option_config == kUniversalCompaction ||
122         option_config == kUniversalCompactionMultiLevel ||
123         option_config == kUniversalSubcompactions ||
124         option_config == kFIFOCompaction ||
125         option_config == kConcurrentSkipList) {
126       return true;
127     }
128 #endif
129 
130     if ((skip_mask & kSkipUniversalCompaction) &&
131         (option_config == kUniversalCompaction ||
132          option_config == kUniversalCompactionMultiLevel ||
133          option_config == kUniversalSubcompactions)) {
134       return true;
135     }
136     if ((skip_mask & kSkipMergePut) && option_config == kMergePut) {
137       return true;
138     }
139     if ((skip_mask & kSkipNoSeekToLast) &&
140         (option_config == kHashLinkList || option_config == kHashSkipList)) {
141       return true;
142     }
143     if ((skip_mask & kSkipPlainTable) &&
144         (option_config == kPlainTableAllBytesPrefix ||
145          option_config == kPlainTableFirstBytePrefix ||
146          option_config == kPlainTableCappedPrefix ||
147          option_config == kPlainTableCappedPrefixNonMmap)) {
148       return true;
149     }
150     if ((skip_mask & kSkipHashIndex) &&
151         (option_config == kBlockBasedTableWithPrefixHashIndex ||
152          option_config == kBlockBasedTableWithWholeKeyHashIndex)) {
153       return true;
154     }
155     if ((skip_mask & kSkipFIFOCompaction) && option_config == kFIFOCompaction) {
156       return true;
157     }
158     if ((skip_mask & kSkipMmapReads) && option_config == kWalDirAndMmapReads) {
159       return true;
160     }
161     return false;
162 }
163 
164 // Switch to a fresh database with the next option configuration to
165 // test.  Return false if there are no more configurations to test.
ChangeOptions(int skip_mask)166 bool DBTestBase::ChangeOptions(int skip_mask) {
167   for (option_config_++; option_config_ < kEnd; option_config_++) {
168     if (ShouldSkipOptions(option_config_, skip_mask)) {
169       continue;
170     }
171     break;
172   }
173 
174   if (option_config_ >= kEnd) {
175     Destroy(last_options_);
176     return false;
177   } else {
178     auto options = CurrentOptions();
179     options.create_if_missing = true;
180     DestroyAndReopen(options);
181     return true;
182   }
183 }
184 
185 // Switch between different compaction styles.
ChangeCompactOptions()186 bool DBTestBase::ChangeCompactOptions() {
187   if (option_config_ == kDefault) {
188     option_config_ = kUniversalCompaction;
189     Destroy(last_options_);
190     auto options = CurrentOptions();
191     options.create_if_missing = true;
192     TryReopen(options);
193     return true;
194   } else if (option_config_ == kUniversalCompaction) {
195     option_config_ = kUniversalCompactionMultiLevel;
196     Destroy(last_options_);
197     auto options = CurrentOptions();
198     options.create_if_missing = true;
199     TryReopen(options);
200     return true;
201   } else if (option_config_ == kUniversalCompactionMultiLevel) {
202     option_config_ = kLevelSubcompactions;
203     Destroy(last_options_);
204     auto options = CurrentOptions();
205     assert(options.max_subcompactions > 1);
206     TryReopen(options);
207     return true;
208   } else if (option_config_ == kLevelSubcompactions) {
209     option_config_ = kUniversalSubcompactions;
210     Destroy(last_options_);
211     auto options = CurrentOptions();
212     assert(options.max_subcompactions > 1);
213     TryReopen(options);
214     return true;
215   } else {
216     return false;
217   }
218 }
219 
220 // Switch between different WAL settings
ChangeWalOptions()221 bool DBTestBase::ChangeWalOptions() {
222   if (option_config_ == kDefault) {
223     option_config_ = kDBLogDir;
224     Destroy(last_options_);
225     auto options = CurrentOptions();
226     Destroy(options);
227     options.create_if_missing = true;
228     TryReopen(options);
229     return true;
230   } else if (option_config_ == kDBLogDir) {
231     option_config_ = kWalDirAndMmapReads;
232     Destroy(last_options_);
233     auto options = CurrentOptions();
234     Destroy(options);
235     options.create_if_missing = true;
236     TryReopen(options);
237     return true;
238   } else if (option_config_ == kWalDirAndMmapReads) {
239     option_config_ = kRecycleLogFiles;
240     Destroy(last_options_);
241     auto options = CurrentOptions();
242     Destroy(options);
243     TryReopen(options);
244     return true;
245   } else {
246     return false;
247   }
248 }
249 
250 // Switch between different filter policy
251 // Jump from kDefault to kFilter to kFullFilter
ChangeFilterOptions()252 bool DBTestBase::ChangeFilterOptions() {
253   if (option_config_ == kDefault) {
254     option_config_ = kFilter;
255   } else if (option_config_ == kFilter) {
256     option_config_ = kFullFilterWithNewTableReaderForCompactions;
257   } else if (option_config_ == kFullFilterWithNewTableReaderForCompactions) {
258     option_config_ = kPartitionedFilterWithNewTableReaderForCompactions;
259   } else {
260     return false;
261   }
262   Destroy(last_options_);
263 
264   auto options = CurrentOptions();
265   options.create_if_missing = true;
266   TryReopen(options);
267   return true;
268 }
269 
270 // Switch between different DB options for file ingestion tests.
ChangeOptionsForFileIngestionTest()271 bool DBTestBase::ChangeOptionsForFileIngestionTest() {
272   if (option_config_ == kDefault) {
273     option_config_ = kUniversalCompaction;
274     Destroy(last_options_);
275     auto options = CurrentOptions();
276     options.create_if_missing = true;
277     TryReopen(options);
278     return true;
279   } else if (option_config_ == kUniversalCompaction) {
280     option_config_ = kUniversalCompactionMultiLevel;
281     Destroy(last_options_);
282     auto options = CurrentOptions();
283     options.create_if_missing = true;
284     TryReopen(options);
285     return true;
286   } else if (option_config_ == kUniversalCompactionMultiLevel) {
287     option_config_ = kLevelSubcompactions;
288     Destroy(last_options_);
289     auto options = CurrentOptions();
290     assert(options.max_subcompactions > 1);
291     TryReopen(options);
292     return true;
293   } else if (option_config_ == kLevelSubcompactions) {
294     option_config_ = kUniversalSubcompactions;
295     Destroy(last_options_);
296     auto options = CurrentOptions();
297     assert(options.max_subcompactions > 1);
298     TryReopen(options);
299     return true;
300   } else if (option_config_ == kUniversalSubcompactions) {
301     option_config_ = kDirectIO;
302     Destroy(last_options_);
303     auto options = CurrentOptions();
304     TryReopen(options);
305     return true;
306   } else {
307     return false;
308   }
309 }
310 
311 // Return the current option configuration.
CurrentOptions(const anon::OptionsOverride & options_override) const312 Options DBTestBase::CurrentOptions(
313     const anon::OptionsOverride& options_override) const {
314   return GetOptions(option_config_, GetDefaultOptions(), options_override);
315 }
316 
CurrentOptions(const Options & default_options,const anon::OptionsOverride & options_override) const317 Options DBTestBase::CurrentOptions(
318     const Options& default_options,
319     const anon::OptionsOverride& options_override) const {
320   return GetOptions(option_config_, default_options, options_override);
321 }
322 
GetDefaultOptions()323 Options DBTestBase::GetDefaultOptions() {
324   Options options;
325   options.write_buffer_size = 4090 * 4096;
326   options.target_file_size_base = 2 * 1024 * 1024;
327   options.max_bytes_for_level_base = 10 * 1024 * 1024;
328   options.max_open_files = 5000;
329   options.wal_recovery_mode = WALRecoveryMode::kTolerateCorruptedTailRecords;
330   options.compaction_pri = CompactionPri::kByCompensatedSize;
331   return options;
332 }
333 
GetOptions(int option_config,const Options & default_options,const anon::OptionsOverride & options_override) const334 Options DBTestBase::GetOptions(
335     int option_config, const Options& default_options,
336     const anon::OptionsOverride& options_override) const {
337   // this redundant copy is to minimize code change w/o having lint error.
338   Options options = default_options;
339   BlockBasedTableOptions table_options;
340   bool set_block_based_table_factory = true;
341 #if !defined(OS_MACOSX) && !defined(OS_WIN) && !defined(OS_SOLARIS) && \
342     !defined(OS_AIX)
343   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearCallBack(
344       "NewRandomAccessFile:O_DIRECT");
345   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearCallBack(
346       "NewWritableFile:O_DIRECT");
347 #endif
348 
349   bool can_allow_mmap = IsMemoryMappedAccessSupported();
350   switch (option_config) {
351 #ifndef ROCKSDB_LITE
352     case kHashSkipList:
353       options.prefix_extractor.reset(NewFixedPrefixTransform(1));
354       options.memtable_factory.reset(NewHashSkipListRepFactory(16));
355       options.allow_concurrent_memtable_write = false;
356       options.unordered_write = false;
357       break;
358     case kPlainTableFirstBytePrefix:
359       options.table_factory.reset(new PlainTableFactory());
360       options.prefix_extractor.reset(NewFixedPrefixTransform(1));
361       options.allow_mmap_reads = can_allow_mmap;
362       options.max_sequential_skip_in_iterations = 999999;
363       set_block_based_table_factory = false;
364       break;
365     case kPlainTableCappedPrefix:
366       options.table_factory.reset(new PlainTableFactory());
367       options.prefix_extractor.reset(NewCappedPrefixTransform(8));
368       options.allow_mmap_reads = can_allow_mmap;
369       options.max_sequential_skip_in_iterations = 999999;
370       set_block_based_table_factory = false;
371       break;
372     case kPlainTableCappedPrefixNonMmap:
373       options.table_factory.reset(new PlainTableFactory());
374       options.prefix_extractor.reset(NewCappedPrefixTransform(8));
375       options.allow_mmap_reads = false;
376       options.max_sequential_skip_in_iterations = 999999;
377       set_block_based_table_factory = false;
378       break;
379     case kPlainTableAllBytesPrefix:
380       options.table_factory.reset(new PlainTableFactory());
381       options.prefix_extractor.reset(NewNoopTransform());
382       options.allow_mmap_reads = can_allow_mmap;
383       options.max_sequential_skip_in_iterations = 999999;
384       set_block_based_table_factory = false;
385       break;
386     case kVectorRep:
387       options.memtable_factory.reset(new VectorRepFactory(100));
388       options.allow_concurrent_memtable_write = false;
389       options.unordered_write = false;
390       break;
391     case kHashLinkList:
392       options.prefix_extractor.reset(NewFixedPrefixTransform(1));
393       options.memtable_factory.reset(
394           NewHashLinkListRepFactory(4, 0, 3, true, 4));
395       options.allow_concurrent_memtable_write = false;
396       options.unordered_write = false;
397       break;
398       case kDirectIO: {
399         options.use_direct_reads = true;
400         options.use_direct_io_for_flush_and_compaction = true;
401         options.compaction_readahead_size = 2 * 1024 * 1024;
402   #if !defined(OS_MACOSX) && !defined(OS_WIN) && !defined(OS_SOLARIS) && \
403       !defined(OS_AIX) && !defined(OS_OPENBSD)
404         ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
405             "NewWritableFile:O_DIRECT", [&](void* arg) {
406               int* val = static_cast<int*>(arg);
407               *val &= ~O_DIRECT;
408             });
409         ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
410             "NewRandomAccessFile:O_DIRECT", [&](void* arg) {
411               int* val = static_cast<int*>(arg);
412               *val &= ~O_DIRECT;
413             });
414         ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->EnableProcessing();
415 #endif
416         break;
417       }
418 #endif  // ROCKSDB_LITE
419     case kMergePut:
420       options.merge_operator = MergeOperators::CreatePutOperator();
421       break;
422     case kFilter:
423       table_options.filter_policy.reset(NewBloomFilterPolicy(10, true));
424       break;
425     case kFullFilterWithNewTableReaderForCompactions:
426       table_options.filter_policy.reset(NewBloomFilterPolicy(10, false));
427       options.new_table_reader_for_compaction_inputs = true;
428       options.compaction_readahead_size = 10 * 1024 * 1024;
429       break;
430     case kPartitionedFilterWithNewTableReaderForCompactions:
431       table_options.filter_policy.reset(NewBloomFilterPolicy(10, false));
432       table_options.partition_filters = true;
433       table_options.index_type =
434           BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch;
435       options.new_table_reader_for_compaction_inputs = true;
436       options.compaction_readahead_size = 10 * 1024 * 1024;
437       break;
438     case kUncompressed:
439       options.compression = kNoCompression;
440       break;
441     case kNumLevel_3:
442       options.num_levels = 3;
443       break;
444     case kDBLogDir:
445       options.db_log_dir = alternative_db_log_dir_;
446       break;
447     case kWalDirAndMmapReads:
448       options.wal_dir = alternative_wal_dir_;
449       // mmap reads should be orthogonal to WalDir setting, so we piggyback to
450       // this option config to test mmap reads as well
451       options.allow_mmap_reads = can_allow_mmap;
452       break;
453     case kManifestFileSize:
454       options.max_manifest_file_size = 50;  // 50 bytes
455       break;
456     case kPerfOptions:
457       options.soft_rate_limit = 2.0;
458       options.delayed_write_rate = 8 * 1024 * 1024;
459       options.report_bg_io_stats = true;
460       // TODO(3.13) -- test more options
461       break;
462     case kUniversalCompaction:
463       options.compaction_style = kCompactionStyleUniversal;
464       options.num_levels = 1;
465       break;
466     case kUniversalCompactionMultiLevel:
467       options.compaction_style = kCompactionStyleUniversal;
468       options.num_levels = 8;
469       break;
470     case kCompressedBlockCache:
471       options.allow_mmap_writes = can_allow_mmap;
472       table_options.block_cache_compressed = NewLRUCache(8 * 1024 * 1024);
473       break;
474     case kInfiniteMaxOpenFiles:
475       options.max_open_files = -1;
476       break;
477     case kxxHashChecksum: {
478       table_options.checksum = kxxHash;
479       break;
480     }
481     case kxxHash64Checksum: {
482       table_options.checksum = kxxHash64;
483       break;
484     }
485     case kFIFOCompaction: {
486       options.compaction_style = kCompactionStyleFIFO;
487       break;
488     }
489     case kBlockBasedTableWithPrefixHashIndex: {
490       table_options.index_type = BlockBasedTableOptions::kHashSearch;
491       options.prefix_extractor.reset(NewFixedPrefixTransform(1));
492       break;
493     }
494     case kBlockBasedTableWithWholeKeyHashIndex: {
495       table_options.index_type = BlockBasedTableOptions::kHashSearch;
496       options.prefix_extractor.reset(NewNoopTransform());
497       break;
498     }
499     case kBlockBasedTableWithPartitionedIndex: {
500       table_options.index_type = BlockBasedTableOptions::kTwoLevelIndexSearch;
501       options.prefix_extractor.reset(NewNoopTransform());
502       break;
503     }
504     case kBlockBasedTableWithPartitionedIndexFormat4: {
505       table_options.format_version = 4;
506       // Format 4 changes the binary index format. Since partitioned index is a
507       // super-set of simple indexes, we are also using kTwoLevelIndexSearch to
508       // test this format.
509       table_options.index_type = BlockBasedTableOptions::kTwoLevelIndexSearch;
510       // The top-level index in partition filters are also affected by format 4.
511       table_options.filter_policy.reset(NewBloomFilterPolicy(10, false));
512       table_options.partition_filters = true;
513       table_options.index_block_restart_interval = 8;
514       break;
515     }
516     case kBlockBasedTableWithIndexRestartInterval: {
517       table_options.index_block_restart_interval = 8;
518       break;
519     }
520     case kOptimizeFiltersForHits: {
521       options.optimize_filters_for_hits = true;
522       set_block_based_table_factory = true;
523       break;
524     }
525     case kRowCache: {
526       options.row_cache = NewLRUCache(1024 * 1024);
527       break;
528     }
529     case kRecycleLogFiles: {
530       options.recycle_log_file_num = 2;
531       break;
532     }
533     case kLevelSubcompactions: {
534       options.max_subcompactions = 4;
535       break;
536     }
537     case kUniversalSubcompactions: {
538       options.compaction_style = kCompactionStyleUniversal;
539       options.num_levels = 8;
540       options.max_subcompactions = 4;
541       break;
542     }
543     case kConcurrentSkipList: {
544       options.allow_concurrent_memtable_write = true;
545       options.enable_write_thread_adaptive_yield = true;
546       break;
547     }
548     case kPipelinedWrite: {
549       options.enable_pipelined_write = true;
550       break;
551     }
552     case kConcurrentWALWrites: {
553       // This options optimize 2PC commit path
554       options.two_write_queues = true;
555       options.manual_wal_flush = true;
556       break;
557     }
558     case kUnorderedWrite: {
559       options.allow_concurrent_memtable_write = false;
560       options.unordered_write = false;
561       break;
562     }
563 
564     default:
565       break;
566   }
567 
568   if (options_override.filter_policy) {
569     table_options.filter_policy = options_override.filter_policy;
570     table_options.partition_filters = options_override.partition_filters;
571     table_options.metadata_block_size = options_override.metadata_block_size;
572   }
573   if (set_block_based_table_factory) {
574     options.table_factory.reset(NewBlockBasedTableFactory(table_options));
575   }
576   options.env = env_;
577   options.create_if_missing = true;
578   options.fail_if_options_file_error = true;
579   return options;
580 }
581 
CreateColumnFamilies(const std::vector<std::string> & cfs,const Options & options)582 void DBTestBase::CreateColumnFamilies(const std::vector<std::string>& cfs,
583                                       const Options& options) {
584   ColumnFamilyOptions cf_opts(options);
585   size_t cfi = handles_.size();
586   handles_.resize(cfi + cfs.size());
587   for (auto cf : cfs) {
588     Status s = db_->CreateColumnFamily(cf_opts, cf, &handles_[cfi++]);
589     ASSERT_OK(s);
590   }
591 }
592 
CreateAndReopenWithCF(const std::vector<std::string> & cfs,const Options & options)593 void DBTestBase::CreateAndReopenWithCF(const std::vector<std::string>& cfs,
594                                        const Options& options) {
595   CreateColumnFamilies(cfs, options);
596   std::vector<std::string> cfs_plus_default = cfs;
597   cfs_plus_default.insert(cfs_plus_default.begin(), kDefaultColumnFamilyName);
598   ReopenWithColumnFamilies(cfs_plus_default, options);
599 }
600 
ReopenWithColumnFamilies(const std::vector<std::string> & cfs,const std::vector<Options> & options)601 void DBTestBase::ReopenWithColumnFamilies(const std::vector<std::string>& cfs,
602                                           const std::vector<Options>& options) {
603   ASSERT_OK(TryReopenWithColumnFamilies(cfs, options));
604 }
605 
ReopenWithColumnFamilies(const std::vector<std::string> & cfs,const Options & options)606 void DBTestBase::ReopenWithColumnFamilies(const std::vector<std::string>& cfs,
607                                           const Options& options) {
608   ASSERT_OK(TryReopenWithColumnFamilies(cfs, options));
609 }
610 
TryReopenWithColumnFamilies(const std::vector<std::string> & cfs,const std::vector<Options> & options)611 Status DBTestBase::TryReopenWithColumnFamilies(
612     const std::vector<std::string>& cfs, const std::vector<Options>& options) {
613   Close();
614   EXPECT_EQ(cfs.size(), options.size());
615   std::vector<ColumnFamilyDescriptor> column_families;
616   for (size_t i = 0; i < cfs.size(); ++i) {
617     column_families.push_back(ColumnFamilyDescriptor(cfs[i], options[i]));
618   }
619   DBOptions db_opts = DBOptions(options[0]);
620   last_options_ = options[0];
621   return DB::Open(db_opts, dbname_, column_families, &handles_, &db_);
622 }
623 
TryReopenWithColumnFamilies(const std::vector<std::string> & cfs,const Options & options)624 Status DBTestBase::TryReopenWithColumnFamilies(
625     const std::vector<std::string>& cfs, const Options& options) {
626   Close();
627   std::vector<Options> v_opts(cfs.size(), options);
628   return TryReopenWithColumnFamilies(cfs, v_opts);
629 }
630 
Reopen(const Options & options)631 void DBTestBase::Reopen(const Options& options) {
632   ASSERT_OK(TryReopen(options));
633 }
634 
Close()635 void DBTestBase::Close() {
636   for (auto h : handles_) {
637     db_->DestroyColumnFamilyHandle(h);
638   }
639   handles_.clear();
640   delete db_;
641   db_ = nullptr;
642 }
643 
DestroyAndReopen(const Options & options)644 void DBTestBase::DestroyAndReopen(const Options& options) {
645   // Destroy using last options
646   Destroy(last_options_);
647   ASSERT_OK(TryReopen(options));
648 }
649 
Destroy(const Options & options,bool delete_cf_paths)650 void DBTestBase::Destroy(const Options& options, bool delete_cf_paths) {
651   std::vector<ColumnFamilyDescriptor> column_families;
652   if (delete_cf_paths) {
653     for (size_t i = 0; i < handles_.size(); ++i) {
654       ColumnFamilyDescriptor cfdescriptor;
655       handles_[i]->GetDescriptor(&cfdescriptor);
656       column_families.push_back(cfdescriptor);
657     }
658   }
659   Close();
660   ASSERT_OK(DestroyDB(dbname_, options, column_families));
661 }
662 
ReadOnlyReopen(const Options & options)663 Status DBTestBase::ReadOnlyReopen(const Options& options) {
664   return DB::OpenForReadOnly(options, dbname_, &db_);
665 }
666 
TryReopen(const Options & options)667 Status DBTestBase::TryReopen(const Options& options) {
668   Close();
669   last_options_.table_factory.reset();
670   // Note: operator= is an unsafe approach here since it destructs
671   // std::shared_ptr in the same order of their creation, in contrast to
672   // destructors which destructs them in the opposite order of creation. One
673   // particular problme is that the cache destructor might invoke callback
674   // functions that use Option members such as statistics. To work around this
675   // problem, we manually call destructor of table_facotry which eventually
676   // clears the block cache.
677   last_options_ = options;
678   return DB::Open(options, dbname_, &db_);
679 }
680 
IsDirectIOSupported()681 bool DBTestBase::IsDirectIOSupported() {
682   return test::IsDirectIOSupported(env_, dbname_);
683 }
684 
IsMemoryMappedAccessSupported() const685 bool DBTestBase::IsMemoryMappedAccessSupported() const {
686   return (!encrypted_env_);
687 }
688 
Flush(int cf)689 Status DBTestBase::Flush(int cf) {
690   if (cf == 0) {
691     return db_->Flush(FlushOptions());
692   } else {
693     return db_->Flush(FlushOptions(), handles_[cf]);
694   }
695 }
696 
Flush(const std::vector<int> & cf_ids)697 Status DBTestBase::Flush(const std::vector<int>& cf_ids) {
698   std::vector<ColumnFamilyHandle*> cfhs;
699   std::for_each(cf_ids.begin(), cf_ids.end(),
700                 [&cfhs, this](int id) { cfhs.emplace_back(handles_[id]); });
701   return db_->Flush(FlushOptions(), cfhs);
702 }
703 
Put(const Slice & k,const Slice & v,WriteOptions wo)704 Status DBTestBase::Put(const Slice& k, const Slice& v, WriteOptions wo) {
705   if (kMergePut == option_config_) {
706     return db_->Merge(wo, k, v);
707   } else {
708     return db_->Put(wo, k, v);
709   }
710 }
711 
Put(int cf,const Slice & k,const Slice & v,WriteOptions wo)712 Status DBTestBase::Put(int cf, const Slice& k, const Slice& v,
713                        WriteOptions wo) {
714   if (kMergePut == option_config_) {
715     return db_->Merge(wo, handles_[cf], k, v);
716   } else {
717     return db_->Put(wo, handles_[cf], k, v);
718   }
719 }
720 
Merge(const Slice & k,const Slice & v,WriteOptions wo)721 Status DBTestBase::Merge(const Slice& k, const Slice& v, WriteOptions wo) {
722   return db_->Merge(wo, k, v);
723 }
724 
Merge(int cf,const Slice & k,const Slice & v,WriteOptions wo)725 Status DBTestBase::Merge(int cf, const Slice& k, const Slice& v,
726                          WriteOptions wo) {
727   return db_->Merge(wo, handles_[cf], k, v);
728 }
729 
Delete(const std::string & k)730 Status DBTestBase::Delete(const std::string& k) {
731   return db_->Delete(WriteOptions(), k);
732 }
733 
Delete(int cf,const std::string & k)734 Status DBTestBase::Delete(int cf, const std::string& k) {
735   return db_->Delete(WriteOptions(), handles_[cf], k);
736 }
737 
SingleDelete(const std::string & k)738 Status DBTestBase::SingleDelete(const std::string& k) {
739   return db_->SingleDelete(WriteOptions(), k);
740 }
741 
SingleDelete(int cf,const std::string & k)742 Status DBTestBase::SingleDelete(int cf, const std::string& k) {
743   return db_->SingleDelete(WriteOptions(), handles_[cf], k);
744 }
745 
SetPreserveDeletesSequenceNumber(SequenceNumber sn)746 bool DBTestBase::SetPreserveDeletesSequenceNumber(SequenceNumber sn) {
747   return db_->SetPreserveDeletesSequenceNumber(sn);
748 }
749 
Get(const std::string & k,const Snapshot * snapshot)750 std::string DBTestBase::Get(const std::string& k, const Snapshot* snapshot) {
751   ReadOptions options;
752   options.verify_checksums = true;
753   options.snapshot = snapshot;
754   std::string result;
755   Status s = db_->Get(options, k, &result);
756   if (s.IsNotFound()) {
757     result = "NOT_FOUND";
758   } else if (!s.ok()) {
759     result = s.ToString();
760   }
761   return result;
762 }
763 
Get(int cf,const std::string & k,const Snapshot * snapshot)764 std::string DBTestBase::Get(int cf, const std::string& k,
765                             const Snapshot* snapshot) {
766   ReadOptions options;
767   options.verify_checksums = true;
768   options.snapshot = snapshot;
769   std::string result;
770   Status s = db_->Get(options, handles_[cf], k, &result);
771   if (s.IsNotFound()) {
772     result = "NOT_FOUND";
773   } else if (!s.ok()) {
774     result = s.ToString();
775   }
776   return result;
777 }
778 
MultiGet(std::vector<int> cfs,const std::vector<std::string> & k,const Snapshot * snapshot,const bool batched)779 std::vector<std::string> DBTestBase::MultiGet(std::vector<int> cfs,
780                                               const std::vector<std::string>& k,
781                                               const Snapshot* snapshot,
782                                               const bool batched) {
783   ReadOptions options;
784   options.verify_checksums = true;
785   options.snapshot = snapshot;
786   std::vector<ColumnFamilyHandle*> handles;
787   std::vector<Slice> keys;
788   std::vector<std::string> result;
789 
790   for (unsigned int i = 0; i < cfs.size(); ++i) {
791     handles.push_back(handles_[cfs[i]]);
792     keys.push_back(k[i]);
793   }
794   std::vector<Status> s;
795   if (!batched) {
796     s = db_->MultiGet(options, handles, keys, &result);
797     for (unsigned int i = 0; i < s.size(); ++i) {
798       if (s[i].IsNotFound()) {
799         result[i] = "NOT_FOUND";
800       } else if (!s[i].ok()) {
801         result[i] = s[i].ToString();
802       }
803     }
804   } else {
805     std::vector<PinnableSlice> pin_values(cfs.size());
806     result.resize(cfs.size());
807     s.resize(cfs.size());
808     db_->MultiGet(options, cfs.size(), handles.data(), keys.data(),
809                   pin_values.data(), s.data());
810     for (unsigned int i = 0; i < s.size(); ++i) {
811       if (s[i].IsNotFound()) {
812         result[i] = "NOT_FOUND";
813       } else if (!s[i].ok()) {
814         result[i] = s[i].ToString();
815       } else {
816         result[i].assign(pin_values[i].data(), pin_values[i].size());
817       }
818     }
819   }
820   return result;
821 }
822 
MultiGet(const std::vector<std::string> & k,const Snapshot * snapshot)823 std::vector<std::string> DBTestBase::MultiGet(const std::vector<std::string>& k,
824                                               const Snapshot* snapshot) {
825   ReadOptions options;
826   options.verify_checksums = true;
827   options.snapshot = snapshot;
828   std::vector<Slice> keys;
829   std::vector<std::string> result;
830   std::vector<Status> statuses(k.size());
831   std::vector<PinnableSlice> pin_values(k.size());
832 
833   for (unsigned int i = 0; i < k.size(); ++i) {
834     keys.push_back(k[i]);
835   }
836   db_->MultiGet(options, dbfull()->DefaultColumnFamily(), keys.size(),
837                 keys.data(), pin_values.data(), statuses.data());
838   result.resize(k.size());
839   for (auto iter = result.begin(); iter != result.end(); ++iter) {
840     iter->assign(pin_values[iter - result.begin()].data(),
841                  pin_values[iter - result.begin()].size());
842   }
843   for (unsigned int i = 0; i < statuses.size(); ++i) {
844     if (statuses[i].IsNotFound()) {
845       result[i] = "NOT_FOUND";
846     }
847   }
848   return result;
849 }
850 
Get(const std::string & k,PinnableSlice * v)851 Status DBTestBase::Get(const std::string& k, PinnableSlice* v) {
852   ReadOptions options;
853   options.verify_checksums = true;
854   Status s = dbfull()->Get(options, dbfull()->DefaultColumnFamily(), k, v);
855   return s;
856 }
857 
GetNumSnapshots()858 uint64_t DBTestBase::GetNumSnapshots() {
859   uint64_t int_num;
860   EXPECT_TRUE(dbfull()->GetIntProperty("rocksdb.num-snapshots", &int_num));
861   return int_num;
862 }
863 
GetTimeOldestSnapshots()864 uint64_t DBTestBase::GetTimeOldestSnapshots() {
865   uint64_t int_num;
866   EXPECT_TRUE(
867       dbfull()->GetIntProperty("rocksdb.oldest-snapshot-time", &int_num));
868   return int_num;
869 }
870 
GetSequenceOldestSnapshots()871 uint64_t DBTestBase::GetSequenceOldestSnapshots() {
872   uint64_t int_num;
873   EXPECT_TRUE(
874       dbfull()->GetIntProperty("rocksdb.oldest-snapshot-sequence", &int_num));
875   return int_num;
876 }
877 
878 // Return a string that contains all key,value pairs in order,
879 // formatted like "(k1->v1)(k2->v2)".
Contents(int cf)880 std::string DBTestBase::Contents(int cf) {
881   std::vector<std::string> forward;
882   std::string result;
883   Iterator* iter = (cf == 0) ? db_->NewIterator(ReadOptions())
884                              : db_->NewIterator(ReadOptions(), handles_[cf]);
885   for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
886     std::string s = IterStatus(iter);
887     result.push_back('(');
888     result.append(s);
889     result.push_back(')');
890     forward.push_back(s);
891   }
892 
893   // Check reverse iteration results are the reverse of forward results
894   unsigned int matched = 0;
895   for (iter->SeekToLast(); iter->Valid(); iter->Prev()) {
896     EXPECT_LT(matched, forward.size());
897     EXPECT_EQ(IterStatus(iter), forward[forward.size() - matched - 1]);
898     matched++;
899   }
900   EXPECT_EQ(matched, forward.size());
901 
902   delete iter;
903   return result;
904 }
905 
AllEntriesFor(const Slice & user_key,int cf)906 std::string DBTestBase::AllEntriesFor(const Slice& user_key, int cf) {
907   Arena arena;
908   auto options = CurrentOptions();
909   InternalKeyComparator icmp(options.comparator);
910   ReadRangeDelAggregator range_del_agg(&icmp,
911                                        kMaxSequenceNumber /* upper_bound */);
912   ScopedArenaIterator iter;
913   if (cf == 0) {
914     iter.set(dbfull()->NewInternalIterator(&arena, &range_del_agg,
915                                            kMaxSequenceNumber));
916   } else {
917     iter.set(dbfull()->NewInternalIterator(&arena, &range_del_agg,
918                                            kMaxSequenceNumber, handles_[cf]));
919   }
920   InternalKey target(user_key, kMaxSequenceNumber, kTypeValue);
921   iter->Seek(target.Encode());
922   std::string result;
923   if (!iter->status().ok()) {
924     result = iter->status().ToString();
925   } else {
926     result = "[ ";
927     bool first = true;
928     while (iter->Valid()) {
929       ParsedInternalKey ikey(Slice(), 0, kTypeValue);
930       if (!ParseInternalKey(iter->key(), &ikey)) {
931         result += "CORRUPTED";
932       } else {
933         if (!last_options_.comparator->Equal(ikey.user_key, user_key)) {
934           break;
935         }
936         if (!first) {
937           result += ", ";
938         }
939         first = false;
940         switch (ikey.type) {
941           case kTypeValue:
942             result += iter->value().ToString();
943             break;
944           case kTypeMerge:
945             // keep it the same as kTypeValue for testing kMergePut
946             result += iter->value().ToString();
947             break;
948           case kTypeDeletion:
949             result += "DEL";
950             break;
951           case kTypeSingleDeletion:
952             result += "SDEL";
953             break;
954           default:
955             assert(false);
956             break;
957         }
958       }
959       iter->Next();
960     }
961     if (!first) {
962       result += " ";
963     }
964     result += "]";
965   }
966   return result;
967 }
968 
969 #ifndef ROCKSDB_LITE
NumSortedRuns(int cf)970 int DBTestBase::NumSortedRuns(int cf) {
971   ColumnFamilyMetaData cf_meta;
972   if (cf == 0) {
973     db_->GetColumnFamilyMetaData(&cf_meta);
974   } else {
975     db_->GetColumnFamilyMetaData(handles_[cf], &cf_meta);
976   }
977   int num_sr = static_cast<int>(cf_meta.levels[0].files.size());
978   for (size_t i = 1U; i < cf_meta.levels.size(); i++) {
979     if (cf_meta.levels[i].files.size() > 0) {
980       num_sr++;
981     }
982   }
983   return num_sr;
984 }
985 
TotalSize(int cf)986 uint64_t DBTestBase::TotalSize(int cf) {
987   ColumnFamilyMetaData cf_meta;
988   if (cf == 0) {
989     db_->GetColumnFamilyMetaData(&cf_meta);
990   } else {
991     db_->GetColumnFamilyMetaData(handles_[cf], &cf_meta);
992   }
993   return cf_meta.size;
994 }
995 
SizeAtLevel(int level)996 uint64_t DBTestBase::SizeAtLevel(int level) {
997   std::vector<LiveFileMetaData> metadata;
998   db_->GetLiveFilesMetaData(&metadata);
999   uint64_t sum = 0;
1000   for (const auto& m : metadata) {
1001     if (m.level == level) {
1002       sum += m.size;
1003     }
1004   }
1005   return sum;
1006 }
1007 
TotalLiveFiles(int cf)1008 size_t DBTestBase::TotalLiveFiles(int cf) {
1009   ColumnFamilyMetaData cf_meta;
1010   if (cf == 0) {
1011     db_->GetColumnFamilyMetaData(&cf_meta);
1012   } else {
1013     db_->GetColumnFamilyMetaData(handles_[cf], &cf_meta);
1014   }
1015   size_t num_files = 0;
1016   for (auto& level : cf_meta.levels) {
1017     num_files += level.files.size();
1018   }
1019   return num_files;
1020 }
1021 
CountLiveFiles()1022 size_t DBTestBase::CountLiveFiles() {
1023   std::vector<LiveFileMetaData> metadata;
1024   db_->GetLiveFilesMetaData(&metadata);
1025   return metadata.size();
1026 }
1027 
NumTableFilesAtLevel(int level,int cf)1028 int DBTestBase::NumTableFilesAtLevel(int level, int cf) {
1029   std::string property;
1030   if (cf == 0) {
1031     // default cfd
1032     EXPECT_TRUE(db_->GetProperty(
1033         "rocksdb.num-files-at-level" + NumberToString(level), &property));
1034   } else {
1035     EXPECT_TRUE(db_->GetProperty(
1036         handles_[cf], "rocksdb.num-files-at-level" + NumberToString(level),
1037         &property));
1038   }
1039   return atoi(property.c_str());
1040 }
1041 
CompressionRatioAtLevel(int level,int cf)1042 double DBTestBase::CompressionRatioAtLevel(int level, int cf) {
1043   std::string property;
1044   if (cf == 0) {
1045     // default cfd
1046     EXPECT_TRUE(db_->GetProperty(
1047         "rocksdb.compression-ratio-at-level" + NumberToString(level),
1048         &property));
1049   } else {
1050     EXPECT_TRUE(db_->GetProperty(
1051         handles_[cf],
1052         "rocksdb.compression-ratio-at-level" + NumberToString(level),
1053         &property));
1054   }
1055   return std::stod(property);
1056 }
1057 
TotalTableFiles(int cf,int levels)1058 int DBTestBase::TotalTableFiles(int cf, int levels) {
1059   if (levels == -1) {
1060     levels = (cf == 0) ? db_->NumberLevels() : db_->NumberLevels(handles_[1]);
1061   }
1062   int result = 0;
1063   for (int level = 0; level < levels; level++) {
1064     result += NumTableFilesAtLevel(level, cf);
1065   }
1066   return result;
1067 }
1068 
1069 // Return spread of files per level
FilesPerLevel(int cf)1070 std::string DBTestBase::FilesPerLevel(int cf) {
1071   int num_levels =
1072       (cf == 0) ? db_->NumberLevels() : db_->NumberLevels(handles_[1]);
1073   std::string result;
1074   size_t last_non_zero_offset = 0;
1075   for (int level = 0; level < num_levels; level++) {
1076     int f = NumTableFilesAtLevel(level, cf);
1077     char buf[100];
1078     snprintf(buf, sizeof(buf), "%s%d", (level ? "," : ""), f);
1079     result += buf;
1080     if (f > 0) {
1081       last_non_zero_offset = result.size();
1082     }
1083   }
1084   result.resize(last_non_zero_offset);
1085   return result;
1086 }
1087 #endif  // !ROCKSDB_LITE
1088 
CountFiles()1089 size_t DBTestBase::CountFiles() {
1090   std::vector<std::string> files;
1091   env_->GetChildren(dbname_, &files);
1092 
1093   std::vector<std::string> logfiles;
1094   if (dbname_ != last_options_.wal_dir) {
1095     env_->GetChildren(last_options_.wal_dir, &logfiles);
1096   }
1097 
1098   return files.size() + logfiles.size();
1099 }
1100 
Size(const Slice & start,const Slice & limit,int cf)1101 uint64_t DBTestBase::Size(const Slice& start, const Slice& limit, int cf) {
1102   Range r(start, limit);
1103   uint64_t size;
1104   if (cf == 0) {
1105     db_->GetApproximateSizes(&r, 1, &size);
1106   } else {
1107     db_->GetApproximateSizes(handles_[1], &r, 1, &size);
1108   }
1109   return size;
1110 }
1111 
Compact(int cf,const Slice & start,const Slice & limit,uint32_t target_path_id)1112 void DBTestBase::Compact(int cf, const Slice& start, const Slice& limit,
1113                          uint32_t target_path_id) {
1114   CompactRangeOptions compact_options;
1115   compact_options.target_path_id = target_path_id;
1116   ASSERT_OK(db_->CompactRange(compact_options, handles_[cf], &start, &limit));
1117 }
1118 
Compact(int cf,const Slice & start,const Slice & limit)1119 void DBTestBase::Compact(int cf, const Slice& start, const Slice& limit) {
1120   ASSERT_OK(
1121       db_->CompactRange(CompactRangeOptions(), handles_[cf], &start, &limit));
1122 }
1123 
Compact(const Slice & start,const Slice & limit)1124 void DBTestBase::Compact(const Slice& start, const Slice& limit) {
1125   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), &start, &limit));
1126 }
1127 
1128 // Do n memtable compactions, each of which produces an sstable
1129 // covering the range [small,large].
MakeTables(int n,const std::string & small,const std::string & large,int cf)1130 void DBTestBase::MakeTables(int n, const std::string& small,
1131                             const std::string& large, int cf) {
1132   for (int i = 0; i < n; i++) {
1133     ASSERT_OK(Put(cf, small, "begin"));
1134     ASSERT_OK(Put(cf, large, "end"));
1135     ASSERT_OK(Flush(cf));
1136     MoveFilesToLevel(n - i - 1, cf);
1137   }
1138 }
1139 
1140 // Prevent pushing of new sstables into deeper levels by adding
1141 // tables that cover a specified range to all levels.
FillLevels(const std::string & smallest,const std::string & largest,int cf)1142 void DBTestBase::FillLevels(const std::string& smallest,
1143                             const std::string& largest, int cf) {
1144   MakeTables(db_->NumberLevels(handles_[cf]), smallest, largest, cf);
1145 }
1146 
MoveFilesToLevel(int level,int cf)1147 void DBTestBase::MoveFilesToLevel(int level, int cf) {
1148   for (int l = 0; l < level; ++l) {
1149     if (cf > 0) {
1150       dbfull()->TEST_CompactRange(l, nullptr, nullptr, handles_[cf]);
1151     } else {
1152       dbfull()->TEST_CompactRange(l, nullptr, nullptr);
1153     }
1154   }
1155 }
1156 
1157 #ifndef ROCKSDB_LITE
DumpFileCounts(const char * label)1158 void DBTestBase::DumpFileCounts(const char* label) {
1159   fprintf(stderr, "---\n%s:\n", label);
1160   fprintf(stderr, "maxoverlap: %" PRIu64 "\n",
1161           dbfull()->TEST_MaxNextLevelOverlappingBytes());
1162   for (int level = 0; level < db_->NumberLevels(); level++) {
1163     int num = NumTableFilesAtLevel(level);
1164     if (num > 0) {
1165       fprintf(stderr, "  level %3d : %d files\n", level, num);
1166     }
1167   }
1168 }
1169 #endif  // !ROCKSDB_LITE
1170 
DumpSSTableList()1171 std::string DBTestBase::DumpSSTableList() {
1172   std::string property;
1173   db_->GetProperty("rocksdb.sstables", &property);
1174   return property;
1175 }
1176 
GetSstFiles(Env * env,std::string path,std::vector<std::string> * files)1177 void DBTestBase::GetSstFiles(Env* env, std::string path,
1178                              std::vector<std::string>* files) {
1179   env->GetChildren(path, files);
1180 
1181   files->erase(
1182       std::remove_if(files->begin(), files->end(), [](std::string name) {
1183         uint64_t number;
1184         FileType type;
1185         return !(ParseFileName(name, &number, &type) && type == kTableFile);
1186       }), files->end());
1187 }
1188 
GetSstFileCount(std::string path)1189 int DBTestBase::GetSstFileCount(std::string path) {
1190   std::vector<std::string> files;
1191   DBTestBase::GetSstFiles(env_, path, &files);
1192   return static_cast<int>(files.size());
1193 }
1194 
1195 // this will generate non-overlapping files since it keeps increasing key_idx
GenerateNewFile(int cf,Random * rnd,int * key_idx,bool nowait)1196 void DBTestBase::GenerateNewFile(int cf, Random* rnd, int* key_idx,
1197                                  bool nowait) {
1198   for (int i = 0; i < KNumKeysByGenerateNewFile; i++) {
1199     ASSERT_OK(Put(cf, Key(*key_idx), RandomString(rnd, (i == 99) ? 1 : 990)));
1200     (*key_idx)++;
1201   }
1202   if (!nowait) {
1203     dbfull()->TEST_WaitForFlushMemTable();
1204     dbfull()->TEST_WaitForCompact();
1205   }
1206 }
1207 
1208 // this will generate non-overlapping files since it keeps increasing key_idx
GenerateNewFile(Random * rnd,int * key_idx,bool nowait)1209 void DBTestBase::GenerateNewFile(Random* rnd, int* key_idx, bool nowait) {
1210   for (int i = 0; i < KNumKeysByGenerateNewFile; i++) {
1211     ASSERT_OK(Put(Key(*key_idx), RandomString(rnd, (i == 99) ? 1 : 990)));
1212     (*key_idx)++;
1213   }
1214   if (!nowait) {
1215     dbfull()->TEST_WaitForFlushMemTable();
1216     dbfull()->TEST_WaitForCompact();
1217   }
1218 }
1219 
1220 const int DBTestBase::kNumKeysByGenerateNewRandomFile = 51;
1221 
GenerateNewRandomFile(Random * rnd,bool nowait)1222 void DBTestBase::GenerateNewRandomFile(Random* rnd, bool nowait) {
1223   for (int i = 0; i < kNumKeysByGenerateNewRandomFile; i++) {
1224     ASSERT_OK(Put("key" + RandomString(rnd, 7), RandomString(rnd, 2000)));
1225   }
1226   ASSERT_OK(Put("key" + RandomString(rnd, 7), RandomString(rnd, 200)));
1227   if (!nowait) {
1228     dbfull()->TEST_WaitForFlushMemTable();
1229     dbfull()->TEST_WaitForCompact();
1230   }
1231 }
1232 
IterStatus(Iterator * iter)1233 std::string DBTestBase::IterStatus(Iterator* iter) {
1234   std::string result;
1235   if (iter->Valid()) {
1236     result = iter->key().ToString() + "->" + iter->value().ToString();
1237   } else {
1238     result = "(invalid)";
1239   }
1240   return result;
1241 }
1242 
OptionsForLogIterTest()1243 Options DBTestBase::OptionsForLogIterTest() {
1244   Options options = CurrentOptions();
1245   options.create_if_missing = true;
1246   options.WAL_ttl_seconds = 1000;
1247   return options;
1248 }
1249 
DummyString(size_t len,char c)1250 std::string DBTestBase::DummyString(size_t len, char c) {
1251   return std::string(len, c);
1252 }
1253 
VerifyIterLast(std::string expected_key,int cf)1254 void DBTestBase::VerifyIterLast(std::string expected_key, int cf) {
1255   Iterator* iter;
1256   ReadOptions ro;
1257   if (cf == 0) {
1258     iter = db_->NewIterator(ro);
1259   } else {
1260     iter = db_->NewIterator(ro, handles_[cf]);
1261   }
1262   iter->SeekToLast();
1263   ASSERT_EQ(IterStatus(iter), expected_key);
1264   delete iter;
1265 }
1266 
1267 // Used to test InplaceUpdate
1268 
1269 // If previous value is nullptr or delta is > than previous value,
1270 //   sets newValue with delta
1271 // If previous value is not empty,
1272 //   updates previous value with 'b' string of previous value size - 1.
updateInPlaceSmallerSize(char * prevValue,uint32_t * prevSize,Slice delta,std::string * newValue)1273 UpdateStatus DBTestBase::updateInPlaceSmallerSize(char* prevValue,
1274                                                   uint32_t* prevSize,
1275                                                   Slice delta,
1276                                                   std::string* newValue) {
1277   if (prevValue == nullptr) {
1278     *newValue = std::string(delta.size(), 'c');
1279     return UpdateStatus::UPDATED;
1280   } else {
1281     *prevSize = *prevSize - 1;
1282     std::string str_b = std::string(*prevSize, 'b');
1283     memcpy(prevValue, str_b.c_str(), str_b.size());
1284     return UpdateStatus::UPDATED_INPLACE;
1285   }
1286 }
1287 
updateInPlaceSmallerVarintSize(char * prevValue,uint32_t * prevSize,Slice delta,std::string * newValue)1288 UpdateStatus DBTestBase::updateInPlaceSmallerVarintSize(char* prevValue,
1289                                                         uint32_t* prevSize,
1290                                                         Slice delta,
1291                                                         std::string* newValue) {
1292   if (prevValue == nullptr) {
1293     *newValue = std::string(delta.size(), 'c');
1294     return UpdateStatus::UPDATED;
1295   } else {
1296     *prevSize = 1;
1297     std::string str_b = std::string(*prevSize, 'b');
1298     memcpy(prevValue, str_b.c_str(), str_b.size());
1299     return UpdateStatus::UPDATED_INPLACE;
1300   }
1301 }
1302 
updateInPlaceLargerSize(char *,uint32_t *,Slice delta,std::string * newValue)1303 UpdateStatus DBTestBase::updateInPlaceLargerSize(char* /*prevValue*/,
1304                                                  uint32_t* /*prevSize*/,
1305                                                  Slice delta,
1306                                                  std::string* newValue) {
1307   *newValue = std::string(delta.size(), 'c');
1308   return UpdateStatus::UPDATED;
1309 }
1310 
updateInPlaceNoAction(char *,uint32_t *,Slice,std::string *)1311 UpdateStatus DBTestBase::updateInPlaceNoAction(char* /*prevValue*/,
1312                                                uint32_t* /*prevSize*/,
1313                                                Slice /*delta*/,
1314                                                std::string* /*newValue*/) {
1315   return UpdateStatus::UPDATE_FAILED;
1316 }
1317 
1318 // Utility method to test InplaceUpdate
validateNumberOfEntries(int numValues,int cf)1319 void DBTestBase::validateNumberOfEntries(int numValues, int cf) {
1320   Arena arena;
1321   auto options = CurrentOptions();
1322   InternalKeyComparator icmp(options.comparator);
1323   ReadRangeDelAggregator range_del_agg(&icmp,
1324                                        kMaxSequenceNumber /* upper_bound */);
1325   // This should be defined after range_del_agg so that it destructs the
1326   // assigned iterator before it range_del_agg is already destructed.
1327   ScopedArenaIterator iter;
1328   if (cf != 0) {
1329     iter.set(dbfull()->NewInternalIterator(&arena, &range_del_agg,
1330                                            kMaxSequenceNumber, handles_[cf]));
1331   } else {
1332     iter.set(dbfull()->NewInternalIterator(&arena, &range_del_agg,
1333                                            kMaxSequenceNumber));
1334   }
1335   iter->SeekToFirst();
1336   ASSERT_EQ(iter->status().ok(), true);
1337   int seq = numValues;
1338   while (iter->Valid()) {
1339     ParsedInternalKey ikey;
1340     ikey.clear();
1341     ASSERT_EQ(ParseInternalKey(iter->key(), &ikey), true);
1342 
1343     // checks sequence number for updates
1344     ASSERT_EQ(ikey.sequence, (unsigned)seq--);
1345     iter->Next();
1346   }
1347   ASSERT_EQ(0, seq);
1348 }
1349 
CopyFile(const std::string & source,const std::string & destination,uint64_t size)1350 void DBTestBase::CopyFile(const std::string& source,
1351                           const std::string& destination, uint64_t size) {
1352   const EnvOptions soptions;
1353   std::unique_ptr<SequentialFile> srcfile;
1354   ASSERT_OK(env_->NewSequentialFile(source, &srcfile, soptions));
1355   std::unique_ptr<WritableFile> destfile;
1356   ASSERT_OK(env_->NewWritableFile(destination, &destfile, soptions));
1357 
1358   if (size == 0) {
1359     // default argument means copy everything
1360     ASSERT_OK(env_->GetFileSize(source, &size));
1361   }
1362 
1363   char buffer[4096];
1364   Slice slice;
1365   while (size > 0) {
1366     uint64_t one = std::min(uint64_t(sizeof(buffer)), size);
1367     ASSERT_OK(srcfile->Read(one, &slice, buffer));
1368     ASSERT_OK(destfile->Append(slice));
1369     size -= slice.size();
1370   }
1371   ASSERT_OK(destfile->Close());
1372 }
1373 
GetAllSSTFiles(uint64_t * total_size)1374 std::unordered_map<std::string, uint64_t> DBTestBase::GetAllSSTFiles(
1375     uint64_t* total_size) {
1376   std::unordered_map<std::string, uint64_t> res;
1377 
1378   if (total_size) {
1379     *total_size = 0;
1380   }
1381   std::vector<std::string> files;
1382   env_->GetChildren(dbname_, &files);
1383   for (auto& file_name : files) {
1384     uint64_t number;
1385     FileType type;
1386     std::string file_path = dbname_ + "/" + file_name;
1387     if (ParseFileName(file_name, &number, &type) && type == kTableFile) {
1388       uint64_t file_size = 0;
1389       env_->GetFileSize(file_path, &file_size);
1390       res[file_path] = file_size;
1391       if (total_size) {
1392         *total_size += file_size;
1393       }
1394     }
1395   }
1396   return res;
1397 }
1398 
ListTableFiles(Env * env,const std::string & path)1399 std::vector<std::uint64_t> DBTestBase::ListTableFiles(Env* env,
1400                                                       const std::string& path) {
1401   std::vector<std::string> files;
1402   std::vector<uint64_t> file_numbers;
1403   env->GetChildren(path, &files);
1404   uint64_t number;
1405   FileType type;
1406   for (size_t i = 0; i < files.size(); ++i) {
1407     if (ParseFileName(files[i], &number, &type)) {
1408       if (type == kTableFile) {
1409         file_numbers.push_back(number);
1410       }
1411     }
1412   }
1413   return file_numbers;
1414 }
1415 
VerifyDBFromMap(std::map<std::string,std::string> true_data,size_t * total_reads_res,bool tailing_iter,std::map<std::string,Status> status)1416 void DBTestBase::VerifyDBFromMap(std::map<std::string, std::string> true_data,
1417                                  size_t* total_reads_res, bool tailing_iter,
1418                                  std::map<std::string, Status> status) {
1419   size_t total_reads = 0;
1420 
1421   for (auto& kv : true_data) {
1422     Status s = status[kv.first];
1423     if (s.ok()) {
1424       ASSERT_EQ(Get(kv.first), kv.second);
1425     } else {
1426       std::string value;
1427       ASSERT_EQ(s, db_->Get(ReadOptions(), kv.first, &value));
1428     }
1429     total_reads++;
1430   }
1431 
1432   // Normal Iterator
1433   {
1434     int iter_cnt = 0;
1435     ReadOptions ro;
1436     ro.total_order_seek = true;
1437     Iterator* iter = db_->NewIterator(ro);
1438     // Verify Iterator::Next()
1439     iter_cnt = 0;
1440     auto data_iter = true_data.begin();
1441     Status s;
1442     for (iter->SeekToFirst(); iter->Valid(); iter->Next(), data_iter++) {
1443       ASSERT_EQ(iter->key().ToString(), data_iter->first);
1444       Status current_status = status[data_iter->first];
1445       if (!current_status.ok()) {
1446         s = current_status;
1447       }
1448       ASSERT_EQ(iter->status(), s);
1449       if (current_status.ok()) {
1450         ASSERT_EQ(iter->value().ToString(), data_iter->second);
1451       }
1452       iter_cnt++;
1453       total_reads++;
1454     }
1455     ASSERT_EQ(data_iter, true_data.end()) << iter_cnt << " / "
1456                                           << true_data.size();
1457     delete iter;
1458 
1459     // Verify Iterator::Prev()
1460     // Use a new iterator to make sure its status is clean.
1461     iter = db_->NewIterator(ro);
1462     iter_cnt = 0;
1463     s = Status::OK();
1464     auto data_rev = true_data.rbegin();
1465     for (iter->SeekToLast(); iter->Valid(); iter->Prev(), data_rev++) {
1466       ASSERT_EQ(iter->key().ToString(), data_rev->first);
1467       Status current_status = status[data_rev->first];
1468       if (!current_status.ok()) {
1469         s = current_status;
1470       }
1471       ASSERT_EQ(iter->status(), s);
1472       if (current_status.ok()) {
1473         ASSERT_EQ(iter->value().ToString(), data_rev->second);
1474       }
1475       iter_cnt++;
1476       total_reads++;
1477     }
1478     ASSERT_EQ(data_rev, true_data.rend()) << iter_cnt << " / "
1479                                           << true_data.size();
1480 
1481     // Verify Iterator::Seek()
1482     for (auto kv : true_data) {
1483       iter->Seek(kv.first);
1484       ASSERT_EQ(kv.first, iter->key().ToString());
1485       ASSERT_EQ(kv.second, iter->value().ToString());
1486       total_reads++;
1487     }
1488     delete iter;
1489   }
1490 
1491   if (tailing_iter) {
1492 #ifndef ROCKSDB_LITE
1493     // Tailing iterator
1494     int iter_cnt = 0;
1495     ReadOptions ro;
1496     ro.tailing = true;
1497     ro.total_order_seek = true;
1498     Iterator* iter = db_->NewIterator(ro);
1499 
1500     // Verify ForwardIterator::Next()
1501     iter_cnt = 0;
1502     auto data_iter = true_data.begin();
1503     for (iter->SeekToFirst(); iter->Valid(); iter->Next(), data_iter++) {
1504       ASSERT_EQ(iter->key().ToString(), data_iter->first);
1505       ASSERT_EQ(iter->value().ToString(), data_iter->second);
1506       iter_cnt++;
1507       total_reads++;
1508     }
1509     ASSERT_EQ(data_iter, true_data.end()) << iter_cnt << " / "
1510                                           << true_data.size();
1511 
1512     // Verify ForwardIterator::Seek()
1513     for (auto kv : true_data) {
1514       iter->Seek(kv.first);
1515       ASSERT_EQ(kv.first, iter->key().ToString());
1516       ASSERT_EQ(kv.second, iter->value().ToString());
1517       total_reads++;
1518     }
1519 
1520     delete iter;
1521 #endif  // ROCKSDB_LITE
1522   }
1523 
1524   if (total_reads_res) {
1525     *total_reads_res = total_reads;
1526   }
1527 }
1528 
VerifyDBInternal(std::vector<std::pair<std::string,std::string>> true_data)1529 void DBTestBase::VerifyDBInternal(
1530     std::vector<std::pair<std::string, std::string>> true_data) {
1531   Arena arena;
1532   InternalKeyComparator icmp(last_options_.comparator);
1533   ReadRangeDelAggregator range_del_agg(&icmp,
1534                                        kMaxSequenceNumber /* upper_bound */);
1535   auto iter =
1536       dbfull()->NewInternalIterator(&arena, &range_del_agg, kMaxSequenceNumber);
1537   iter->SeekToFirst();
1538   for (auto p : true_data) {
1539     ASSERT_TRUE(iter->Valid());
1540     ParsedInternalKey ikey;
1541     ASSERT_TRUE(ParseInternalKey(iter->key(), &ikey));
1542     ASSERT_EQ(p.first, ikey.user_key);
1543     ASSERT_EQ(p.second, iter->value());
1544     iter->Next();
1545   };
1546   ASSERT_FALSE(iter->Valid());
1547   iter->~InternalIterator();
1548 }
1549 
1550 #ifndef ROCKSDB_LITE
1551 
GetNumberOfSstFilesForColumnFamily(DB * db,std::string column_family_name)1552 uint64_t DBTestBase::GetNumberOfSstFilesForColumnFamily(
1553     DB* db, std::string column_family_name) {
1554   std::vector<LiveFileMetaData> metadata;
1555   db->GetLiveFilesMetaData(&metadata);
1556   uint64_t result = 0;
1557   for (auto& fileMetadata : metadata) {
1558     result += (fileMetadata.column_family_name == column_family_name);
1559   }
1560   return result;
1561 }
1562 #endif  // ROCKSDB_LITE
1563 
1564 }  // namespace ROCKSDB_NAMESPACE
1565