From acf77e1bfee9ebb0867ac277927ed2e37276c493 Mon Sep 17 00:00:00 2001
From: Peter Dillinger <peterd@meta.com>
Date: Tue, 30 Jan 2024 16:16:04 -0800
Subject: [PATCH] Fix possible crash test segfault in
 FileExpectedStateManager::Restore() (#12314)

Summary:
`replayer` could be `nullptr` if `!s.ok()` from an earlier failure. Also consider status returned from `record->Accept()`.

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

Test Plan: blackbox_crash_test run

Reviewed By: hx235

Differential Revision: D53241506

Pulled By: pdillinger

fbshipit-source-id: fd330417c23391ca819c3ee0f69e4156d81934dc
---
 db_stress_tool/expected_state.cc | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/db_stress_tool/expected_state.cc b/db_stress_tool/expected_state.cc
index ce02ce49c79..428aa09f9a4 100644
--- a/db_stress_tool/expected_state.cc
+++ b/db_stress_tool/expected_state.cc
@@ -660,26 +660,26 @@ Status FileExpectedStateManager::Restore(DB* db) {
     if (s.ok()) {
       s = replayer->Prepare();
     }
-    for (;;) {
+    for (; s.ok();) {
       std::unique_ptr<TraceRecord> record;
       s = replayer->Next(&record);
       if (!s.ok()) {
+        if (s.IsCorruption() && handler->IsDone()) {
+          // There could be a corruption reading the tail record of the trace
+          // due to `db_stress` crashing while writing it. It shouldn't matter
+          // as long as we already found all the write ops we need to catch up
+          // the expected state.
+          s = Status::OK();
+        }
+        if (s.IsIncomplete()) {
+          // OK because `Status::Incomplete` is expected upon finishing all the
+          // trace records.
+          s = Status::OK();
+        }
         break;
       }
       std::unique_ptr<TraceRecordResult> res;
-      record->Accept(handler.get(), &res);
-    }
-    if (s.IsCorruption() && handler->IsDone()) {
-      // There could be a corruption reading the tail record of the trace due to
-      // `db_stress` crashing while writing it. It shouldn't matter as long as
-      // we already found all the write ops we need to catch up the expected
-      // state.
-      s = Status::OK();
-    }
-    if (s.IsIncomplete()) {
-      // OK because `Status::Incomplete` is expected upon finishing all the
-      // trace records.
-      s = Status::OK();
+      s = record->Accept(handler.get(), &res);
     }
   }