From c9b915cd494b79c9a48a4dbb14ab388bcd918642 Mon Sep 17 00:00:00 2001 From: abbycin Date: Tue, 12 Aug 2025 17:13:06 +0800 Subject: [PATCH] adpat rocksdb --- Cargo.lock | 4 +- Cargo.toml | 9 +- rocksdb/.clang-format | 58 +++++++++ rocksdb/.clang-tidy | 14 +++ rocksdb/.gitignore | 3 + rocksdb/CMakeLists.txt | 32 ++++- rocksdb/CMakePresets.json | 40 ++++++ rocksdb/instant.h | 54 +++----- rocksdb/main.cpp | 252 ++++++++++++++++++++++++-------------- rocksdb/vcpkg.json | 5 + scripts/{x.sh => mace.sh} | 10 +- scripts/plot.py | 15 ++- scripts/rocksdb.sh | 41 +++++++ src/main.rs | 57 +++++---- 14 files changed, 412 insertions(+), 182 deletions(-) create mode 100644 rocksdb/.clang-format create mode 100644 rocksdb/.clang-tidy create mode 100644 rocksdb/.gitignore create mode 100644 rocksdb/CMakePresets.json create mode 100644 rocksdb/vcpkg.json rename scripts/{x.sh => mace.sh} (84%) create mode 100755 scripts/rocksdb.sh diff --git a/Cargo.lock b/Cargo.lock index f76b287..ff02ba2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -192,7 +192,6 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "io" version = "0.0.2" -source = "git+https://git.o2c.fun/abby/mace.git?branch=task-64#0a5ebdbdca9e1097da83e600a788288bfa4f6b85" dependencies = [ "libc", ] @@ -236,8 +235,7 @@ checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" [[package]] name = "mace" -version = "0.0.9" -source = "git+https://git.o2c.fun/abby/mace.git?branch=task-64#0a5ebdbdca9e1097da83e600a788288bfa4f6b85" +version = "0.0.10" dependencies = [ "crc32c", "crossbeam-epoch", diff --git a/Cargo.toml b/Cargo.toml index b7ee3d9..53ab474 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,10 +4,17 @@ version = "0.1.0" edition = "2024" [dependencies] -mace = { git = "https://git.o2c.fun/abby/mace.git", branch = "task-64" } +# mace = { git = "https://git.o2c.fun/abby/mace.git", branch = "task-64-1" } +mace = { path = "/home/workspace/gits/github/mace"} clap = { version = "4.5.42", features = ["derive"] } rand = "0.9.2" [profile.release] lto = true opt-level = 3 + +[profile.release-with-symbol] +inherits = "release" +# debug = true +debug = "line-tables-only" +panic = "abort" diff --git a/rocksdb/.clang-format b/rocksdb/.clang-format new file mode 100644 index 0000000..a8207b4 --- /dev/null +++ b/rocksdb/.clang-format @@ -0,0 +1,58 @@ +# Generated from CLion C/C++ Code Style settings +--- +Language: Cpp +BasedOnStyle: LLVM +AccessModifierOffset: -4 +AlignConsecutiveAssignments: false +AlignConsecutiveDeclarations: false +AlignOperands: true +AlignTrailingComments: false +AlwaysBreakTemplateDeclarations: Yes +BraceWrapping: + AfterCaseLabel: false + AfterClass: false + AfterControlStatement: false + AfterEnum: false + AfterFunction: false + AfterNamespace: false + AfterStruct: false + AfterUnion: false + AfterExternBlock: false + BeforeCatch: false + BeforeElse: false + BeforeLambdaBody: false + BeforeWhile: false + SplitEmptyFunction: true + SplitEmptyRecord: true + SplitEmptyNamespace: true +BreakBeforeBraces: Custom +BreakConstructorInitializers: AfterColon +BreakConstructorInitializersBeforeComma: false +ColumnLimit: 120 +ConstructorInitializerAllOnOneLineOrOnePerLine: false +ContinuationIndentWidth: 8 +IncludeCategories: + - Regex: '^<.*' + Priority: 1 + - Regex: '^".*' + Priority: 2 + - Regex: '.*' + Priority: 3 +IncludeIsMainRegex: '([-_](test|unittest))?$' +IndentCaseLabels: true +IndentWidth: 4 +InsertNewlineAtEOF: true +MacroBlockBegin: '' +MacroBlockEnd: '' +MaxEmptyLinesToKeep: 2 +NamespaceIndentation: All +SpaceAfterCStyleCast: true +SpaceAfterTemplateKeyword: false +SpaceBeforeRangeBasedForLoopColon: false +SpaceInEmptyParentheses: false +SpacesInAngles: false +SpacesInConditionalStatement: false +SpacesInCStyleCastParentheses: false +SpacesInParentheses: false +TabWidth: 4 +... diff --git a/rocksdb/.clang-tidy b/rocksdb/.clang-tidy new file mode 100644 index 0000000..7d8db4b --- /dev/null +++ b/rocksdb/.clang-tidy @@ -0,0 +1,14 @@ +Checks: > + -*, + clang-diagnostic-*, + -llvmlibc-restrict-system-libc-headers, + llvm-*, + misc-*, + -misc-const-correctness, + -misc-unused-parameters, + -misc-non-private-member-variables-in-classes, + -misc-no-recursion, + -misc-use-anonymous-namespace, + readability-identifier-naming, + -llvm-header-guard, + -llvm-include-order, diff --git a/rocksdb/.gitignore b/rocksdb/.gitignore new file mode 100644 index 0000000..2b02800 --- /dev/null +++ b/rocksdb/.gitignore @@ -0,0 +1,3 @@ +build/ +.cache +compile_commands.json diff --git a/rocksdb/CMakeLists.txt b/rocksdb/CMakeLists.txt index 74a1317..c747aae 100644 --- a/rocksdb/CMakeLists.txt +++ b/rocksdb/CMakeLists.txt @@ -1,10 +1,30 @@ -cmake_minimum_required(VERSION 4.0) -project(rockseb_bench) +cmake_minimum_required(VERSION 3.20) +project(rocksdb_bench CXX) set(CMAKE_CXX_STANDARD 23) -add_executable(rockseb_bench main.cpp +if (CMAKE_BUILD_TYPE MATCHES Release) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -flto") +endif() +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g") + +find_package(CLI11 CONFIG REQUIRED) + +add_executable(rocksdb_bench main.cpp instant.h) -set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3") -target_compile_options(rockseb_bench PRIVATE "-flto") -target_link_libraries(rockseb_bench rocksdb) + +target_include_directories(rocksdb_bench PRIVATE ${CMAKE_CURRENT_LIST_DIR}) +target_link_libraries(rocksdb_bench rocksdb CLI11::CLI11) + +file(GLOB ALL_SOURCE_FILES *.cpp *.h) +add_custom_target(format + COMMAND clang-format --style=file -i ${ALL_SOURCE_FILES} + VERBATIM +) + +add_custom_target(tidy + COMMAND clang-tidy -p=${CMAKE_BINARY_DIR} --config-file=${CMAKE_SOURCE_DIR}/.clang-tidy ${ALL_SOURCE_FILES} -extra-arg='--std=c++23' + COMMENT "Running Clang-Tidy" + VERBATIM +) diff --git a/rocksdb/CMakePresets.json b/rocksdb/CMakePresets.json new file mode 100644 index 0000000..c161e9d --- /dev/null +++ b/rocksdb/CMakePresets.json @@ -0,0 +1,40 @@ +{ + "version": 3, + "configurePresets":[ + { + "name": "basic", + "hidden": true, + "cacheVariables": { + "CMAKE_C_COMPILER": "gcc", + "CMAKe_CXX_COMPILER": "g++", + "CMAKE_TOOLCHAIN_FILE": "$env{VCPKG_ROOT}/scripts/buildsystems/vcpkg.cmake", + "CMAKE_EXPORT_COMPILE_COMMANDS": "ON", + "ENABLE_PROFILING": "ON" + } + }, + { + "name": "debug", + "inherits": "basic", + "binaryDir": "${sourceDir}/build/debug", + "cacheVariables": { + "CMAKE_BUILD_TYPE": "Debug", + "CMAKE_INSTALL_PREFIX": "${sourceDir}/dist/debug" + } + }, + { + "name": "release", + "inherits": "basic", + "binaryDir": "${sourceDir}/build/release", + "cacheVariables": { + "CMAKE_BUILD_TYPE": "Release", + "CMAKE_INSTALL_PREFIX": "${sourceDir}/dist/release" + } + } + ], + "buildPresets": [ + { + "name": "release", + "configurePreset": "release" + } + ] +} diff --git a/rocksdb/instant.h b/rocksdb/instant.h index b3f1c3a..9ccafce 100644 --- a/rocksdb/instant.h +++ b/rocksdb/instant.h @@ -3,53 +3,29 @@ #include -namespace nm -{ -class Instant { -public: - static Instant now() - { - return { std::chrono::steady_clock::now() }; - } +namespace nm { + class Instant { + public: + static Instant now() { return {std::chrono::steady_clock::now()}; } - auto elapse_usec() - { - return diff().count() / 1e3; - } + auto elapse_usec() { return diff().count() / 1e3; } - auto elapse_ms() - { - return diff().count() / 1e6; - } + auto elapse_ms() { return diff().count() / 1e6; } - auto elapse_sec() - { - return diff().count() / 1e9; - } + auto elapse_sec() { return diff().count() / 1e9; } - auto elapse_min() - { - return elapse_sec() / 60.0; - } + auto elapse_min() { return elapse_sec() / 60.0; } - void reset() - { - *this = now(); - } + void reset() { *this = now(); } -private: - using timepoint = std::chrono::steady_clock::time_point; - Instant(timepoint now) : tp_ { now } - { - } + private: + using timepoint = std::chrono::steady_clock::time_point; + Instant(timepoint now) : tp_{now} {} - auto diff() -> std::chrono::duration - { - return std::chrono::steady_clock::now() - tp_; - } + auto diff() -> std::chrono::duration { return std::chrono::steady_clock::now() - tp_; } - timepoint tp_; -}; + timepoint tp_; + }; } // namespace nm #endif // INSTANT_20231013102058 diff --git a/rocksdb/main.cpp b/rocksdb/main.cpp index 55f975a..e31ee3d 100644 --- a/rocksdb/main.cpp +++ b/rocksdb/main.cpp @@ -1,3 +1,7 @@ +#include +#include +#include +#include #include #include #include @@ -5,115 +9,177 @@ #include #include -#include #include #include #include -#include -#include #include +#include "CLI/CLI.hpp" #include "instant.h" -int main() -{ - std::string db_root = "/home/abby/rocksdb_tmp"; - std::filesystem::remove_all(db_root); - rocksdb::Options options; - options.disable_auto_compactions = true; - options.create_if_missing = true; - options.max_write_buffer_number = 10; - options.target_file_size_base = 64 << 20; - options.write_buffer_size = 64 << 20; - options.level0_file_num_compaction_trigger = 500; - options.max_bytes_for_level_base = 2 << 30; - constexpr size_t count = 100000; - constexpr size_t workers = 4; - auto ropt = rocksdb::ReadOptions(); - auto wopt = rocksdb::WriteOptions(); - std::vector wg; - std::vector> keys {}; - rocksdb::OptimisticTransactionDB *db; - std::cout << "db_root " << db_root << '\n'; - auto s = rocksdb::OptimisticTransactionDB::Open(options, db_root, &db); - assert(s.ok()); - std::string val(1024, 'x'); +struct Args { + size_t threads; + size_t iterations; + size_t key_size; + size_t value_size; + size_t insert_ratio; + bool random; + std::string mode; + std::string path; +}; - for (size_t tid = 0; tid < workers; ++tid) { - std::vector key {}; +int main(int argc, char *argv[]) { + CLI::App app{"rocksdb bench"}; + Args args{ + .threads = 4, + .iterations = 100000, + .key_size = 16, + .value_size = 1024, + .insert_ratio = 30, + .mode = "insert", + .path = "/tmp/rocksdb_tmp", + }; - for (size_t i = 0; i < count; ++i) { - auto tmp = std::format("key_{}_{}", tid, i); - tmp.resize(1024, 'x'); - key.push_back(std::move(tmp)); - } - keys.emplace_back(std::move(key)); - } + app.add_option("-m,--mode", args.mode, "Mode: insert, get, mixed"); + 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("-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"); + app.add_option("--random", args.random, "Shuffle insert keys"); - for (size_t tid = 0; tid < workers; ++tid) { - auto ks = &keys[tid]; - wg.emplace_back( - [&] - { - for (auto &x : *ks) { - auto kv = db->BeginTransaction(wopt); - kv->Put(x, x); - kv->Commit(); - delete kv; - } - }); - } + CLI11_PARSE(app, argc, argv); - for (auto &w : wg) { - w.join(); - } + if (args.path.empty()) { + std::println("path is empty"); + return 1; + } - delete db; - wg.clear(); + if (std::filesystem::exists(args.path)) { + std::println("path `{}` already exists", args.path); + return 1; + } - auto sts = - rocksdb::OptimisticTransactionDB::Open(options, db_root, &db); - assert(sts.ok()); - std::barrier g { workers }; - auto b = nm::Instant::now(); - std::mutex mtx {}; - std::atomic operation { 0 }; + if (args.mode != "insert" && args.mode != "get" && args.mode != "mixed") { + std::println("Error: Invalid mode"); + return 1; + } - for (size_t tid = 0; tid < workers; ++tid) { - auto ks = &keys[tid]; - wg.emplace_back( - [&] - { - g.arrive_and_wait(); + if (args.key_size < 16 || args.value_size < 16) { + std::println("Error: key_size or value_size too small, must >= 16"); + return 1; + } - std::string val; - if (mtx.try_lock()) { - b = nm::Instant::now(); - mtx.unlock(); - } + if (args.insert_ratio > 100) { + std::println("Error: Insert ratio must be between 0 and 100"); + return 1; + } - for (auto &x : *ks) { - auto kv = db->BeginTransaction(wopt); - kv->Get(ropt, x, &val); - kv->Commit(); - delete kv; - } - operation.fetch_add(count, - std::memory_order_relaxed); - }); - } + rocksdb::Options options; + options.disable_auto_compactions = true; + options.create_if_missing = true; + options.max_write_buffer_number = 10; + options.target_file_size_base = 64 << 20; + options.write_buffer_size = 64 << 20; + options.level0_file_num_compaction_trigger = 500; + options.max_bytes_for_level_base = 2 << 30; + auto ropt = rocksdb::ReadOptions(); + auto wopt = rocksdb::WriteOptions(); + std::vector wg; + std::vector> keys{}; + std::atomic total_op{0}; + rocksdb::OptimisticTransactionDB *db; + auto b = nm::Instant::now(); + std::mutex mtx{}; + auto s = rocksdb::OptimisticTransactionDB::Open(options, args.path, &db); + assert(s.ok()); + std::barrier barrier{static_cast(args.threads)}; - for (auto &w : wg) { - w.join(); - } + std::random_device rd{}; + std::mt19937 gen(rd()); + std::uniform_int_distribution dist(0, 100); - printf("thread %ld\niterations %ld\nkey_size %ld\nvalue_size %ld\nops: " - "%.2f\n", - workers, - count, - keys[0][0].size(), - val.size(), - static_cast(operation.load(std::memory_order_relaxed)) / - b.elapse_sec()); - delete db; + + std::string val(args.value_size, 'x'); + for (size_t tid = 0; tid < args.threads; ++tid) { + std::vector key{}; + for (size_t i = 0; i < args.iterations; ++i) { + auto tmp = std::format("key_{}_{}", tid, i); + tmp.resize(args.key_size, 'x'); + key.emplace_back(std::move(tmp)); + } + if (args.random) { + std::shuffle(keys.begin(), keys.end(), gen); + } + keys.emplace_back(std::move(key)); + } + + + if (args.mode == "get") { + auto *kv = db->BeginTransaction(wopt); + for (size_t tid = 0; tid < args.threads; ++tid) { + auto *tk = &keys[tid]; + for (auto &key: *tk) { + kv->Put(key, val); + } + } + kv->Commit(); + delete kv; + } + + for (size_t tid = 0; tid < args.threads; ++tid) { + auto *tk = &keys[tid]; + wg.emplace_back([&] { + std::string rval(args.value_size, '0'); + barrier.arrive_and_wait(); + if (mtx.try_lock()) { + b = nm::Instant::now(); + mtx.unlock(); + } + + if (args.mode == "insert") { + for (auto &key: *tk) { + auto *kv = db->BeginTransaction(wopt); + kv->Put(key, val); + kv->Commit(); + delete kv; + } + + } else if (args.mode == "get") { + for (auto &key: *tk) { + auto *kv = db->BeginTransaction(wopt); + kv->Get(ropt, key, &rval); + kv->Commit(); + delete kv; + } + } else if (args.mode == "mixed") { + for (auto &key: *tk) { + auto is_insert = dist(gen) < args.insert_ratio; + auto *kv = db->BeginTransaction(wopt); + if (is_insert) { + kv->Put(key, val); + } else { + kv->Get(ropt, key, &rval); // not found + } + kv->Commit(); + delete kv; + } + } + total_op.fetch_add(args.iterations, std::memory_order::relaxed); + }); + } + + for (auto &w: wg) { + w.join(); + } + size_t ratio = [&args] -> size_t { + if (args.mode == "mixed") + return args.insert_ratio; + return args.mode == "insert" ? 100 : 0; + }(); + double ops = static_cast(total_op.load(std::memory_order_relaxed)) / b.elapse_sec(); + std::println("{},{},{},{},{},{:.2f}", args.mode, args.threads, args.key_size, args.value_size, ratio, ops); + delete db; + std::filesystem::remove_all(args.path); } diff --git a/rocksdb/vcpkg.json b/rocksdb/vcpkg.json new file mode 100644 index 0000000..f1d17f9 --- /dev/null +++ b/rocksdb/vcpkg.json @@ -0,0 +1,5 @@ +{ + "dependencies": [ + "cli11" + ] +} diff --git a/scripts/x.sh b/scripts/mace.sh similarity index 84% rename from scripts/x.sh rename to scripts/mace.sh index 7735967..f0feac5 100755 --- a/scripts/x.sh +++ b/scripts/mace.sh @@ -21,20 +21,20 @@ function samples() { ./target/release/kv_bench --path /home/abby/mace_bench --threads $i --iterations 100000 --mode get --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} if test $? -ne 0 then - echo "insert threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" + echo "get threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" exit 1 fi ./target/release/kv_bench --path /home/abby/mace_bench --threads $i --iterations 100000 --mode mixed --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} --insert-ratio 30 if test $? -ne 0 then - echo "insert threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" + echo "mixed threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" exit 1 fi done done } -echo mode,threads,key_size,value_size,insert_ratio,ops > scripts/x.csv -samples 2>> scripts/x.csv +echo mode,threads,key_size,value_size,insert_ratio,ops > scripts/mace.csv +samples 2>> scripts/mace.csv popd -./bin/python plot.py +./bin/python plot.py mace.csv diff --git a/scripts/plot.py b/scripts/plot.py index b98f05a..2a82a1f 100644 --- a/scripts/plot.py +++ b/scripts/plot.py @@ -1,14 +1,17 @@ import pandas as pd import matplotlib.pyplot as plt from adjustText import adjust_text +import sys def real_mode(m): if m == "mixed": - return "mixed (70% read, 30% write)" - return m + return "MIXED (70% Get, 30% Put)" + return m.upper() +name = sys.argv[1] +prefix = name.split(".")[0] # 读取数据 -df = pd.read_csv("./x.csv") +df = pd.read_csv(f"./{name}") # 按 mode 分组 modes = df["mode"].unique() @@ -40,11 +43,11 @@ for mode in modes: adjust_text(texts, arrowprops=dict(arrowstyle="->", color='gray')) # 设置图表样式 - plt.title(f"Performance: {real_mode(mode).upper()}", fontsize=16) + plt.title(f"{prefix.upper()}: {real_mode(mode)}", fontsize=16) plt.xlabel("Threads", fontsize=14) plt.ylabel("OPS", fontsize=14) plt.grid(True, linestyle="--", alpha=0.6) plt.legend() plt.tight_layout() - plt.savefig(f"{mode}.png") - plt.close() \ No newline at end of file + plt.savefig(f"{prefix}_{mode}.png") + plt.close() diff --git a/scripts/rocksdb.sh b/scripts/rocksdb.sh new file mode 100755 index 0000000..c0df2a5 --- /dev/null +++ b/scripts/rocksdb.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash + +pushd . +cd ../rocksdb +cmake --preset release 1>/dev/null 2>/dev/null +cmake --build --preset release 1>/dev/null 2>/dev/null + +function samples() { + kv_sz=(16 16 100 1024 1024 1024) + # set -x + + for ((i = 1; i <= $(nproc); i *= 2)) + do + for ((j = 0; j < ${#kv_sz[@]}; j += 2)) + do + ./build/release/rocksdb_bench --path /home/abby/rocksdb_tmp --threads $i --iterations 100000 --mode insert --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} + if test $? -ne 0 + then + echo "insert threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" + exit 1 + fi + ./build/release/rocksdb_bench --path /home/abby/rocksdb_tmp --threads $i --iterations 100000 --mode get --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} + if test $? -ne 0 + then + echo "get threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" + exit 1 + fi + ./build/release/rocksdb_bench --path /home/abby/rocksdb_tmp --threads $i --iterations 100000 --mode mixed --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} --insert-ratio 30 + if test $? -ne 0 + then + echo "mixed threads $i ksz ${kv_sz[j]} vsz ${kv_sz[j+1]} fail" + exit 1 + fi + done + done +} + +echo mode,threads,key_size,value_size,insert_ratio,ops > ../scripts/rocksdb.csv +samples 1>> ../scripts/rocksdb.csv +popd +./bin/python plot.py rocksdb.csv diff --git a/src/main.rs b/src/main.rs index 3843bd4..f02bb46 100644 --- a/src/main.rs +++ b/src/main.rs @@ -59,26 +59,36 @@ fn main() { return; } + let mut keys: Vec>> = Vec::with_capacity(args.threads); let mut opt = Options::new(path); opt.sync_on_write = false; opt.tmp_store = true; - // make sure there's no remote indirection - opt.max_inline_size = 4096; // opt.cache_capacity = 3 << 30; // this is very important for large key-value store let db = Mace::new(opt.validate().unwrap()).unwrap(); + let mut rng = rand::rng(); let value = Arc::new(vec![b'0'; args.value_size]); - if args.mode == "get" { - for tid in 0..args.threads { - for i in 0..args.iterations { - let key = format!("key_{tid}_{i}"); - let mut tmp = key.into_bytes(); - tmp.resize(args.key_size, b'x'); - let pre_tx = db.begin().unwrap(); - pre_tx.put(&tmp, &*value).unwrap(); - pre_tx.commit().unwrap(); - } + for tid in 0..args.threads { + let mut tk = Vec::with_capacity(args.iterations); + for i in 0..args.iterations { + let mut key = format!("key_{tid}_{i}").into_bytes(); + key.resize(args.key_size, b'x'); + tk.push(key); } + if args.random { + tk.shuffle(&mut rng); + } + keys.push(tk); + } + + if args.mode == "get" { + let pre_tx = db.begin().unwrap(); + (0..args.threads).for_each(|tid| { + for i in 0..args.iterations { + pre_tx.put(&keys[tid][i], &*value).unwrap(); + } + }); + pre_tx.commit().unwrap(); } let barrier = Arc::new(std::sync::Barrier::new(args.threads)); @@ -88,6 +98,7 @@ fn main() { let h: Vec> = (0..args.threads) .map(|tid| { let db = db.clone(); + let tk: &Vec> = unsafe { std::mem::transmute(&keys[tid]) }; let total_ops = total_ops.clone(); let barrier = Arc::clone(&barrier); let mode = args.mode.clone(); @@ -96,18 +107,6 @@ fn main() { let val = value.clone(); std::thread::spawn(move || { - let mut keys: Vec> = Vec::with_capacity(args.iterations); - for i in 0..args.iterations { - let key = format!("key_{tid}_{i}"); - let mut tmp = key.into_bytes(); - tmp.resize(args.key_size, b'x'); - keys.push(tmp); - } - - let mut rng = rand::rng(); - if args.random { - keys.shuffle(&mut rng); - } barrier.wait(); { @@ -118,21 +117,21 @@ fn main() { match mode.as_str() { "insert" => { - for key in &keys { + for key in tk { let tx = db.begin().unwrap(); tx.put(key.as_slice(), val.as_slice()).unwrap(); tx.commit().unwrap(); } } "get" => { - for key in &keys { + for key in tk { let tx = db.view().unwrap(); tx.get(key).unwrap(); } } "mixed" => { - for key in &keys { - let is_insert = rng.random_range(0..100) < insert_ratio; + for key in tk { + let is_insert = rand::random_range(0..100) < insert_ratio; if is_insert { let tx = db.begin().unwrap(); @@ -140,7 +139,7 @@ fn main() { tx.commit().unwrap(); } else { let tx = db.view().unwrap(); - let _ = tx.get(key); // may not insert + let _ = tx.get(key); // not found } } }