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