From 55877d88935b8e98c7bb280eb0ec280a27fb7906 Mon Sep 17 00:00:00 2001
From: anand76 <anand1976@users.noreply.github.com>
Date: Tue, 30 Jul 2024 12:31:02 -0700
Subject: [PATCH] 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
---
 .../bug_fixes/prevent_duplicate_txn_name.md         |  1 +
 utilities/transactions/pessimistic_transaction.cc   |  7 ++++---
 .../transactions/pessimistic_transaction_db.cc      | 13 ++++++++++---
 utilities/transactions/pessimistic_transaction_db.h |  3 ++-
 4 files changed, 17 insertions(+), 7 deletions(-)
 create mode 100644 unreleased_history/bug_fixes/prevent_duplicate_txn_name.md

diff --git a/unreleased_history/bug_fixes/prevent_duplicate_txn_name.md b/unreleased_history/bug_fixes/prevent_duplicate_txn_name.md
new file mode 100644
index 0000000000..152869d7d2
--- /dev/null
+++ b/unreleased_history/bug_fixes/prevent_duplicate_txn_name.md
@@ -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.
diff --git a/utilities/transactions/pessimistic_transaction.cc b/utilities/transactions/pessimistic_transaction.cc
index b572014dad..45f84ef9f3 100644
--- a/utilities/transactions/pessimistic_transaction.cc
+++ b/utilities/transactions/pessimistic_transaction.cc
@@ -1247,14 +1247,15 @@ Status PessimisticTransaction::SetName(const TransactionName& name) {
   if (txn_state_ == STARTED) {
     if (name_.length()) {
       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) {
       s = Status::InvalidArgument(
           "Transaction name length must be between 1 and 512 chars.");
     } else {
       name_ = name;
-      txn_db_impl_->RegisterTransaction(this);
+      s = txn_db_impl_->RegisterTransaction(this);
+      if (!s.ok()) {
+        name_.clear();
+      }
     }
   } else {
     s = Status::InvalidArgument("Transaction is beyond state for naming.");
diff --git a/utilities/transactions/pessimistic_transaction_db.cc b/utilities/transactions/pessimistic_transaction_db.cc
index 7aca2bb523..ea38e17eea 100644
--- a/utilities/transactions/pessimistic_transaction_db.cc
+++ b/utilities/transactions/pessimistic_transaction_db.cc
@@ -723,6 +723,11 @@ void PessimisticTransactionDB::ReinitializeTransaction(
 Transaction* PessimisticTransactionDB::GetTransactionByName(
     const TransactionName& name) {
   std::lock_guard<std::mutex> lock(name_map_mutex_);
+  return GetTransactionByNameLocked(name);
+}
+
+Transaction* PessimisticTransactionDB::GetTransactionByNameLocked(
+    const TransactionName& name) {
   auto it = transactions_.find(name);
   if (it == transactions_.end()) {
     return nullptr;
@@ -755,13 +760,15 @@ void PessimisticTransactionDB::SetDeadlockInfoBufferSize(uint32_t target_size) {
   lock_manager_->Resize(target_size);
 }
 
-void PessimisticTransactionDB::RegisterTransaction(Transaction* txn) {
+Status PessimisticTransactionDB::RegisterTransaction(Transaction* txn) {
   assert(txn);
   assert(txn->GetName().length() > 0);
-  assert(GetTransactionByName(txn->GetName()) == nullptr);
   assert(txn->GetState() == Transaction::STARTED);
   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) {
diff --git a/utilities/transactions/pessimistic_transaction_db.h b/utilities/transactions/pessimistic_transaction_db.h
index 77670e7d4a..bec29182ab 100644
--- a/utilities/transactions/pessimistic_transaction_db.h
+++ b/utilities/transactions/pessimistic_transaction_db.h
@@ -173,7 +173,7 @@ class PessimisticTransactionDB : public TransactionDB {
 
   Transaction* GetTransactionByName(const TransactionName& name) override;
 
-  void RegisterTransaction(Transaction* txn);
+  Status RegisterTransaction(Transaction* txn);
   void UnregisterTransaction(Transaction* txn);
 
   // not thread safe. current use case is during recovery (single thread)
@@ -239,6 +239,7 @@ class PessimisticTransactionDB : public TransactionDB {
   friend class WriteUnpreparedTransactionTest_MarkLogWithPrepSection_Test;
 
   Transaction* BeginInternalTransaction(const WriteOptions& options);
+  Transaction* GetTransactionByNameLocked(const TransactionName& name);
 
   std::shared_ptr<LockManager> lock_manager_;