From a8a4ed52a4f4ab57d214ed4984306ea89da32f3f Mon Sep 17 00:00:00 2001
From: Yanqin Jin <yanqin@fb.com>
Date: Thu, 17 Nov 2022 20:43:50 -0800
Subject: [PATCH] Test Merge with timestamps in stress test (#10948)

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

Test Plan: make crash_test_with_ts

Reviewed By: ltamasi

Differential Revision: D41390854

Pulled By: riversand963

fbshipit-source-id: 599e114da8e2b2bbff5628fb8c67fa0393a31c05
---
 db_stress_tool/db_stress_test_base.cc   | 19 ++++++++++++-------
 db_stress_tool/no_batched_ops_stress.cc |  9 +++++++--
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc
index ccf59a14f..627fba88d 100644
--- a/db_stress_tool/db_stress_test_base.cc
+++ b/db_stress_tool/db_stress_test_base.cc
@@ -517,9 +517,18 @@ void StressTest::PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
 
       shared->Put(cf_idx, k, value_base, true /* pending */);
 
+      std::string ts;
+      if (FLAGS_user_timestamp_size > 0) {
+        ts = GetNowNanos();
+      }
+
       if (FLAGS_use_merge) {
         if (!FLAGS_use_txn) {
-          s = db_->Merge(write_opts, cfh, key, v);
+          if (FLAGS_user_timestamp_size > 0) {
+            s = db_->Merge(write_opts, cfh, key, ts, v);
+          } else {
+            s = db_->Merge(write_opts, cfh, key, v);
+          }
         } else {
 #ifndef ROCKSDB_LITE
           Transaction* txn;
@@ -538,7 +547,6 @@ void StressTest::PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
       } else {
         if (!FLAGS_use_txn) {
           if (FLAGS_user_timestamp_size > 0) {
-            const std::string ts = GetNowNanos();
             s = db_->Put(write_opts, cfh, key, ts, v);
           } else {
             s = db_->Put(write_opts, cfh, key, v);
@@ -2960,7 +2968,8 @@ void StressTest::MaybeUseOlderTimestampForRangeScan(ThreadState* thread,
   ts_slice = ts_str;
   read_opts.timestamp = &ts_slice;
 
-  if (!thread->rand.OneInOpt(3)) {
+  // TODO (yanqin): support Merge with iter_start_ts
+  if (!thread->rand.OneInOpt(3) || FLAGS_use_merge || FLAGS_use_full_merge_v1) {
     return;
   }
 
@@ -2981,10 +2990,6 @@ void CheckAndSetOptionsForUserTimestamp(Options& options) {
             static_cast<int>(cmp->timestamp_size()));
     exit(1);
   }
-  if (FLAGS_use_merge || FLAGS_use_full_merge_v1) {
-    fprintf(stderr, "Merge does not support timestamp yet.\n");
-    exit(1);
-  }
   if (FLAGS_use_txn) {
     fprintf(stderr, "TransactionDB does not support timestamp yet.\n");
     exit(1);
diff --git a/db_stress_tool/no_batched_ops_stress.cc b/db_stress_tool/no_batched_ops_stress.cc
index 67b661d34..bf01b788f 100644
--- a/db_stress_tool/no_batched_ops_stress.cc
+++ b/db_stress_tool/no_batched_ops_stress.cc
@@ -61,7 +61,8 @@ class NonBatchedOpsStressTest : public StressTest {
           static_cast<int>(VerificationMethod::kNumberOfMethods);
 
       const VerificationMethod method =
-          static_cast<VerificationMethod>(thread->rand.Uniform(num_methods));
+          static_cast<VerificationMethod>(thread->rand.Uniform(
+              (FLAGS_user_timestamp_size > 0) ? num_methods - 1 : num_methods));
 
       if (method == VerificationMethod::kIterator) {
         std::unique_ptr<Iterator> iter(
@@ -805,7 +806,11 @@ class NonBatchedOpsStressTest : public StressTest {
 
     if (FLAGS_use_merge) {
       if (!FLAGS_use_txn) {
-        s = db_->Merge(write_opts, cfh, k, v);
+        if (FLAGS_user_timestamp_size == 0) {
+          s = db_->Merge(write_opts, cfh, k, v);
+        } else {
+          s = db_->Merge(write_opts, cfh, k, write_ts, v);
+        }
       } else {
 #ifndef ROCKSDB_LITE
         Transaction* txn;