From 0177b0afc1e497a5b4e4cfa1337587865d126542 Mon Sep 17 00:00:00 2001
From: abbycin <abbytsing@gmail.com>
Date: Tue, 10 Mar 2026 09:19:46 +0800
Subject: [PATCH] Align transaction benchmark semantics

---
 rocksdb/main.cpp | 19 +++++++++++++++++++
 src/main.rs      |  6 +++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/rocksdb/main.cpp b/rocksdb/main.cpp
index 61bb91c..ffa784c 100644
--- a/rocksdb/main.cpp
+++ b/rocksdb/main.cpp
@@ -560,6 +560,8 @@ static bool run_one_op(OpKind op, rocksdb::OptimisticTransactionDB *db, rocksdb:
         }
 
         auto *txn = db->BeginTransaction(wopt);
+        txn->SetSnapshot();
+        ropt.snapshot = txn->GetSnapshot();
         auto st = txn->Get(ropt, handle, key, &out);
         auto cst = txn->Commit();
         delete txn;
@@ -597,6 +599,21 @@ static bool run_one_op(OpKind op, rocksdb::OptimisticTransactionDB *db, rocksdb:
         }
 
         auto *txn = db->BeginTransaction(wopt);
+        if (spec.insert_only) {
+            auto pst = txn->Put(handle, key.value(), value);
+            auto cst = txn->Commit();
+            delete txn;
+            return pst.ok() && cst.ok();
+        }
+
+        rocksdb::ReadOptions update_ropt;
+        txn->SetSnapshot();
+        update_ropt.snapshot = txn->GetSnapshot();
+        auto gst = txn->GetForUpdate(update_ropt, handle, key.value(), static_cast<std::string *>(nullptr));
+        if (!gst.ok()) {
+            delete txn;
+            return false;
+        }
         auto pst = txn->Put(handle, key.value(), value);
         auto cst = txn->Commit();
         delete txn;
@@ -641,6 +658,8 @@ static bool run_one_op(OpKind op, rocksdb::OptimisticTransactionDB *db, rocksdb:
     }
 
     auto *txn = db->BeginTransaction(wopt);
+    txn->SetSnapshot();
+    ropt.snapshot = txn->GetSnapshot();
     auto *iter = txn->GetIterator(ropt);
     iter->Seek(prefix.value());
     size_t scanned = 0;
diff --git a/src/main.rs b/src/main.rs
index b2cc01a..972f7fc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1110,7 +1110,11 @@ fn run_one_op(
 
             if let Some(key) = key_opt {
                 if let Ok(tx) = bucket.begin() {
-                    let write_ok = tx.upsert(key.as_slice(), value.as_slice()).is_ok();
+                    let write_ok = if spec.insert_only {
+                        tx.upsert(key.as_slice(), value.as_slice()).is_ok()
+                    } else {
+                        tx.update(key.as_slice(), value.as_slice()).is_ok()
+                    };
                     if !write_ok {
                         false
                     } else {