From db0960800a2971e94c73d7e7eb1bd2e9b8d3c54c Mon Sep 17 00:00:00 2001
From: Levi Tamasi <ltamasi@meta.com>
Date: Wed, 22 May 2024 11:30:33 -0700
Subject: [PATCH] Add Transaction::PutEntity to the stress tests (#12688)

Summary:
Pull Request resolved: https://github.com/facebook/rocksdb/pull/12688

As a first step of covering the wide-column transaction APIs, the patch adds `PutEntity` to the optimistic and pessimistic transaction stress tests (for the latter, only when the WriteCommitted policy is utilized). Other APIs and the multi-operation transaction test will be covered by subsequent PRs.

Reviewed By: jaykorean

Differential Revision: D57675781

fbshipit-source-id: bfe062ec5f6ab48641cd99a70f239ce4aa39299c
---
 db_stress_tool/db_stress_test_base.cc   | 17 +++++++++++++++--
 db_stress_tool/db_stress_tool.cc        |  6 +++---
 db_stress_tool/no_batched_ops_stress.cc | 16 +++++++++++-----
 tools/db_crashtest.py                   |  7 +++----
 4 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc
index 1179ba9c92..d4102fb831 100644
--- a/db_stress_tool/db_stress_test_base.cc
+++ b/db_stress_tool/db_stress_test_base.cc
@@ -617,8 +617,21 @@ void StressTest::PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
 
       if (FLAGS_use_put_entity_one_in > 0 &&
           (value_base % FLAGS_use_put_entity_one_in) == 0) {
-        s = db_->PutEntity(write_opts, cfh, key,
-                           GenerateWideColumns(value_base, v));
+        if (!FLAGS_use_txn) {
+          if (FLAGS_use_attribute_group) {
+            s = db_->PutEntity(write_opts, key,
+                               GenerateAttributeGroups({cfh}, value_base, v));
+          } else {
+            s = db_->PutEntity(write_opts, cfh, key,
+                               GenerateWideColumns(value_base, v));
+          }
+        } else {
+          s = ExecuteTransaction(
+              write_opts, /*thread=*/nullptr, [&](Transaction& txn) {
+                return txn.PutEntity(cfh, key,
+                                     GenerateWideColumns(value_base, v));
+              });
+        }
       } else if (FLAGS_use_merge) {
         if (!FLAGS_use_txn) {
           if (FLAGS_user_timestamp_size > 0) {
diff --git a/db_stress_tool/db_stress_tool.cc b/db_stress_tool/db_stress_tool.cc
index e7756a277c..34a45a603c 100644
--- a/db_stress_tool/db_stress_tool.cc
+++ b/db_stress_tool/db_stress_tool.cc
@@ -301,11 +301,11 @@ int db_stress_tool(int argc, char** argv) {
   }
 
   if (FLAGS_use_put_entity_one_in > 0 &&
-      (FLAGS_use_full_merge_v1 || FLAGS_use_txn || FLAGS_test_multi_ops_txns ||
+      (FLAGS_use_full_merge_v1 || FLAGS_test_multi_ops_txns ||
        FLAGS_user_timestamp_size > 0)) {
     fprintf(stderr,
-            "Wide columns are incompatible with V1 Merge, transactions, and "
-            "user-defined timestamps\n");
+            "Wide columns are incompatible with V1 Merge, the multi-op "
+            "transaction test, and user-defined timestamps\n");
     exit(1);
   }
 
diff --git a/db_stress_tool/no_batched_ops_stress.cc b/db_stress_tool/no_batched_ops_stress.cc
index c18cf506b6..0d0e0c5e14 100644
--- a/db_stress_tool/no_batched_ops_stress.cc
+++ b/db_stress_tool/no_batched_ops_stress.cc
@@ -1483,12 +1483,18 @@ class NonBatchedOpsStressTest : public StressTest {
 
     if (FLAGS_use_put_entity_one_in > 0 &&
         (value_base % FLAGS_use_put_entity_one_in) == 0) {
-      if (FLAGS_use_attribute_group) {
-        s = db_->PutEntity(write_opts, k,
-                           GenerateAttributeGroups({cfh}, value_base, v));
+      if (!FLAGS_use_txn) {
+        if (FLAGS_use_attribute_group) {
+          s = db_->PutEntity(write_opts, k,
+                             GenerateAttributeGroups({cfh}, value_base, v));
+        } else {
+          s = db_->PutEntity(write_opts, cfh, k,
+                             GenerateWideColumns(value_base, v));
+        }
       } else {
-        s = db_->PutEntity(write_opts, cfh, k,
-                           GenerateWideColumns(value_base, v));
+        s = ExecuteTransaction(write_opts, thread, [&](Transaction& txn) {
+          return txn.PutEntity(cfh, k, GenerateWideColumns(value_base, v));
+        });
       }
     } else if (FLAGS_use_timed_put_one_in > 0 &&
                ((value_base + kLargePrimeForCommonFactorSkew) %
diff --git a/tools/db_crashtest.py b/tools/db_crashtest.py
index 94bdd7bb06..afded2211e 100644
--- a/tools/db_crashtest.py
+++ b/tools/db_crashtest.py
@@ -493,8 +493,6 @@ txn_params = {
     # pipeline write is not currnetly compatible with WritePrepared txns
     "enable_pipelined_write": 0,
     "create_timestamped_snapshot_one_in": random.choice([0, 20]),
-    # PutEntity in transactions is not yet implemented
-    "use_put_entity_one_in": 0,
     # Should not be used with TransactionDB which uses snapshot.
     "inplace_update_support": 0,
     # TimedPut is not supported in transaction
@@ -508,8 +506,6 @@ optimistic_txn_params = {
     "occ_validation_policy": random.randint(0, 1),
     "share_occ_lock_buckets": random.randint(0, 1),
     "occ_lock_bucket_count": lambda: random.choice([10, 100, 500]),
-    # PutEntity in transactions is not yet implemented
-    "use_put_entity_one_in": 0,
     # Should not be used with OptimisticTransactionDB which uses snapshot.
     "inplace_update_support": 0,
     # TimedPut is not supported in transaction
@@ -801,6 +797,9 @@ def finalize_and_sanitize(src_params):
     if dest_params.get("use_txn") == 1 and dest_params.get("txn_write_policy", 0) != 0:
         dest_params["sync_fault_injection"] = 0
         dest_params["manual_wal_flush_one_in"] = 0
+        # Wide-column pessimistic transaction APIs are initially supported for
+        # WriteCommitted only
+        dest_params["use_put_entity_one_in"] = 0
     # Wide column stress tests require FullMergeV3
     if dest_params["use_put_entity_one_in"] != 0:
         dest_params["use_full_merge_v1"] = 0