bench_wbwi.cc

#include <arpa/inet.h>
#include <pthread.h>
#include <stdlib.h>
#include <time.h>
#include <cstdio>
#include <string>

#include "rocksdb/cache.h"
#include "rocksdb/db.h"
#include "rocksdb/options.h"
#include "rocksdb/perf_context.h"
#include "rocksdb/slice.h"
#include "rocksdb/statistics.h"
#include "rocksdb/utilities/checkpoint.h"
#include "rocksdb/utilities/convenience.h"
#include "rocksdb/utilities/transaction_db.h"
#include "rocksdb/utilities/write_batch_with_index.h"

using namespace rocksdb;

std::string kDBPath = "/tmp/rocksdb-example";
rocksdb::TransactionDB *db;

typedef unsigned char uchar;

inline void store_big_uint4(uchar *dst, uint32_t n) {
  uint32_t src = htonl(n);
  memcpy(dst, &src, 4);
}

inline uint32_t read_big_uint4(const uchar *b) {
  return (((uint32_t)(b[0]) << 24) | ((uint32_t)(b[1]) << 16) |
          ((uint32_t)(b[2]) << 8) | (uint32_t)(b[3]));
}

static rocksdb::BlockBasedTableOptions table_options;
rocksdb::ColumnFamilyOptions default_cf_opts;

int main(int argc, char **argv) {
  if (argc < 3) {
    printf("Arguments: insert|update num_trans\n");
    exit(1);
  }
  char *mode = argv[1];
  int num_trans = atoi(argv[2]);
  srand(0);

  rocksdb::DBOptions db_options;
  db_options.create_if_missing = true;
  //db_options.statistics = CreateDBStatistics();

  rocksdb::TransactionDBOptions txn_db_options;

  table_options.block_cache = rocksdb::NewLRUCache(10 * 1024 * 1024);
  table_options.format_version = 2;
  table_options.block_size = 16384;
  table_options.cache_index_and_filter_blocks = 1;

  default_cf_opts.table_factory.reset(
      rocksdb::NewBlockBasedTableFactory(table_options));
  default_cf_opts.write_buffer_size = 128 * 1024 * 1024;

  rocksdb::Options main_opts(db_options, default_cf_opts);

  // open DB
  Status s = rocksdb::TransactionDB::Open(main_opts, txn_db_options, kDBPath, &db);
  assert(s.ok());

  if (!strcmp(mode, "insert")) {
    TransactionOptions txn_opts;
    WriteOptions write_opts;
    Transaction *txn = db->BeginTransaction(write_opts, txn_opts);
    txn->SetSnapshot();
    // rocksdb::WriteBatchBase *wb = txn->GetWriteBatch();
    for (int i = 1; i <= num_trans; i++) {
      uchar key[128] = {0};
      uchar value[128] = {0};
      uint id2_1 = 0;
      uint id2_2 = i / 5 + 1;
      uint id_type_1 = 0;
      uint id_type_2 = 123456789;
      int data1 = rand();
      int data2 = rand();
      store_big_uint4(key, 260);
      store_big_uint4(key + 4, i);

      store_big_uint4(value, id2_1);
      store_big_uint4(value + 4, id2_2);
      store_big_uint4(value + 8, id_type_1);
      store_big_uint4(value + 12, id_type_2);
      store_big_uint4(value + 16, data1);
      store_big_uint4(value + 20, data2);

      Slice key_slice = Slice((char *)key, 8);
      Slice value_slice = Slice((char *)value, 24);

      s = txn->Put(key_slice, value_slice);
      assert(s.ok());
    }
  } else if (!strcmp(mode, "update")) {
    rocksdb::ReadOptions read_opts;
    TransactionOptions txn_opts;
    WriteOptions write_opts;
    Transaction *txn = db->BeginTransaction(write_opts, txn_opts);
    txn->SetSnapshot();
    read_opts.snapshot = txn->GetSnapshot();
    std::string retrieved_record;
    for (int i = 1; i <= num_trans; i++) {
      uchar key[128] = {0};
      uchar value[128] = {0};
      int data1 = rand();
      int data2 = rand();
      store_big_uint4(key, 260);
      store_big_uint4(key + 4, i);

      store_big_uint4(value, 1);
      store_big_uint4(value + 4, data1);
      store_big_uint4(value + 8, data2);

      Slice key_slice = Slice((char *)key, 8);
      Slice value_slice = Slice((char *)value, 12);

      s = txn->GetForUpdate(read_opts, key_slice, &retrieved_record);
      assert(s.ok());

      s = txn->Put(key_slice, value_slice);
      assert(s.ok());
    }
  }

  delete db;

  return 0;
}