Skip to content

Instantly share code, notes, and snippets.

@IslamAbdelRahman
Created March 16, 2016 01:54
Show Gist options
  • Save IslamAbdelRahman/b120f4fba8d6ff7d58d2 to your computer and use it in GitHub Desktop.
Save IslamAbdelRahman/b120f4fba8d6ff7d58d2 to your computer and use it in GitHub Desktop.
#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;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment