diff --git a/utilities/transactions/write_unprepared_transaction_test.cc b/utilities/transactions/write_unprepared_transaction_test.cc index 92835bba89..fba67e57a4 100644 --- a/utilities/transactions/write_unprepared_transaction_test.cc +++ b/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 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) { diff --git a/utilities/transactions/write_unprepared_txn.cc b/utilities/transactions/write_unprepared_txn.cc index 444c1c9b63..32c62ae69b 100644 --- a/utilities/transactions/write_unprepared_txn.cc +++ b/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