ceph/src/rocksdb/util/transaction_test_util.h

   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 {
  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   explicit RandomTransactionInserter(
  39       Random64* rand, const WriteOptions& write_options = WriteOptions(),
  40       const ReadOptions& read_options = ReadOptions(), uint64_t num_keys = 1000,
  41       uint16_t num_sets = 3);
  42
  43   ~RandomTransactionInserter();
  44
  45   // Increment a key in each set using a Transaction on a TransactionDB.
  46   //
  47   // Returns true if the transaction succeeded OR if any error encountered was
  48   // expected (eg a write-conflict). Error status may be obtained by calling
  49   // GetLastStatus();
  50   bool TransactionDBInsert(
  51       TransactionDB* db,
  52       const TransactionOptions& txn_options = TransactionOptions());
  53
  54   // Increment a key in each set using a Transaction on an
  55   // OptimisticTransactionDB
  56   //
  57   // Returns true if the transaction succeeded OR if any error encountered was
  58   // expected (eg a write-conflict). Error status may be obtained by calling
  59   // GetLastStatus();
  60   bool OptimisticTransactionDBInsert(
  61       OptimisticTransactionDB* db,
  62       const OptimisticTransactionOptions& txn_options =
  63           OptimisticTransactionOptions());
  64   // Increment a key in each set without using a transaction.  If this function
  65   // is called in parallel, then Verify() may fail.
  66   //
  67   // Returns true if the write succeeds.
  68   // Error status may be obtained by calling GetLastStatus().
  69   bool DBInsert(DB* db);
  70
  71   // Get the ikey'th key from set set_i
  72   static Status DBGet(DB* db, Transaction* txn, ReadOptions& read_options,
  73                       uint16_t set_i, uint64_t ikey, bool get_for_update,
  74                       uint64_t* int_value, std::string* full_key,
  75                       bool* unexpected_error);
  76
  77   // Returns OK if Invariant is true.
  78   static Status Verify(DB* db, uint16_t num_sets, uint64_t num_keys_per_set = 0,
  79                        bool take_snapshot = false, Random64* rand = nullptr);
  80
  81   // Returns the status of the previous Insert operation
  82   Status GetLastStatus() { return last_status_; }
  83
  84   // Returns the number of successfully written calls to
  85   // TransactionDBInsert/OptimisticTransactionDBInsert/DBInsert
  86   uint64_t GetSuccessCount() { return success_count_; }
  87
  88   // Returns the number of calls to
  89   // TransactionDBInsert/OptimisticTransactionDBInsert/DBInsert that did not
  90   // write any data.
  91   uint64_t GetFailureCount() { return failure_count_; }
  92
  93   // Returns the sum of user keys/values Put() to the DB.
  94   size_t GetBytesInserted() { return bytes_inserted_; }
  95
  96  private:
  97   // Input options
  98   Random64* rand_;
  99   const WriteOptions write_options_;
 100   ReadOptions read_options_;
 101   const uint64_t num_keys_;
 102   const uint16_t num_sets_;
 103
 104   // Number of successful insert batches performed
 105   uint64_t success_count_ = 0;
 106
 107   // Number of failed insert batches attempted
 108   uint64_t failure_count_ = 0;
 109
 110   size_t bytes_inserted_ = 0;
 111
 112   // Status returned by most recent insert operation
 113   Status last_status_;
 114
 115   // optimization: re-use allocated transaction objects.
 116   Transaction* txn_ = nullptr;
 117   Transaction* optimistic_txn_ = nullptr;
 118
 119   std::atomic<int> txn_id_;
 120
 121   bool DoInsert(DB* db, Transaction* txn, bool is_optimistic);
 122 };
 123
 124 }  // namespace rocksdb
 125
 126 #endif  // ROCKSDB_LITE