SERVER-67466 Internal transactions API may be memory unsafe when outer OperationContext is interrupted

2 files changed, 8 insertions, 5 deletions

diff --git a/src/mongo/db/transaction_api.cpp b/src/mongo/db/transaction_api.cpp
index 0a950dbbe54..4b8bd6cbb0a 100644
--- a/src/mongo/db/transaction_api.cpp
+++ b/src/mongo/db/transaction_api.cpp
@@ -101,10 +101,16 @@ StatusWith<CommitResult> SyncTransactionWithRetries::runNoThrow(OperationContext
         return yieldStatus;
     }
 
-    auto txnResult = _txn->run(std::move(callback)).getNoThrow(opCtx);
+    auto txnFuture = _txn->run(std::move(callback));
+    auto txnResult = txnFuture.getNoThrow(opCtx);
+
     // Cancel the source to guarantee the transaction will terminate if our opCtx was interrupted.
     _source.cancel();
 
+    // Wait for transaction to complete before returning so variables referenced by its callback are
+    // guaranteed to be in scope even if the API caller's opCtx was interrupted.
+    txnFuture.wait();
+
     // Post transaction processing, which must also happen inline.
     OperationTimeTracker::get(opCtx)->updateOperationTime(_txn->getOperationTime());
     repl::ReplClientInfo::forClient(opCtx->getClient())
diff --git a/src/mongo/db/transaction_api_test.cpp b/src/mongo/db/transaction_api_test.cpp
index bcfb8ba5815..90a8daf4366 100644
--- a/src/mongo/db/transaction_api_test.cpp
+++ b/src/mongo/db/transaction_api_test.cpp
@@ -1203,6 +1203,7 @@ TEST_F(TxnAPITest, UnyieldsAfterCancellation) {
     auto killerThread = stdx::thread([&txnApiStarted, &opCtxKilled, opCtx = opCtx()] {
         txnApiStarted.countDownAndWait();
         opCtx->markKilled();
+        opCtxKilled.countDownAndWait();
     });
 
     auto swResult = txnWithRetries().runNoThrow(
@@ -1230,11 +1231,7 @@ TEST_F(TxnAPITest, UnyieldsAfterCancellation) {
     ASSERT_EQ(resourceYielder()->timesYielded(), 1);
     ASSERT_EQ(resourceYielder()->timesUnyielded(), 1);
 
-    opCtxKilled.countDownAndWait();
     killerThread.join();
-
-    // Wait until the API callback has completed before letting the barriers go out of scope.
-    waitForAllEarlierTasksToComplete();
 }
 
 TEST_F(TxnAPITest, ClientSession_UsesNonRetryableInternalSession) {