abby/kv_bench
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix test

Browse Source
This commit is contained in:
abbycin
2025-12-20 16:13:03 +08:00
parent 75e8b90cb9
commit ccb108c33d
12 changed files with 117 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rocksdb/main.cpp
+66 -8
View File
@@ -1,7 +1,8 @@
#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <exception>
#include <fmt/base.h>
#include <fmt/format.h>
#include <memory>
@@ -10,6 +11,7 @@
#include <rocksdb/db.h>
#include <rocksdb/env.h>
#include <rocksdb/options.h>
#include <rocksdb/slice.h>
#include <rocksdb/table.h>
#include <rocksdb/utilities/optimistic_transaction_db.h>
#include <rocksdb/utilities/transaction.h>
@@ -23,11 +25,17 @@
#include "CLI/CLI.hpp"
#include "instant.h"
template<class T>
static void black_box(const T &t) {
asm volatile("" ::"m"(t) : "memory");
}
struct Args {
size_t threads;
size_t iterations;
size_t key_size;
size_t value_size;
size_t blob_size;
size_t insert_ratio;
bool random;
std::string mode;
@@ -41,6 +49,7 @@ int main(int argc, char *argv[]) {
.iterations = 100000,
.key_size = 16,
.value_size = 1024,
.blob_size = 8192,
.insert_ratio = 30,
.mode = "insert",
.path = "/tmp/rocksdb_tmp",
@@ -50,6 +59,7 @@ int main(int argc, char *argv[]) {
app.add_option("-t,--threads", args.threads, "Threads");
app.add_option("-k,--key-size", args.key_size, "Key Size");
app.add_option("-v,--value-size", args.value_size, "Value Size");
app.add_option("-b,--blob-size", args.value_size, "Blob Size");
app.add_option("-i,--iterations", args.iterations, "Iterations");
app.add_option("-r,--insert-ratio", args.insert_ratio, "Insert Ratio for mixed mode");
app.add_option("-p,--path", args.path, "DataBase Home");
@@ -82,9 +92,25 @@ int main(int argc, char *argv[]) {
return 1;
}
auto find_upper_bound = [](std::string prefix) {
std::string upper_bound_key = prefix;
for (int i = upper_bound_key.length() - 1; i >= 0; --i) {
if ((unsigned char) upper_bound_key[i] != 0xff) {
upper_bound_key[i] = (unsigned char) upper_bound_key[i] + 1;
upper_bound_key.resize(i + 1);
break;
}
if (i == 0) {
upper_bound_key = "";
break;
}
}
return upper_bound_key;
};
rocksdb::ColumnFamilyOptions cfo{};
cfo.enable_blob_files = true;
cfo.min_blob_size = 8192;
cfo.min_blob_size = args.blob_size;
// use 1GB block cache
auto cache = rocksdb::NewLRUCache(1 << 30);
rocksdb::BlockBasedTableOptions table_options{};
@@ -122,11 +148,11 @@ int main(int argc, char *argv[]) {
std::mt19937 gen(rd());
std::uniform_int_distribution<int> dist(0, 100);
std::string val(args.value_size, 'x');
auto keys_per_thread = args.iterations / args.threads;
for (size_t tid = 0; tid < args.threads; ++tid) {
std::vector<std::string> key{};
for (size_t i = 0; i < args.iterations; ++i) {
for (size_t i = 0; i < keys_per_thread; ++i) {
auto tmp = std::format("key_{}_{}", tid, i);
tmp.resize(args.key_size, 'x');
key.emplace_back(std::move(tmp));
@@ -137,7 +163,6 @@ int main(int argc, char *argv[]) {
keys.emplace_back(std::move(key));
}
auto *handle = handles[0];
if (args.mode == "get" || args.mode == "scan") {
@@ -156,15 +181,37 @@ int main(int argc, char *argv[]) {
// re-open db
s = rocksdb::OptimisticTransactionDB::Open(options, args.path, cfd, &handles, &db);
assert(s.ok());
handle = handles[0];
// simulate common use cases
std::uniform_int_distribution<int> dist(0, args.threads - 1);
for (size_t i = 0; i < keys_per_thread; ++i) {
auto tid = dist(gen);
auto k = std::format("key_{}_{}", tid, i);
k.resize(args.key_size, 'x');
auto s = db->Get(rocksdb::ReadOptions(), k, &val);
if (!s.ok()) {
std::terminate();
}
}
}
handle = handles[0];
auto *snapshot = db->GetSnapshot();
for (size_t tid = 0; tid < args.threads; ++tid) {
auto *tk = &keys[tid];
wg.emplace_back([&, tid] {
std::string rval(args.value_size, '0');
auto prefix = std::format("key_{}", tid);
auto ropt = rocksdb::ReadOptions();
auto upper_bound = find_upper_bound(prefix);
auto upper_bound_slice = rocksdb::Slice(upper_bound);
if (!upper_bound.empty()) {
ropt.iterate_upper_bound = &upper_bound_slice;
}
auto *tk = &keys[tid];
ropt.prefix_same_as_start = true;
ropt.snapshot = snapshot;
size_t round = 0;
barrier.arrive_and_wait();
if (mtx.try_lock()) {
@@ -174,6 +221,7 @@ int main(int argc, char *argv[]) {
if (args.mode == "insert") {
for (auto &key: *tk) {
round += 1;
auto *kv = db->BeginTransaction(wopt);
kv->Put(handle, key, val);
kv->Commit();
@@ -182,6 +230,7 @@ int main(int argc, char *argv[]) {
} else if (args.mode == "get") {
for (auto &key: *tk) {
round += 1;
auto *kv = db->BeginTransaction(wopt);
kv->Get(ropt, handle, key, &rval);
kv->Commit();
@@ -189,6 +238,7 @@ int main(int argc, char *argv[]) {
}
} else if (args.mode == "mixed") {
for (auto &key: *tk) {
round += 1;
auto is_insert = dist(gen) < args.insert_ratio;
auto *kv = db->BeginTransaction(wopt);
if (is_insert) {
@@ -202,12 +252,19 @@ int main(int argc, char *argv[]) {
} else if (args.mode == "scan") {
auto *iter = db->NewIterator(ropt);
iter->Seek(prefix);
size_t n = 0;
while (iter->Valid()) {
round += 1;
auto k = iter->key();
auto v = iter->value();
black_box(k);
black_box(v);
iter->Next();
n += 1;
}
delete iter;
}
total_op.fetch_add(args.iterations, std::memory_order::relaxed);
total_op.fetch_add(round, std::memory_order::relaxed);
});
}
@@ -222,6 +279,7 @@ int main(int argc, char *argv[]) {
uint64_t ops = total_op.load(std::memory_order_relaxed) / b.elapse_sec();
fmt::println("{},{},{},{},{},{},{}", args.mode, args.threads, args.key_size, args.value_size, ratio, (uint64_t) ops,
(uint64_t) b.elapse_ms());
db->ReleaseSnapshot(snapshot);
delete handle;
delete db;
std::filesystem::remove_all(args.path);