Make transaction name conflict check more robust (#12895)

Summary:
The `PessimisticTransaction::SetName()` code checks for an existing txn of the given name before registering the new txn. However, this is not atomic, which could result in a race condition if two txns try to register with the same name. Both might succeed and lead to unpredictable behavior. This PR makes the test and set atomic.

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

Reviewed By: pdillinger

Differential Revision: D60460482

Pulled By: anand1976

fbshipit-source-id: e8afeb2356e1b8f4e8df785cb73532739f82579d
This commit is contained in:
anand76 2024-07-30 12:31:02 -07:00 committed by Facebook GitHub Bot
parent 9058fd037c
commit 55877d8893
4 changed files with 17 additions and 7 deletions

View File

@ -0,0 +1 @@
Fix a race condition in pessimistic transactions that could allow multiple transactions with the same name to be registered simultaneously, resulting in a crash or other unpredictable behavior.

View File

@ -1247,14 +1247,15 @@ Status PessimisticTransaction::SetName(const TransactionName& name) {
if (txn_state_ == STARTED) { if (txn_state_ == STARTED) {
if (name_.length()) { if (name_.length()) {
s = Status::InvalidArgument("Transaction has already been named."); s = Status::InvalidArgument("Transaction has already been named.");
} else if (txn_db_impl_->GetTransactionByName(name) != nullptr) {
s = Status::InvalidArgument("Transaction name must be unique.");
} else if (name.length() < 1 || name.length() > 512) { } else if (name.length() < 1 || name.length() > 512) {
s = Status::InvalidArgument( s = Status::InvalidArgument(
"Transaction name length must be between 1 and 512 chars."); "Transaction name length must be between 1 and 512 chars.");
} else { } else {
name_ = name; name_ = name;
txn_db_impl_->RegisterTransaction(this); s = txn_db_impl_->RegisterTransaction(this);
if (!s.ok()) {
name_.clear();
}
} }
} else { } else {
s = Status::InvalidArgument("Transaction is beyond state for naming."); s = Status::InvalidArgument("Transaction is beyond state for naming.");

View File

@ -723,6 +723,11 @@ void PessimisticTransactionDB::ReinitializeTransaction(
Transaction* PessimisticTransactionDB::GetTransactionByName( Transaction* PessimisticTransactionDB::GetTransactionByName(
const TransactionName& name) { const TransactionName& name) {
std::lock_guard<std::mutex> lock(name_map_mutex_); std::lock_guard<std::mutex> lock(name_map_mutex_);
return GetTransactionByNameLocked(name);
}
Transaction* PessimisticTransactionDB::GetTransactionByNameLocked(
const TransactionName& name) {
auto it = transactions_.find(name); auto it = transactions_.find(name);
if (it == transactions_.end()) { if (it == transactions_.end()) {
return nullptr; return nullptr;
@ -755,13 +760,15 @@ void PessimisticTransactionDB::SetDeadlockInfoBufferSize(uint32_t target_size) {
lock_manager_->Resize(target_size); lock_manager_->Resize(target_size);
} }
void PessimisticTransactionDB::RegisterTransaction(Transaction* txn) { Status PessimisticTransactionDB::RegisterTransaction(Transaction* txn) {
assert(txn); assert(txn);
assert(txn->GetName().length() > 0); assert(txn->GetName().length() > 0);
assert(GetTransactionByName(txn->GetName()) == nullptr);
assert(txn->GetState() == Transaction::STARTED); assert(txn->GetState() == Transaction::STARTED);
std::lock_guard<std::mutex> lock(name_map_mutex_); std::lock_guard<std::mutex> lock(name_map_mutex_);
transactions_[txn->GetName()] = txn; if (!transactions_.insert({txn->GetName(), txn}).second) {
return Status::InvalidArgument("Duplicate txn name " + txn->GetName());
}
return Status::OK();
} }
void PessimisticTransactionDB::UnregisterTransaction(Transaction* txn) { void PessimisticTransactionDB::UnregisterTransaction(Transaction* txn) {

View File

@ -173,7 +173,7 @@ class PessimisticTransactionDB : public TransactionDB {
Transaction* GetTransactionByName(const TransactionName& name) override; Transaction* GetTransactionByName(const TransactionName& name) override;
void RegisterTransaction(Transaction* txn); Status RegisterTransaction(Transaction* txn);
void UnregisterTransaction(Transaction* txn); void UnregisterTransaction(Transaction* txn);
// not thread safe. current use case is during recovery (single thread) // not thread safe. current use case is during recovery (single thread)
@ -239,6 +239,7 @@ class PessimisticTransactionDB : public TransactionDB {
friend class WriteUnpreparedTransactionTest_MarkLogWithPrepSection_Test; friend class WriteUnpreparedTransactionTest_MarkLogWithPrepSection_Test;
Transaction* BeginInternalTransaction(const WriteOptions& options); Transaction* BeginInternalTransaction(const WriteOptions& options);
Transaction* GetTransactionByNameLocked(const TransactionName& name);
std::shared_ptr<LockManager> lock_manager_; std::shared_ptr<LockManager> lock_manager_;