From d4da4128646c05c29d57fd788f4eb5062dcd7405 Mon Sep 17 00:00:00 2001
From: anand76 <anand76@devvm1373.frc2.facebook.com>
Date: Fri, 20 Dec 2019 23:10:58 -0800
Subject: [PATCH] Add Transaction::MultiGet to db_stress (#6227)

Summary:
Call Transaction::MultiGet from TestMultiGet() in db_stress. We add some Puts/Merges/Deletes into the transaction in order to exercise MultiGetFromBatchAndDB. There is no data verification on read, just check status. There is currently no read data verification in db_stress as it requires synchronization with writes. It needs to be tackled separately.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/6227

Test Plan: make crash_test_with_txn

Differential Revision: D19204611

Pulled By: anand1976

fbshipit-source-id: 770d0e30d002e88626c264c58103f1d709bb060c
---
 db_stress_tool/db_stress_test_base.cc   | 11 +++++
 db_stress_tool/db_stress_test_base.h    |  2 +
 db_stress_tool/no_batched_ops_stress.cc | 60 ++++++++++++++++++++++++-
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc
index be65b198b..c0ecb926a 100644
--- a/db_stress_tool/db_stress_test_base.cc
+++ b/db_stress_tool/db_stress_test_base.cc
@@ -470,6 +470,17 @@ Status StressTest::CommitTxn(Transaction* txn) {
   delete txn;
   return s;
 }
+
+Status StressTest::RollbackTxn(Transaction* txn) {
+  if (!FLAGS_use_txn) {
+    return Status::InvalidArgument(
+        "RollbackTxn when FLAGS_use_txn is not"
+        " set");
+  }
+  Status s = txn->Rollback();
+  delete txn;
+  return s;
+}
 #endif
 
 void StressTest::OperateDb(ThreadState* thread) {
diff --git a/db_stress_tool/db_stress_test_base.h b/db_stress_tool/db_stress_test_base.h
index f16af3b0b..0f0dbbc82 100644
--- a/db_stress_tool/db_stress_test_base.h
+++ b/db_stress_tool/db_stress_test_base.h
@@ -52,6 +52,8 @@ class StressTest {
   Status NewTxn(WriteOptions& write_opts, Transaction** txn);
 
   Status CommitTxn(Transaction* txn);
+
+  Status RollbackTxn(Transaction* txn);
 #endif
 
   virtual void MaybeClearOneColumnFamily(ThreadState* /* thread */) {}
diff --git a/db_stress_tool/no_batched_ops_stress.cc b/db_stress_tool/no_batched_ops_stress.cc
index da7796e36..938c75ced 100644
--- a/db_stress_tool/no_batched_ops_stress.cc
+++ b/db_stress_tool/no_batched_ops_stress.cc
@@ -166,12 +166,68 @@ class NonBatchedOpsStressTest : public StressTest {
     std::vector<Status> statuses(num_keys);
     ColumnFamilyHandle* cfh = column_families_[rand_column_families[0]];
 
+    // Create a transaction in order to write some data. The purpose is to
+    // exercise WriteBatchWithIndex::MultiGetFromBatchAndDB. The transaction
+    // will be rolled back once MultiGet returns.
+#ifndef ROCKSDB_LITE
+    Transaction* txn = nullptr;
+    if (FLAGS_use_txn) {
+      WriteOptions wo;
+      NewTxn(wo, &txn);
+    }
+#endif
     for (size_t i = 0; i < num_keys; ++i) {
       key_str.emplace_back(Key(rand_keys[i]));
       keys.emplace_back(key_str.back());
+#ifndef ROCKSDB_LITE
+      if (FLAGS_use_txn) {
+        // With a 1 in 10 probability, insert the just added key in the batch
+        // into the transaction. This will create an overlap with the MultiGet
+        // keys and exercise some corner cases in the code
+        if (thread->rand.OneIn(10)) {
+          int op = thread->rand.Uniform(2);
+          Status s;
+          switch (op) {
+            case 0:
+            case 1: {
+              uint32_t value_base =
+                  thread->rand.Next() % thread->shared->UNKNOWN_SENTINEL;
+              char value[100];
+              size_t sz = GenerateValue(value_base, value, sizeof(value));
+              Slice v(value, sz);
+              if (op == 0) {
+                s = txn->Put(cfh, keys.back(), v);
+              } else {
+                s = txn->Merge(cfh, keys.back(), v);
+              }
+              break;
+            }
+            case 2:
+              s = txn->Delete(cfh, keys.back());
+              break;
+            default:
+              assert(false);
+          }
+          if (!s.ok()) {
+            fprintf(stderr, "Transaction put: %s\n", s.ToString().c_str());
+            std::terminate();
+          }
+        }
+      }
+#endif
+    }
+
+    if (!FLAGS_use_txn) {
+      db_->MultiGet(read_opts, cfh, num_keys, keys.data(), values.data(),
+                    statuses.data());
+    } else {
+#ifndef ROCKSDB_LITE
+      txn->MultiGet(read_opts, cfh, num_keys, keys.data(), values.data(),
+                    statuses.data());
+      RollbackTxn(txn);
+#endif
     }
-    db_->MultiGet(read_opts, cfh, num_keys, keys.data(), values.data(),
-                  statuses.data());
+
     for (const auto& s : statuses) {
       if (s.ok()) {
         // found case