db_stress: Rebuild expected state from DB when trace is truncated in blackbox crash tests

Summary:
In blackbox crash tests, db_stress can be killed at any point — including
while it is actively appending records to the trace file. This leaves the
recovered DB in a consistent state (thanks to WAL recovery), but the trace
file shorter than expected: fewer operations were recorded than the DB's
recovered sequence number implies.

Previously, when FinishInitDb detected this situation it called
FileExpectedStateManager::Restore() which ended up returning
Status::Corruption("Trace ended before replaying all expected write ops").
db_stress treated that as a hard failure and called exit(1), causing spurious
test failures that had nothing to do with actual data corruption.

This PR fixes the two-part problem:

1. expected_state.cc – FileExpectedStateManager::Restore():
   When the trace replay finishes (either because Prepare() signals the trace
   is already at EOF, or because Next() exhausts all records) but the handler
   has not replayed all expected write operations, return Status::Incomplete()
   instead of Status::Corruption(). Incomplete means "the trace ran out early,
   which is expected in a blackbox crash"; Corruption is reserved for actual
   data integrity problems.

2. db_stress_test_base.cc – StressTest::FinishInitDb() + new
   RebuildExpectedStateFromDb():
   When Restore() returns Incomplete, fall back to scanning the recovered DB
   directly and rebuilding the expected state from its actual contents (the
   same approach used when there is no history at all). If the rebuild scan
   fails, or if Restore() returns any other non-OK status, still exit(1).

The new RebuildExpectedStateFromDb() method iterates every column family,
calls ClearColumnFamily() + SyncPut() for each key found, and validates wide
column consistency and key-range sanity along the way.

Test Plan:
Covered by existing blackbox crash test modes (blackbox, cf_consistency, etc.)
which can now survive a truncated trace without false failures.

Reviewers:

Subscribers:

Tasks:

Tags:

Author: xingbowang

db_stress_tool/db_stress_test_base.cc

@@ -466,15 +466,24 @@ void StressTest::FinishInitDb(SharedState* shared) {
   }
 
   if (shared->HasHistory()) {
-    // The way it works right now is, if there's any history, that means the
-    // previous run mutating the DB had all its operations traced, in which case
-    // we should always be able to `Restore()` the expected values to match the
-    // `db_`'s current seqno.
+    // Normally startup replays trace records to catch the expected state up to
+    // `db_`'s recovered seqno. Blackbox crash tests can kill db_stress while
+    // the trace file is being appended, so treat a truncated replay suffix as a
+    // signal to rebuild the latest expected state directly from the DB.
     Status s = shared->Restore(db_);
     if (!s.ok()) {
-      fprintf(stderr, "Error restoring historical expected values: %s\n",
-              s.ToString().c_str());
-      exit(1);
+      if (s.IsIncomplete()) {
+        fprintf(stdout,
+                "Historical expected-state restore was incomplete during "
+                "crash recovery; rebuilding from recovered DB: %s\n",
+                s.ToString().c_str());
+        s = RebuildExpectedStateFromDb(shared);
+      }
+      if (!s.ok()) {
+        fprintf(stderr, "Error restoring historical expected values: %s\n",
+                s.ToString().c_str());
+        exit(1);
+      }
     }
   }
   if (FLAGS_use_txn && !FLAGS_use_optimistic_txn) {
@@ -499,6 +508,60 @@ void StressTest::FinishInitDb(SharedState* shared) {
   }
 }
 
+Status StressTest::RebuildExpectedStateFromDb(SharedState* shared) {
+  assert(shared);
+
+  ReadOptions read_opts;
+  read_opts.total_order_seek = true;
+
+  for (size_t cf_idx = 0; cf_idx < column_families_.size(); ++cf_idx) {
+    const int cf = static_cast<int>(cf_idx);
+    ColumnFamilyHandle* const cfh = column_families_[cf_idx];
+    assert(cfh);
+
+    shared->LockColumnFamily(cf);
+    shared->ClearColumnFamily(cf);
+
+    std::unique_ptr<Iterator> iter(db_->NewIterator(read_opts, cfh));
+    for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
+      if (!VerifyWideColumns(iter->value(), iter->columns())) {
+        Status s = Status::Corruption(
+            "Wide columns inconsistent while rebuilding expected state",
+            cfh->GetName());
+        shared->UnlockColumnFamily(cf);
+        return s;
+      }
+
+      uint64_t key_id = 0;
+      if (!GetIntVal(iter->key().ToString(), &key_id)) {
+        Status s = Status::Corruption(
+            "Unable to parse key while rebuilding expected state",
+            iter->key().ToString(/* hex */ true));
+        shared->UnlockColumnFamily(cf);
+        return s;
+      }
+      if (key_id >= static_cast<uint64_t>(shared->GetMaxKey())) {
+        Status s = Status::Corruption(
+            "Out-of-range key while rebuilding expected state",
+            std::to_string(key_id));
+        shared->UnlockColumnFamily(cf);
+        return s;
+      }
+
+      shared->SyncPut(cf, static_cast<int64_t>(key_id),
+                      GetValueBase(iter->value()));
+    }
+
+    Status s = iter->status();
+    shared->UnlockColumnFamily(cf);
+    if (!s.ok()) {
+      return s;
+    }
+  }
+
+  return Status::OK();
+}
+
 void StressTest::TrackExpectedState(SharedState* shared) {
   // When data loss is simulated, recovery from potential data loss is a prefix
   // recovery that requires tracing

db_stress_tool/db_stress_test_base.h

@@ -144,6 +144,9 @@ class StressTest {
   void PreloadDbAndReopenAsReadOnly(int64_t number_of_keys,
                                     SharedState* shared);
 
+  // Rebuilds the latest expected state from the recovered DB contents.
+  Status RebuildExpectedStateFromDb(SharedState* shared);
+
   Status SetOptions(ThreadState* thread);
 
   // For transactionsDB, there can be txns prepared but not yet committeed

db_stress_tool/expected_state.cc

@@ -770,10 +770,15 @@ Status FileExpectedStateManager::Restore(DB* db) {
                                  std::move(trace_reader), &replayer);
     }
 
+    bool reached_trace_end = false;
     if (s.ok()) {
       s = replayer->Prepare();
+      if (s.IsIncomplete()) {
+        reached_trace_end = true;
+        s = Status::OK();
+      }
     }
-    for (; s.ok();) {
+    for (; s.ok() && !reached_trace_end;) {
       std::unique_ptr<TraceRecord> record;
       s = replayer->Next(&record);
       if (!s.ok()) {
@@ -789,13 +794,18 @@ Status FileExpectedStateManager::Restore(DB* db) {
           // trace records.
           s = Status::OK();
         }
+        reached_trace_end = true;
         break;
       }
       std::unique_ptr<TraceRecordResult> res;
       s = record->Accept(handler.get(), &res);
     }
     if (s.ok() && !handler->IsDone()) {
-      s = Status::Corruption(
+      // Blackbox crash tests can terminate db_stress while it is appending the
+      // trace, leaving a valid recovered DB but a shorter trace suffix. Treat
+      // that as an incomplete restore so startup can rebuild expected state
+      // from the live DB instead of reporting a false corruption.
+      s = Status::Incomplete(
           "Trace ended before replaying all expected write ops",
           std::to_string(handler->NumWriteOps()) + " < " +
               std::to_string(seqno - saved_seqno_));