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 #ifdef GFLAGS
11 #pragma once
12 #include "db_stress_tool/db_stress_common.h"
13 #include "db_stress_tool/db_stress_shared_state.h"
14 
15 namespace ROCKSDB_NAMESPACE {
16 class Transaction;
17 class TransactionDB;
18 
19 class StressTest {
20  public:
21   StressTest();
22 
23   virtual ~StressTest();
24 
25   std::shared_ptr<Cache> NewCache(size_t capacity);
26 
27   bool BuildOptionsTable();
28 
29   void InitDb();
30   void InitReadonlyDb(SharedState*);
31 
32   // Return false if verification fails.
33   bool VerifySecondaries();
34 
35   void OperateDb(ThreadState* thread);
36   virtual void VerifyDb(ThreadState* thread) const = 0;
ContinuouslyVerifyDb(ThreadState *)37   virtual void ContinuouslyVerifyDb(ThreadState* /*thread*/) const {}
38 
39   void PrintStatistics();
40 
41  protected:
42   Status AssertSame(DB* db, ColumnFamilyHandle* cf,
43                     ThreadState::SnapshotState& snap_state);
44 
45   // Currently PreloadDb has to be single-threaded.
46   void PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
47                                     SharedState* shared);
48 
49   Status SetOptions(ThreadState* thread);
50 
51 #ifndef ROCKSDB_LITE
52   Status NewTxn(WriteOptions& write_opts, Transaction** txn);
53 
54   Status CommitTxn(Transaction* txn);
55 
56   Status RollbackTxn(Transaction* txn);
57 #endif
58 
MaybeClearOneColumnFamily(ThreadState *)59   virtual void MaybeClearOneColumnFamily(ThreadState* /* thread */) {}
60 
ShouldAcquireMutexOnKey()61   virtual bool ShouldAcquireMutexOnKey() const { return false; }
62 
GenerateColumnFamilies(const int,int rand_column_family)63   virtual std::vector<int> GenerateColumnFamilies(
64       const int /* num_column_families */, int rand_column_family) const {
65     return {rand_column_family};
66   }
67 
GenerateKeys(int64_t rand_key)68   virtual std::vector<int64_t> GenerateKeys(int64_t rand_key) const {
69     return {rand_key};
70   }
71 
72   virtual Status TestGet(ThreadState* thread, const ReadOptions& read_opts,
73                          const std::vector<int>& rand_column_families,
74                          const std::vector<int64_t>& rand_keys) = 0;
75 
76   virtual std::vector<Status> TestMultiGet(
77       ThreadState* thread, const ReadOptions& read_opts,
78       const std::vector<int>& rand_column_families,
79       const std::vector<int64_t>& rand_keys) = 0;
80 
81   virtual Status TestPrefixScan(ThreadState* thread,
82                                 const ReadOptions& read_opts,
83                                 const std::vector<int>& rand_column_families,
84                                 const std::vector<int64_t>& rand_keys) = 0;
85 
86   virtual Status TestPut(ThreadState* thread, WriteOptions& write_opts,
87                          const ReadOptions& read_opts,
88                          const std::vector<int>& cf_ids,
89                          const std::vector<int64_t>& keys, char (&value)[100],
90                          std::unique_ptr<MutexLock>& lock) = 0;
91 
92   virtual Status TestDelete(ThreadState* thread, WriteOptions& write_opts,
93                             const std::vector<int>& rand_column_families,
94                             const std::vector<int64_t>& rand_keys,
95                             std::unique_ptr<MutexLock>& lock) = 0;
96 
97   virtual Status TestDeleteRange(ThreadState* thread, WriteOptions& write_opts,
98                                  const std::vector<int>& rand_column_families,
99                                  const std::vector<int64_t>& rand_keys,
100                                  std::unique_ptr<MutexLock>& lock) = 0;
101 
102   virtual void TestIngestExternalFile(
103       ThreadState* thread, const std::vector<int>& rand_column_families,
104       const std::vector<int64_t>& rand_keys,
105       std::unique_ptr<MutexLock>& lock) = 0;
106 
107   // Issue compact range, starting with start_key, whose integer value
108   // is rand_key.
109   virtual void TestCompactRange(ThreadState* thread, int64_t rand_key,
110                                 const Slice& start_key,
111                                 ColumnFamilyHandle* column_family);
112 
113   // Calculate a hash value for all keys in range [start_key, end_key]
114   // at a certain snapshot.
115   uint32_t GetRangeHash(ThreadState* thread, const Snapshot* snapshot,
116                         ColumnFamilyHandle* column_family,
117                         const Slice& start_key, const Slice& end_key);
118 
119   // Return a column family handle that mirrors what is pointed by
120   // `column_family_id`, which will be used to validate data to be correct.
121   // By default, the column family itself will be returned.
GetControlCfh(ThreadState *,int column_family_id)122   virtual ColumnFamilyHandle* GetControlCfh(ThreadState* /* thread*/,
123                                             int column_family_id) {
124     return column_families_[column_family_id];
125   }
126 
127 #ifndef ROCKSDB_LITE
128   // Generated a list of keys that close to boundaries of SST keys.
129   // If there isn't any SST file in the DB, return empty list.
130   std::vector<std::string> GetWhiteBoxKeys(ThreadState* thread, DB* db,
131                                            ColumnFamilyHandle* cfh,
132                                            size_t num_keys);
133 #else   // !ROCKSDB_LITE
GetWhiteBoxKeys(ThreadState *,DB *,ColumnFamilyHandle *,size_t)134   std::vector<std::string> GetWhiteBoxKeys(ThreadState*, DB*,
135                                            ColumnFamilyHandle*, size_t) {
136     // Not supported in LITE mode.
137     return {};
138   }
139 #endif  // !ROCKSDB_LITE
140 
141   // Given a key K, this creates an iterator which scans to K and then
142   // does a random sequence of Next/Prev operations.
143   virtual Status TestIterate(ThreadState* thread, const ReadOptions& read_opts,
144                              const std::vector<int>& rand_column_families,
145                              const std::vector<int64_t>& rand_keys);
146 
147   // Enum used by VerifyIterator() to identify the mode to validate.
148   enum LastIterateOp {
149     kLastOpSeek,
150     kLastOpSeekForPrev,
151     kLastOpNextOrPrev,
152     kLastOpSeekToFirst,
153     kLastOpSeekToLast
154   };
155 
156   // Compare the two iterator, iter and cmp_iter are in the same position,
157   // unless iter might be made invalidate or undefined because of
158   // upper or lower bounds, or prefix extractor.
159   // Will flag failure if the verification fails.
160   // diverged = true if the two iterator is already diverged.
161   // True if verification passed, false if not.
162   // op_logs is the information to print when validation fails.
163   void VerifyIterator(ThreadState* thread, ColumnFamilyHandle* cmp_cfh,
164                       const ReadOptions& ro, Iterator* iter, Iterator* cmp_iter,
165                       LastIterateOp op, const Slice& seek_key,
166                       const std::string& op_logs, bool* diverged);
167 
168   virtual Status TestBackupRestore(ThreadState* thread,
169                                    const std::vector<int>& rand_column_families,
170                                    const std::vector<int64_t>& rand_keys);
171 
172   virtual Status TestCheckpoint(ThreadState* thread,
173                                 const std::vector<int>& rand_column_families,
174                                 const std::vector<int64_t>& rand_keys);
175 
176   void TestCompactFiles(ThreadState* thread, ColumnFamilyHandle* column_family);
177 
178   Status TestFlush(const std::vector<int>& rand_column_families);
179 
180   Status TestPauseBackground(ThreadState* thread);
181 
182   void TestAcquireSnapshot(ThreadState* thread, int rand_column_family,
183                            const std::string& keystr, uint64_t i);
184 
185   Status MaybeReleaseSnapshots(ThreadState* thread, uint64_t i);
186 #ifndef ROCKSDB_LITE
187   Status VerifyGetLiveFiles() const;
188   Status VerifyGetSortedWalFiles() const;
189   Status VerifyGetCurrentWalFile() const;
190 
191   virtual Status TestApproximateSize(
192       ThreadState* thread, uint64_t iteration,
193       const std::vector<int>& rand_column_families,
194       const std::vector<int64_t>& rand_keys);
195 #endif  // !ROCKSDB_LITE
196 
197   void VerificationAbort(SharedState* shared, std::string msg, Status s) const;
198 
199   void VerificationAbort(SharedState* shared, std::string msg, int cf,
200                          int64_t key) const;
201 
202   void PrintEnv() const;
203 
204   void Open();
205 
206   void Reopen(ThreadState* thread);
207 
208   std::shared_ptr<Cache> cache_;
209   std::shared_ptr<Cache> compressed_cache_;
210   std::shared_ptr<const FilterPolicy> filter_policy_;
211   DB* db_;
212 #ifndef ROCKSDB_LITE
213   TransactionDB* txn_db_;
214 #endif
215   Options options_;
216   std::vector<ColumnFamilyHandle*> column_families_;
217   std::vector<std::string> column_family_names_;
218   std::atomic<int> new_column_family_name_;
219   int num_times_reopened_;
220   std::unordered_map<std::string, std::vector<std::string>> options_table_;
221   std::vector<std::string> options_index_;
222   std::atomic<bool> db_preload_finished_;
223 
224   // Fields used for stress-testing secondary instance in the same process
225   std::vector<DB*> secondaries_;
226   std::vector<std::vector<ColumnFamilyHandle*>> secondary_cfh_lists_;
227 
228   // Fields used for continuous verification from another thread
229   DB* cmp_db_;
230   std::vector<ColumnFamilyHandle*> cmp_cfhs_;
231 };
232 
233 }  // namespace ROCKSDB_NAMESPACE
234 #endif  // GFLAGS
235