txn_bench.cc

#include <arpa/inet.h>
#include <pthread.h>
#include <stdlib.h>
#include <time.h>
#include <sys/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 = "/dev/shm/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]));
}

double gettimeofday_sec() {
  struct timeval tv;
  gettimeofday(&tv, NULL);
  return tv.tv_sec + (double)tv.tv_usec * 1e-6;
}

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;

  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;
  default_cf_opts.disable_auto_compactions = true;

  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());
  TransactionOptions txn_opts;
  WriteOptions write_opts;
  Transaction *txn = db->BeginTransaction(write_opts, txn_opts);

  srand(0);
  int num_ok = 0;
  auto t1 = gettimeofday_sec();
  if (!strcmp(mode, "put")) {
    txn->SetSnapshot();

    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);
      num_ok += s.ok();
      assert(s.ok());
    }
  } else if (!strcmp(mode, "get_for_update")) {
    rocksdb::ReadOptions read_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);
      num_ok += s.ok();
      assert(s.ok());
    }
  }
  auto t2 = gettimeofday_sec();

  printf("%d OK Ops | took %10.3f seconds\n", num_ok, t2 - t1);

  s = txn->Commit();
  if (!s.ok()) {
    printf("Error while commit\n");
    return 1;
  }

  delete txn;
  delete db;

  return 0;
}