share iterations among threads

rocksdb/main.cpp

@@ -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();
+            }
+        }
+    }
+
+    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);

scripts/mace.sh

@@ -1,5 +1,11 @@
 #!/usr/bin/env bash
 
+if [ "$#" -ne 1 ]
+then
+  printf "\033[m$0 path\033[0m\n"
+  exit 1
+fi
+
 pushd .
 cd ..
 cargo build --release 1>/dev/null 2> /dev/null
@@ -8,31 +14,32 @@ function samples() {
         export RUST_BACKTRACE=full
         kv_sz=(16 16 100 1024 1024 1024 16 10240)
         # set -x
+        db_root=$1
 
-        cnt=10000
+        cnt=100000
         for ((i = 1; i <= $(nproc); i *= 2))
         do
                 for ((j = 0; j < ${#kv_sz[@]}; j += 2))
                 do
-                        ./target/release/kv_bench --path /home/abby/mace_bench --threads $i --iterations $cnt --mode insert --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]}
+                        ./target/release/kv_bench --path $db_root --threads $i --iterations $cnt --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
-                        ./target/release/kv_bench --path /home/abby/mace_bench --threads $i --iterations $cnt --mode get --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]}
+                        ./target/release/kv_bench --path $db_root --threads $i --iterations $cnt --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
-                        ./target/release/kv_bench --path /home/abby/mace_bench --threads $i --iterations $cnt --mode mixed --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} --insert-ratio 30
+                        ./target/release/kv_bench --path $db_root --threads $i --iterations $cnt --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
-                        ./target/release/kv_bench --path /home/abby/mace_bench --threads $i --iterations $cnt --mode scan --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} --insert-ratio 30
+                        ./target/release/kv_bench --path $db_root --threads $i --iterations $cnt --mode scan --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"
@@ -43,6 +50,6 @@ function samples() {
 }
 
 echo mode,threads,key_size,value_size,insert_ratio,ops,elasped > scripts/mace.csv
-samples 2>> scripts/mace.csv
+samples $1 2>> scripts/mace.csv
 popd
 ./bin/python plot.py mace.csv

scripts/rocksdb.sh

@@ -1,5 +1,11 @@
 #!/usr/bin/env bash
 
+if [ "$#" -ne 1 ]
+then
+  printf "\033[m$0 path\033[0m\n"
+  exit 1
+fi
+
 pushd .
 cd ../rocksdb
 cmake --preset release 1>/dev/null 2>/dev/null
@@ -8,30 +14,31 @@ cmake --build --preset release 1>/dev/null 2>/dev/null
 function samples() {
         kv_sz=(16 16 100 1024 1024 1024 16 10240)
         # set -x
-        cnt=10000
+        db_root=$1
+        cnt=100000
         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 $cnt --mode insert --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]}
+                        ./build/release/rocksdb_bench --path $db_root --threads $i --iterations $cnt --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 $cnt --mode get --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]}
+                        ./build/release/rocksdb_bench --path $db_root --threads $i --iterations $cnt --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 $cnt --mode mixed --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} --insert-ratio 30
+                        ./build/release/rocksdb_bench --path $db_root --threads $i --iterations $cnt --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
-                        ./build/release/rocksdb_bench --path /home/abby/rocksdb_tmp --threads $i --iterations $cnt --mode scan --key-size ${kv_sz[j]} --value-size ${kv_sz[j+1]} --insert-ratio 30
+                        ./build/release/rocksdb_bench --path $db_root --threads $i --iterations $cnt --mode scan --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"
@@ -42,6 +49,6 @@ function samples() {
 }
 
 echo mode,threads,key_size,value_size,insert_ratio,ops,elapsed > ../scripts/rocksdb.csv
-samples 1>> ../scripts/rocksdb.csv
+samples $1 1>> ../scripts/rocksdb.csv
 popd
 ./bin/python plot.py rocksdb.csv

src/main.rs

@@ -89,9 +89,10 @@ fn main() {
 
     let mut rng = rand::rng();
     let value = Arc::new(vec![b'0'; args.value_size]);
+    let keys_per_thread = args.iterations / args.threads;
     for tid in 0..args.threads {
-        let mut tk = Vec::with_capacity(args.iterations);
-        for i in 0..args.iterations {
+        let mut tk = Vec::with_capacity(keys_per_thread);
+        for i in 0..keys_per_thread {
             let mut key = format!("key_{tid}_{i}").into_bytes();
             key.resize(args.key_size, b'x');
             tk.push(key);
@@ -105,8 +106,8 @@ fn main() {
     if args.mode == "get" || args.mode == "scan" {
         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();
+            for k in &keys[tid] {
+                pre_tx.put(k, &*value).unwrap();
             }
         });
         pre_tx.commit().unwrap();
@@ -114,6 +115,15 @@ fn main() {
         // re-open db
         saved.tmp_store = true;
         db = Mace::new(saved.validate().unwrap()).unwrap();
+
+        // simulate common use cases
+        for i in 0..keys_per_thread {
+            let tid = rng.random_range(0..args.threads);
+            let mut k = format!("key_{tid}_{i}").into_bytes();
+            k.resize(args.key_size, b'x');
+            let view = db.view().unwrap();
+            view.get(&k).unwrap();
+        }
     }
 
     let barrier = Arc::new(std::sync::Barrier::new(args.threads));
@@ -134,6 +144,7 @@ fn main() {
 
             std::thread::spawn(move || {
                 // coreid::bind_core(tid);
+                let mut round = 0;
                 barrier.wait();
 
                 {
@@ -141,10 +152,10 @@ fn main() {
                         *guard = Instant::now();
                     }
                 }
-
                 match mode.as_str() {
                     "insert" => {
                         for key in tk {
+                            round += 1;
                             let tx = db.begin().unwrap();
                             tx.put(key.as_slice(), val.as_slice()).unwrap();
                             tx.commit().unwrap();
@@ -152,13 +163,16 @@ fn main() {
                     }
                     "get" => {
                         for key in tk {
+                            round += 1;
                             let tx = db.view().unwrap();
-                            tx.get(key).unwrap();
+                            let x = tx.get(key).unwrap();
+                            std::hint::black_box(x);
                         }
                     }
                     "mixed" => {
                         for key in tk {
                             let is_insert = rand::random_range(0..100) < insert_ratio;
+                            round += 1;
 
                             if is_insert {
                                 let tx = db.begin().unwrap();
@@ -166,7 +180,8 @@ fn main() {
                                 tx.commit().unwrap();
                             } else {
                                 let tx = db.view().unwrap();
-                                let _ = tx.get(key); // not found
+                                let x = tx.get(key); // not found
+                                let _ = std::hint::black_box(x);
                             }
                         }
                     }
@@ -174,13 +189,14 @@ fn main() {
                         let view = db.view().unwrap();
                         let iter = view.seek(prefix);
                         for x in iter {
+                            round += 1;
                             std::hint::black_box(x);
                         }
                     }
                     _ => panic!("Invalid mode"),
                 }
 
-                total_ops.fetch_add(args.iterations, std::sync::atomic::Ordering::Relaxed);
+                total_ops.fetch_add(round, std::sync::atomic::Ordering::Relaxed);
             })
         })
         .collect();
@@ -194,16 +210,6 @@ fn main() {
     let total = total_ops.load(std::sync::atomic::Ordering::Relaxed);
     let ops = (total as f64 / duration.as_secs_f64()) as usize;
 
-    // println!("{:<20} {}", "Test Mode:", args.mode);
-    // println!("{:<20} {}", "Threads:", args.threads);
-    // println!("{:<20} {}", "Iterations", args.iterations);
-    // println!("{:<20} {}B", "Key Size:", args.key_size);
-    // println!("{:<20} {}B", "Value Size:", args.value_size);
-    // println!("{:<20} {ops}", "Total Ops:");
-    // if args.mode == "mixed" {
-    //     println!("{:<20} {}%", "Insert Ratio:", args.insert_ratio);
-    // }
-
     let ratio = if args.mode == "mixed" {
         args.insert_ratio
     } else if args.mode == "insert" {
@@ -211,9 +217,8 @@ fn main() {
     } else {
         0
     };
-    // eprintln!("mode,threads,key_size,value_size,insert_ratio,ops");
     eprintln!(
-        "{},{},{},{},{},{:.2},{}",
+        "{},{},{},{},{},{},{}",
         args.mode,
         args.threads,
         args.key_size,