ceph/src/rocksdb/tools/db_repl_stress.cc

   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 #ifndef ROCKSDB_LITE
   7 #ifndef GFLAGS
   8 #include <cstdio>
   9 int main() {
  10   fprintf(stderr, "Please install gflags to run rocksdb tools\n");
  11   return 1;
  12 }
  13 #else
  14
  15 #include <atomic>
  16 #include <cstdio>
  17
  18 #include "db/write_batch_internal.h"
  19 #include "rocksdb/db.h"
  20 #include "rocksdb/types.h"
  21 #include "util/gflags_compat.h"
  22 #include "util/testutil.h"
  23
  24 // Run a thread to perform Put's.
  25 // Another thread uses GetUpdatesSince API to keep getting the updates.
  26 // options :
  27 // --num_inserts = the num of inserts the first thread should perform.
  28 // --wal_ttl = the wal ttl for the run.
  29
  30 using namespace rocksdb;
  31
  32 using GFLAGS_NAMESPACE::ParseCommandLineFlags;
  33 using GFLAGS_NAMESPACE::SetUsageMessage;
  34
  35 struct DataPumpThread {
  36   size_t no_records;
  37   DB* db;  // Assumption DB is Open'ed already.
  38 };
  39
  40 static std::string RandomString(Random* rnd, int len) {
  41   std::string r;
  42   test::RandomString(rnd, len, &r);
  43   return r;
  44 }
  45
  46 static void DataPumpThreadBody(void* arg) {
  47   DataPumpThread* t = reinterpret_cast<DataPumpThread*>(arg);
  48   DB* db = t->db;
  49   Random rnd(301);
  50   size_t i = 0;
  51   while (i++ < t->no_records) {
  52     if (!db->Put(WriteOptions(), Slice(RandomString(&rnd, 500)),
  53                  Slice(RandomString(&rnd, 500)))
  54              .ok()) {
  55       fprintf(stderr, "Error in put\n");
  56       exit(1);
  57     }
  58   }
  59 }
  60
  61 struct ReplicationThread {
  62   std::atomic<bool> stop;
  63   DB* db;
  64   volatile size_t no_read;
  65 };
  66
  67 static void ReplicationThreadBody(void* arg) {
  68   ReplicationThread* t = reinterpret_cast<ReplicationThread*>(arg);
  69   DB* db = t->db;
  70   std::unique_ptr<TransactionLogIterator> iter;
  71   SequenceNumber currentSeqNum = 1;
  72   while (!t->stop.load(std::memory_order_acquire)) {
  73     iter.reset();
  74     Status s;
  75     while (!db->GetUpdatesSince(currentSeqNum, &iter).ok()) {
  76       if (t->stop.load(std::memory_order_acquire)) {
  77         return;
  78       }
  79     }
  80     fprintf(stderr, "Refreshing iterator\n");
  81     for (; iter->Valid(); iter->Next(), t->no_read++, currentSeqNum++) {
  82       BatchResult res = iter->GetBatch();
  83       if (res.sequence != currentSeqNum) {
  84         fprintf(stderr, "Missed a seq no. b/w %ld and %ld\n",
  85                 (long)currentSeqNum, (long)res.sequence);
  86         exit(1);
  87       }
  88     }
  89   }
  90 }
  91
  92 DEFINE_uint64(num_inserts, 1000,
  93               "the num of inserts the first thread should"
  94               " perform.");
  95 DEFINE_uint64(wal_ttl_seconds, 1000, "the wal ttl for the run(in seconds)");
  96 DEFINE_uint64(wal_size_limit_MB, 10,
  97               "the wal size limit for the run"
  98               "(in MB)");
  99
 100 int main(int argc, const char** argv) {
 101   SetUsageMessage(
 102       std::string("\nUSAGE:\n") + std::string(argv[0]) +
 103       " --num_inserts=<num_inserts> --wal_ttl_seconds=<WAL_ttl_seconds>" +
 104       " --wal_size_limit_MB=<WAL_size_limit_MB>");
 105   ParseCommandLineFlags(&argc, const_cast<char***>(&argv), true);
 106
 107   Env* env = Env::Default();
 108   std::string default_db_path;
 109   env->GetTestDirectory(&default_db_path);
 110   default_db_path += "db_repl_stress";
 111   Options options;
 112   options.create_if_missing = true;
 113   options.WAL_ttl_seconds = FLAGS_wal_ttl_seconds;
 114   options.WAL_size_limit_MB = FLAGS_wal_size_limit_MB;
 115   DB* db;
 116   DestroyDB(default_db_path, options);
 117
 118   Status s = DB::Open(options, default_db_path, &db);
 119
 120   if (!s.ok()) {
 121     fprintf(stderr, "Could not open DB due to %s\n", s.ToString().c_str());
 122     exit(1);
 123   }
 124
 125   DataPumpThread dataPump;
 126   dataPump.no_records = FLAGS_num_inserts;
 127   dataPump.db = db;
 128   env->StartThread(DataPumpThreadBody, &dataPump);
 129
 130   ReplicationThread replThread;
 131   replThread.db = db;
 132   replThread.no_read = 0;
 133   replThread.stop.store(false, std::memory_order_release);
 134
 135   env->StartThread(ReplicationThreadBody, &replThread);
 136   while (replThread.no_read < FLAGS_num_inserts)
 137     ;
 138   replThread.stop.store(true, std::memory_order_release);
 139   if (replThread.no_read < dataPump.no_records) {
 140     // no. read should be => than inserted.
 141     fprintf(stderr,
 142             "No. of Record's written and read not same\nRead : %" ROCKSDB_PRIszt
 143             " Written : %" ROCKSDB_PRIszt "\n",
 144             replThread.no_read, dataPump.no_records);
 145     exit(1);
 146   }
 147   fprintf(stderr, "Successful!\n");
 148   exit(0);
 149 }
 150
 151 #endif  // GFLAGS
 152
 153 #else  // ROCKSDB_LITE
 154 #include <stdio.h>
 155 int main(int /*argc*/, char** /*argv*/) {
 156   fprintf(stderr, "Not supported in lite mode.\n");
 157   return 1;
 158 }
 159 #endif  // ROCKSDB_LITE