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

refine test

Browse Source
This commit is contained in:
abbycin 2026-03-03 21:32:10 +08:00
parent 80b4745118
commit abf82f735c
Signed by: abby
GPG Key ID: B636E0F0307EF8EB
5 changed files with 60 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
rocksdb/main.cpp
View File

@ -17,9 +17,11 @@
#include <rocksdb/utilities/transaction.h>
#include <rocksdb/utilities/transaction_db.h>
#include <algorithm>
#include <barrier>
#include <filesystem>
#include <format>
#include <numeric>
#include <string>
#include <pthread.h>
@ -47,6 +49,13 @@ static void bind_core(size_t tid) {
(void) pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &set);
}
static void require_ok(const rocksdb::Status &st, const char *what) {
if (!st.ok()) {
fmt::println(stderr, "{} failed: {}", what, st.ToString());
std::abort();
}
}
struct Args {
size_t threads;
size_t iterations;
@ -159,19 +168,17 @@ int main(int argc, char *argv[]) {
auto wopt = rocksdb::WriteOptions();
wopt.no_slowdown = true;
// wopt.disableWAL = true;
std::vector<std::thread> wg;
std::atomic<uint64_t> total_op{0};
rocksdb::OptimisticTransactionDB *db;
auto b = nm::Instant::now();
std::vector<rocksdb::ColumnFamilyHandle *> handles{};
auto s = rocksdb::OptimisticTransactionDB::Open(options, args.path, cfd, &handles, &db);
assert(s.ok());
require_ok(s, "open db");
std::barrier ready_barrier{static_cast<ptrdiff_t>(args.threads + 1)};
std::barrier start_barrier{static_cast<ptrdiff_t>(args.threads + 1)};
std::random_device rd{};
std::mt19937 gen(rd());
std::string val(args.value_size, 'x');
@ -189,9 +196,9 @@ int main(int argc, char *argv[]) {
for (size_t j = 0; j < batch_size && (i + j) < count; ++j) {
auto key = std::format("key_{}_{}", tid, i + j);
key.resize(args.key_size, 'x');
kv->Put(handle, key, val);
require_ok(kv->Put(handle, key, val), "fill put");
}
kv->Commit();
require_ok(kv->Commit(), "fill commit");
delete kv;
}
});
@ -204,7 +211,7 @@ int main(int argc, char *argv[]) {
handles.clear();
// re-open db
s = rocksdb::OptimisticTransactionDB::Open(options, args.path, cfd, &handles, &db);
assert(s.ok());
require_ok(s, "reopen db");
handle = handles[0];
}
@ -241,19 +248,18 @@ int main(int argc, char *argv[]) {
key.resize(args.key_size, 'x');
round += 1;
auto *kv = db->BeginTransaction(wopt);
kv->Put(handle, key, val);
kv->Commit();
require_ok(kv->Put(handle, key, val), "insert put");
require_ok(kv->Commit(), "insert commit");
delete kv;
}
} else if (args.mode == "get") {
// rocksdb has no dedicated read-only txn in this bench path, use direct get for fair read-path
// comparison with mace view
for (size_t i: indices) {
auto key = std::format("key_{}_{}", tid, i);
key.resize(args.key_size, 'x');
round += 1;
auto *kv = db->BeginTransaction(wopt);
kv->Get(ropt, handle, key, &rval);
kv->Commit();
delete kv;
require_ok(db->Get(ropt, handle, key, &rval), "get");
}
} else if (args.mode == "mixed") {
for (size_t i: indices) {
@ -263,11 +269,14 @@ int main(int argc, char *argv[]) {
auto is_insert = mixed_dist(thread_gen) < static_cast<int>(args.insert_ratio);
auto *kv = db->BeginTransaction(wopt);
if (is_insert) {
kv->Put(handle, key, val);
require_ok(kv->Put(handle, key, val), "mixed put");
} else {
kv->Get(ropt, handle, key, &rval);
auto st = kv->Get(ropt, handle, key, &rval);
if (!st.ok() && !st.IsNotFound()) {
require_ok(st, "mixed get");
}
kv->Commit();
}
require_ok(kv->Commit(), "mixed commit");
delete kv;
}
} else if (args.mode == "scan") {
@ -281,6 +290,7 @@ int main(int argc, char *argv[]) {
black_box(v);
iter->Next();
}
require_ok(iter->status(), "scan iterate");
delete iter;
}
total_op.fetch_add(round, std::memory_order::relaxed);
@ -299,7 +309,13 @@ int main(int argc, char *argv[]) {
return args.insert_ratio;
return args.mode == "insert" ? 100 : 0;
}();
uint64_t ops = total_op.load(std::memory_order_relaxed) / b.elapse_sec();
const auto elapsed_us = b.elapse_usec();
uint64_t ops = 0;
const auto total = total_op.load(std::memory_order_relaxed);
if (elapsed_us > 0) {
ops = static_cast<uint64_t>(static_cast<double>(total) * 1000000.0 / elapsed_us);
}
if (args.mode == "insert") {
if (args.random) {
args.mode = "random_insert";
@ -307,8 +323,8 @@ int main(int argc, char *argv[]) {
args.mode = "sequential_insert";
}
}
fmt::println("{},{},{},{},{},{},{}", args.mode, args.threads, args.key_size, args.value_size, ratio, (uint64_t) ops,
(uint64_t) b.elapse_ms());
fmt::println("{},{},{},{},{},{},{}", args.mode, args.threads, args.key_size, args.value_size, ratio, ops,
static_cast<uint64_t>(elapsed_us));
delete handle;
delete db;
std::filesystem::remove_all(args.path);

2
scripts/mace.sh
View File

@ -47,7 +47,7 @@ function samples() {
done
}
echo mode,threads,key_size,value_size,insert_ratio,ops,elasped > "${script_dir}/mace.csv"
echo mode,threads,key_size,value_size,insert_ratio,ops,elapsed_us > "${script_dir}/mace.csv"
samples "$1" 1>> "${script_dir}/mace.csv"
if [ -x "${script_dir}/bin/python" ]; then
(cd "${script_dir}" && "${script_dir}/bin/python" plot.py mace.csv)

26
scripts/plot.py
View File

@ -3,6 +3,7 @@ import matplotlib.pyplot as plt
from adjustText import adjust_text
import sys
def real_mode(m):
if m == "mixed":
return "Mixed (70% Get, 30% Insert)"
@ -12,20 +13,29 @@ def real_mode(m):
return "Sequential Scan"
return m.capitalize()
name = sys.argv[1]
prefix = name.split(".")[0]
# 读取数据
# read benchmark data
# keep compatibility with older csv files that used elapsed/elasped
# and normalize to elapsed_us
df = pd.read_csv(f"./{name}")
if "elapsed_us" not in df.columns:
if "elapsed" in df.columns:
df = df.rename(columns={"elapsed": "elapsed_us"})
elif "elasped" in df.columns:
df = df.rename(columns={"elasped": "elapsed_us"})
# 按 mode 分组
# group by mode
modes = df["mode"].unique()
for mode in modes:
plt.figure(figsize=(16, 9))
subset = df[df["mode"] == mode]
# 按 key_size/value_size 分组
# group by key/value size
key_value_combinations = subset.groupby(["key_size", "value_size"])
texts = []
@ -34,19 +44,17 @@ for mode in modes:
x = group["threads"]
y = group["ops"]
# 绘制折线
# draw line
line, = plt.plot(x, y, marker="o", label=label)
# 添加文本标签
# add labels
for xi, yi, ops in zip(x, y, group["ops"]):
texts.append(
plt.text(xi, yi, f"{int(ops)}", color=line.get_color(), fontsize=12)
)
# 自动调整文本位置
adjust_text(texts, arrowprops=dict(arrowstyle="->", color='gray'))
adjust_text(texts, arrowprops=dict(arrowstyle="->", color="gray"))
# 设置图表样式
plt.title(f"{prefix.upper()}: {real_mode(mode)}", fontsize=16)
plt.xlabel("Threads", fontsize=14)
plt.ylabel("OPS", fontsize=14)

2
scripts/rocksdb.sh
View File

@ -47,7 +47,7 @@ function samples() {
done
}
echo mode,threads,key_size,value_size,insert_ratio,ops,elapsed > "${script_dir}/rocksdb.csv"
echo mode,threads,key_size,value_size,insert_ratio,ops,elapsed_us > "${script_dir}/rocksdb.csv"
samples "$1" 1>> "${script_dir}/rocksdb.csv"
if [ -x "${script_dir}/bin/python" ]; then
(cd "${script_dir}" && "${script_dir}/bin/python" plot.py rocksdb.csv)

16
src/main.rs
View File

@ -236,9 +236,13 @@ fn main() {
x.join().unwrap();
}
let duration = start_time.elapsed();
let elapsed_us = start_time.elapsed().as_micros() as u64;
let total = total_ops.load(std::sync::atomic::Ordering::Relaxed);
let ops = (total as f64 / duration.as_secs_f64()) as usize;
let ops = if elapsed_us == 0 {
0
} else {
((total as u128 * 1_000_000u128) / elapsed_us as u128) as usize
};
let ratio = if args.mode == "mixed" {
args.insert_ratio
@ -257,13 +261,7 @@ fn main() {
}
println!(
"{},{},{},{},{},{},{}",
mode,
args.threads,
args.key_size,
args.value_size,
ratio,
ops,
duration.as_millis()
mode, args.threads, args.key_size, args.value_size, ratio, ops, elapsed_us
);
drop(db);
#[cfg(feature = "custom_alloc")]