Fix stress test key validation with UDT

db_stress_tool/db_stress_common.h

@@ -36,6 +36,7 @@
 #include <thread>
 
 #include "db/db_impl/db_impl.h"
+#include "db/dbformat.h"
 #include "db/version_set.h"
 #include "db/wide/wide_columns_helper.h"
 #include "db_stress_tool/db_stress_env_wrapper.h"
@@ -670,6 +671,23 @@ inline std::string Key(int64_t val) {
   return key;
 }
 
+// Helper function to strip user-defined timestamp from a key if timestamps
+// are enabled. This is used when comparing or parsing iterator keys that
+// include the timestamp suffix.
+inline Slice MaybeStripTimestamp(const Slice& key) {
+  if (FLAGS_user_timestamp_size > 0) {
+    return StripTimestampFromUserKey(key, FLAGS_user_timestamp_size);
+  }
+  return key;
+}
+
+// Helper function to get the user key portion as a string, stripping the
+// timestamp if user-defined timestamps are enabled. This is suitable for
+// passing to GetIntVal() when parsing iterator keys.
+inline std::string GetUserKeyForParsing(const Slice& key) {
+  return MaybeStripTimestamp(key).ToString();
+}
+
 // Given a string key, map it to an index into the expected values buffer
 inline bool GetIntVal(std::string big_endian_key, uint64_t* key_p) {
   size_t size_key = big_endian_key.size();

db_stress_tool/db_stress_test_base.cc

@@ -582,7 +582,7 @@ Status StressTest::AssertSame(DB* db, ColumnFamilyHandle* cf,
         new std::vector<bool>(FLAGS_max_key));
     for (iterator->SeekToFirst(); iterator->Valid(); iterator->Next()) {
       uint64_t key_val;
-      if (GetIntVal(iterator->key().ToString(), &key_val)) {
+      if (GetIntVal(GetUserKeyForParsing(iterator->key()), &key_val)) {
         (*tmp_bitvec.get())[key_val] = true;
       }
     }
@@ -839,7 +839,7 @@ void StressTest::ProcessRecoveredPreparedTxnsHelper(Transaction* txn,
         txn->GetWriteBatch()->NewIterator(column_families_[i]));
     for (wbwi_iter->SeekToFirst(); wbwi_iter->Valid(); wbwi_iter->Next()) {
       uint64_t key_val;
-      if (GetIntVal(wbwi_iter->Entry().key.ToString(), &key_val)) {
+      if (GetIntVal(GetUserKeyForParsing(wbwi_iter->Entry().key), &key_val)) {
         shared->SyncPendingPut(static_cast<int>(i) /* cf_idx */, key_val);
       }
     }
@@ -3203,7 +3203,7 @@ void StressTest::TestAcquireSnapshot(ThreadState* thread,
     std::unique_ptr<Iterator> iterator(db_->NewIterator(ropt));
     for (iterator->SeekToFirst(); iterator->Valid(); iterator->Next()) {
       uint64_t key_val;
-      if (GetIntVal(iterator->key().ToString(), &key_val)) {
+      if (GetIntVal(GetUserKeyForParsing(iterator->key()), &key_val)) {
         (*key_vec)[key_val] = true;
       }
     }

db_stress_tool/no_batched_ops_stress.cc

@@ -123,7 +123,7 @@ class NonBatchedOpsStressTest : public StressTest {
           std::string from_db;
 
           if (iter->Valid()) {
-            const int diff = iter->key().compare(k);
+            const int diff = MaybeStripTimestamp(iter->key()).compare(k);
 
             if (diff > 0) {
               s = Status::NotFound();
@@ -2616,7 +2616,7 @@ class NonBatchedOpsStressTest : public StressTest {
       }
 
       // iter is valid, the range (last_key, current key) was skipped
-      GetIntVal(iter->key().ToString(), &curr);
+      GetIntVal(GetUserKeyForParsing(iter->key()), &curr);
       if (static_cast<int64_t>(curr) <= last_key) {
         thread->shared->SetVerificationFailure();
         fprintf(stderr,
@@ -2688,7 +2688,7 @@ class NonBatchedOpsStressTest : public StressTest {
       }
 
       // the range (current key, last key) was skipped
-      GetIntVal(iter->key().ToString(), &curr);
+      GetIntVal(GetUserKeyForParsing(iter->key()), &curr);
       if (last_key <= static_cast<int64_t>(curr)) {
         thread->shared->SetVerificationFailure();
         fprintf(stderr,
@@ -2763,7 +2763,7 @@ class NonBatchedOpsStressTest : public StressTest {
           return Status::OK();
         }
       } else if (iter->Valid()) {
-        GetIntVal(iter->key().ToString(), &curr);
+        GetIntVal(GetUserKeyForParsing(iter->key()), &curr);
         if (static_cast<int64_t>(curr) < mid) {
           thread->shared->SetVerificationFailure();
           fprintf(stderr,
@@ -2787,7 +2787,7 @@ class NonBatchedOpsStressTest : public StressTest {
           return Status::OK();
         }
       } else if (iter->Valid()) {
-        GetIntVal(iter->key().ToString(), &curr);
+        GetIntVal(GetUserKeyForParsing(iter->key()), &curr);
         if (mid < static_cast<int64_t>(curr)) {
           thread->shared->SetVerificationFailure();
           fprintf(stderr,
@@ -2819,7 +2819,7 @@ class NonBatchedOpsStressTest : public StressTest {
         return Status::OK();
       }
 
-      GetIntVal(iter->key().ToString(), &curr);
+      GetIntVal(GetUserKeyForParsing(iter->key()), &curr);
       if (static_cast<int64_t>(curr) < lb) {
         iter->Next();
         op_logs += "N";
@@ -2859,7 +2859,7 @@ class NonBatchedOpsStressTest : public StressTest {
             break;
           }
           uint64_t next = 0;
-          GetIntVal(iter->key().ToString(), &next);
+          GetIntVal(GetUserKeyForParsing(iter->key()), &next);
           if (next <= curr) {
             thread->shared->SetVerificationFailure();
             fprintf(stderr,
@@ -2884,7 +2884,7 @@ class NonBatchedOpsStressTest : public StressTest {
             break;
           }
           uint64_t prev = 0;
-          GetIntVal(iter->key().ToString(), &prev);
+          GetIntVal(GetUserKeyForParsing(iter->key()), &prev);
           if (curr <= prev) {
             thread->shared->SetVerificationFailure();
             fprintf(stderr,