Fix WriteUnprepared transaction rollback after commit write failure (#14920)

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

When a WriteUnprepared transaction's commit write fails while two_write_queues
is enabled (e.g. due to an I/O error), CommitInternal() unconditionally cleared
unprep_seqs_ and the flushed/unflushed save points before returning the error.
That left the transaction with no record of its prepared sequence numbers, so a
subsequent Rollback() -- the natural way for a caller to recover from a failed
commit -- constructed a WriteUnpreparedCommitEntryPreReleaseCallback over an
empty unprep_seqs_ and tripped assert(unprep_seqs.size() > 0) in its
constructor. In release builds this would instead corrupt commit-cache
accounting rather than aborting.

Fix: only clear unprep_seqs_ and the save points on the commit success path. On
the two-write-queues commit-write failure path we still remove the prepared
entry that AddPreparedCallback may have added for commit_batch_seq (that seq is
never stored in unprep_seqs_, so a later Rollback cannot clean it up), but
unprep_seqs_ is otherwise preserved so the transaction can still be rolled back
once the caller resumes the DB from the write error.

Originally authored by Devmate trigger

Reviewed By: xingbowang

Differential Revision: D110271443

fbshipit-source-id: 7d991e01ddc74ccd170cbc25f6af160c5f3d5663

utilities/transactions/write_unprepared_transaction_test.cc

@@ -1017,6 +1017,58 @@ TEST_P(WriteUnpreparedTransactionTest, RangeTombstoneOwnDeletionsAndRollback) {
   }
 }
 
+TEST_P(WriteUnpreparedTransactionTest,
+       RollbackPreparedAfterCommitWriteFailure) {
+  // Verify that a WriteUnprepared transaction can be rolled back after a
+  // commit write failure. Previously, CommitInternal() would clear
+  // unprep_seqs_ even on failure, making subsequent Rollback() hit
+  // assertion `unprep_seqs.size() > 0`.
+  options.max_bgerror_resume_count = 0;
+  ASSERT_OK(ReOpen());
+
+  // Use fault injection to cause commit write to fail.
+  fault_fs->SetFileTypesExcludedFromFaultInjection({FileType::kInfoLogFile});
+
+  WriteOptions woptions;
+  TransactionOptions txn_options;
+  // Set a low flush threshold to trigger unprepared batches being written to
+  // DB, populating unprep_seqs_.
+  txn_options.write_batch_flush_threshold = 1;
+  std::unique_ptr<Transaction> txn(db->BeginTransaction(woptions, txn_options));
+  ASSERT_NE(nullptr, txn);
+  ASSERT_OK(txn->SetName("wup_commit_fail_rollback"));
+  ASSERT_OK(txn->Put(Slice("key1"), Slice("value1")));
+  // The flush threshold of 1 should cause the batch to be flushed as
+  // unprepared, populating unprep_seqs_.
+  ASSERT_OK(txn->Put(Slice("key2"), Slice("value2")));
+  ASSERT_OK(txn->Prepare());
+
+  // Inject a write error so that Commit() fails.
+  fault_fs->SetThreadLocalErrorContext(FaultInjectionIOType::kWrite,
+                                       /*seed=*/0, /*one_in=*/1,
+                                       /*retryable=*/true,
+                                       /*has_data_loss=*/false);
+  fault_fs->EnableThreadLocalErrorInjection(FaultInjectionIOType::kWrite);
+  const Status commit_s = txn->Commit();
+  fault_fs->DisableThreadLocalErrorInjection(FaultInjectionIOType::kWrite);
+
+  ASSERT_NOK(commit_s);
+  ASSERT_TRUE(FaultInjectionTestFS::IsInjectedError(commit_s))
+      << commit_s.ToString();
+  fault_fs->GetAndResetInjectedThreadLocalErrorCount(
+      FaultInjectionIOType::kWrite);
+
+  // Resume the DB to clear the error state.
+  ASSERT_OK(db->Resume());
+
+  // Rollback should succeed without hitting the assertion failure.
+  ASSERT_OK(txn->Rollback());
+  txn.reset();
+
+  // Verify the DB is usable after rollback.
+  ASSERT_OK(ReOpenNoDelete());
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

utilities/transactions/write_unprepared_txn.cc

@@ -608,13 +608,18 @@ Status WriteUnpreparedTxn::CommitInternal() {
       for (const auto& seq : unprep_seqs_) {
         wpt_db_->RemovePrepared(seq.first, seq.second);
       }
-    }
-    if (UNLIKELY(!do_one_write)) {
+      unprep_seqs_.clear();
+      flushed_save_points_.reset(nullptr);
+      unflushed_save_points_.reset(nullptr);
+    } else if (UNLIKELY(!do_one_write)) {
+      // The commit write failed after add_prepared_callback may have added a
+      // prepared entry for commit_batch_seq. Since commit_batch_seq is never
+      // stored in unprep_seqs_ on this path, a follow-up Rollback() cannot
+      // remove it, so clean it up here. unprep_seqs_ is otherwise preserved so
+      // the transaction can still be rolled back after the caller recovers
+      // from the write error.
       wpt_db_->RemovePrepared(commit_batch_seq, commit_batch_cnt);
     }
-    unprep_seqs_.clear();
-    flushed_save_points_.reset(nullptr);
-    unflushed_save_points_.reset(nullptr);
     return s;
   }  // else do the 2nd write to publish seq