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 #pragma once 7 8 #ifndef ROCKSDB_LITE 9 10 #include "rocksdb/options.h" 11 #include "port/port.h" 12 #include "rocksdb/utilities/optimistic_transaction_db.h" 13 #include "rocksdb/utilities/transaction_db.h" 14 15 namespace ROCKSDB_NAMESPACE { 16 17 class DB; 18 class Random64; 19 20 // Utility class for stress testing transactions. Can be used to write many 21 // transactions in parallel and then validate that the data written is logically 22 // consistent. This class assumes the input DB is initially empty. 23 // 24 // Each call to TransactionDBInsert()/OptimisticTransactionDBInsert() will 25 // increment the value of a key in #num_sets sets of keys. Regardless of 26 // whether the transaction succeeds, the total sum of values of keys in each 27 // set is an invariant that should remain equal. 28 // 29 // After calling TransactionDBInsert()/OptimisticTransactionDBInsert() many 30 // times, Verify() can be called to validate that the invariant holds. 31 // 32 // To test writing Transaction in parallel, multiple threads can create a 33 // RandomTransactionInserter with similar arguments using the same DB. 34 class RandomTransactionInserter { 35 public: 36 // num_keys is the number of keys in each set. 37 // num_sets is the number of sets of keys. 38 // cmt_delay_ms is the delay between prepare (if there is any) and commit 39 // first_id is the id of the first transaction 40 explicit RandomTransactionInserter( 41 Random64* rand, const WriteOptions& write_options = WriteOptions(), 42 const ReadOptions& read_options = ReadOptions(), uint64_t num_keys = 1000, 43 uint16_t num_sets = 3, const uint64_t cmt_delay_ms = 0, 44 const uint64_t first_id = 0); 45 46 ~RandomTransactionInserter(); 47 48 // Increment a key in each set using a Transaction on a TransactionDB. 49 // 50 // Returns true if the transaction succeeded OR if any error encountered was 51 // expected (eg a write-conflict). Error status may be obtained by calling 52 // GetLastStatus(); 53 bool TransactionDBInsert( 54 TransactionDB* db, 55 const TransactionOptions& txn_options = TransactionOptions()); 56 57 // Increment a key in each set using a Transaction on an 58 // OptimisticTransactionDB 59 // 60 // Returns true if the transaction succeeded OR if any error encountered was 61 // expected (eg a write-conflict). Error status may be obtained by calling 62 // GetLastStatus(); 63 bool OptimisticTransactionDBInsert( 64 OptimisticTransactionDB* db, 65 const OptimisticTransactionOptions& txn_options = 66 OptimisticTransactionOptions()); 67 // Increment a key in each set without using a transaction. If this function 68 // is called in parallel, then Verify() may fail. 69 // 70 // Returns true if the write succeeds. 71 // Error status may be obtained by calling GetLastStatus(). 72 bool DBInsert(DB* db); 73 74 // Get the ikey'th key from set set_i 75 static Status DBGet(DB* db, Transaction* txn, ReadOptions& read_options, 76 uint16_t set_i, uint64_t ikey, bool get_for_update, 77 uint64_t* int_value, std::string* full_key, 78 bool* unexpected_error); 79 80 // Returns OK if Invariant is true. 81 static Status Verify(DB* db, uint16_t num_sets, uint64_t num_keys_per_set = 0, 82 bool take_snapshot = false, Random64* rand = nullptr, 83 uint64_t delay_ms = 0); 84 85 // Returns the status of the previous Insert operation GetLastStatus()86 Status GetLastStatus() { return last_status_; } 87 88 // Returns the number of successfully written calls to 89 // TransactionDBInsert/OptimisticTransactionDBInsert/DBInsert GetSuccessCount()90 uint64_t GetSuccessCount() { return success_count_; } 91 92 // Returns the number of calls to 93 // TransactionDBInsert/OptimisticTransactionDBInsert/DBInsert that did not 94 // write any data. GetFailureCount()95 uint64_t GetFailureCount() { return failure_count_; } 96 97 // Returns the sum of user keys/values Put() to the DB. GetBytesInserted()98 size_t GetBytesInserted() { return bytes_inserted_; } 99 100 private: 101 // Input options 102 Random64* rand_; 103 const WriteOptions write_options_; 104 ReadOptions read_options_; 105 const uint64_t num_keys_; 106 const uint16_t num_sets_; 107 108 // Number of successful insert batches performed 109 uint64_t success_count_ = 0; 110 111 // Number of failed insert batches attempted 112 uint64_t failure_count_ = 0; 113 114 size_t bytes_inserted_ = 0; 115 116 // Status returned by most recent insert operation 117 Status last_status_; 118 119 // optimization: re-use allocated transaction objects. 120 Transaction* txn_ = nullptr; 121 Transaction* optimistic_txn_ = nullptr; 122 123 uint64_t txn_id_; 124 // The delay between ::Prepare and ::Commit 125 const uint64_t cmt_delay_ms_; 126 127 bool DoInsert(DB* db, Transaction* txn, bool is_optimistic); 128 }; 129 130 } // namespace ROCKSDB_NAMESPACE 131 132 #endif // ROCKSDB_LITE 133