fix rocksdb iteration calc

rocksdb/main.cpp

@@ -10,6 +10,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>
@@ -19,6 +20,7 @@
 #include <filesystem>
 #include <format>
 #include <string>
+#include <syncstream>
 
 #include "CLI/CLI.hpp"
 #include "instant.h"
@@ -87,6 +89,22 @@ 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;
@@ -172,6 +190,13 @@ int main(int argc, char *argv[]) {
             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;
+            }
+            ropt.prefix_same_as_start = true;
+            size_t round = 0;
 
             barrier.arrive_and_wait();
             if (mtx.try_lock()) {
@@ -181,6 +206,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();
@@ -189,6 +215,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();
@@ -196,6 +223,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) {
@@ -209,14 +237,17 @@ 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;
                     black_box(iter->key());
                     black_box(iter->value());
                     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);
         });
     }
 

src/main.rs

@@ -134,10 +134,10 @@ fn main() {
             let st = start_time.clone();
             let val = value.clone();
             let prefix = format!("key_{tid}");
-            println!("prefix {}", prefix);
 
             std::thread::spawn(move || {
                 // coreid::bind_core(tid);
+                let mut round = 0;
                 barrier.wait();
 
                 {
@@ -145,10 +145,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();
@@ -156,13 +156,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();
@@ -170,7 +173,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);
                             }
                         }
                     }
@@ -178,13 +182,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();
@@ -217,7 +222,7 @@ fn main() {
     };
     // eprintln!("mode,threads,key_size,value_size,insert_ratio,ops");
     eprintln!(
-        "{},{},{},{},{},{:.2},{}",
+        "{},{},{},{},{},{},{}",
         args.mode,
         args.threads,
         args.key_size,