From a6ee297ac94e7237879588b52356033ec2dd0af7 Mon Sep 17 00:00:00 2001
From: Levi Tamasi <ltamasi@meta.com>
Date: Fri, 8 Nov 2024 15:37:19 -0800
Subject: [PATCH] Save the key before calling PrepareValue() in the stress
 tests (#13129)

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

The `PrepareValue()` call on an iterator can fail, for example due to our stress tests' read fault injection. Such a failure invalidates the iterator, which makes it illegal to call methods like `key()` on it and leads to assertion violations. The patch fixes this by saving the key before calling `PrepareValue()`, so we can still print it for debugging purposes in case the call fails.

Reviewed By: jowlyzhang

Differential Revision: D65689225

fbshipit-source-id: c2bf298366def0ba3b3c089ee58e28609ecdfab4
---
 db_stress_tool/batched_ops_stress.cc    | 7 ++++++-
 db_stress_tool/db_stress_test_base.cc   | 6 +++++-
 db_stress_tool/no_batched_ops_stress.cc | 6 +++++-
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/db_stress_tool/batched_ops_stress.cc b/db_stress_tool/batched_ops_stress.cc
index d624af9c66..5ecef42eee 100644
--- a/db_stress_tool/batched_ops_stress.cc
+++ b/db_stress_tool/batched_ops_stress.cc
@@ -615,10 +615,15 @@ class BatchedOpsStressTest : public StressTest {
                iters[i]->key().starts_with(prefix_slices[i]));
 
         if (ro_copies[i].allow_unprepared_value) {
+          // Save key in case PrepareValue fails and invalidates the iterator
+          const std::string prepare_value_key =
+              iters[i]->key().ToString(/* hex */ true);
+
           if (!iters[i]->PrepareValue()) {
             fprintf(stderr,
                     "prefix scan error: PrepareValue failed for key %s\n",
-                    iters[i]->key().ToString(/* hex */ true).c_str());
+                    prepare_value_key.c_str());
+            continue;
           }
         }
 
diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc
index 2a7ea050ce..5dd4fe8782 100644
--- a/db_stress_tool/db_stress_test_base.cc
+++ b/db_stress_tool/db_stress_test_base.cc
@@ -2001,13 +2001,17 @@ void StressTest::VerifyIterator(
 
   if (!*diverged && iter->Valid()) {
     if (ro.allow_unprepared_value) {
+      // Save key in case PrepareValue fails and invalidates the iterator
+      const std::string prepare_value_key =
+          iter->key().ToString(/* hex */ true);
+
       if (!iter->PrepareValue()) {
         fprintf(
             stderr,
             "Iterator failed to prepare value for key %s %s under specified "
             "iterator ReadOptions: %s (Empty string or missing field indicates "
             "default option or value is used)\n",
-            iter->key().ToString(true).c_str(), op_logs.c_str(),
+            prepare_value_key.c_str(), op_logs.c_str(),
             read_opt_oss.str().c_str());
         *diverged = true;
       }
diff --git a/db_stress_tool/no_batched_ops_stress.cc b/db_stress_tool/no_batched_ops_stress.cc
index 62bf3ca859..9485bca777 100644
--- a/db_stress_tool/no_batched_ops_stress.cc
+++ b/db_stress_tool/no_batched_ops_stress.cc
@@ -2301,12 +2301,16 @@ class NonBatchedOpsStressTest : public StressTest {
       assert(iter->Valid());
 
       if (ro.allow_unprepared_value) {
+        // Save key in case PrepareValue fails and invalidates the iterator
+        const std::string prepare_value_key =
+            iter->key().ToString(/* hex */ true);
+
         if (!iter->PrepareValue()) {
           shared->SetVerificationFailure();
 
           fprintf(stderr,
                   "Verification failed for key %s: failed to prepare value\n",
-                  Slice(iter->key()).ToString(/* hex */ true).c_str());
+                  prepare_value_key.c_str());
           fprintf(stderr, "Column family: %s, op_logs: %s\n",
                   cfh->GetName().c_str(), op_logs.c_str());