SERVER-91839 Unyielding transaction participant resources should fail if resources not stashed (#24133)

GitOrigin-RevId: 0e826100543b1c0e7de29a52fbcfdeb150d834fa

src/mongo/db/service_entry_point_shard_role.cpp

@@ -1034,11 +1034,6 @@ void CheckoutSessionAndInvokeCommand::_checkOutSession() {
             }
         }
 
-        if (opCtx->isStartingMultiDocumentTransaction()) {
-            service_entry_point_shard_role_helpers::waitForReadConcern(
-                opCtx, _ecd->getInvocation(), execContext.getRequest());
-        }
-
         // Release the transaction lock resources and abort storage transaction for unprepared
         // transactions on failure to unstash the transaction resources to opCtx. We don't want
         // to have this error guard for beginOrContinue as it can abort the transaction for any
@@ -1052,6 +1047,11 @@ void CheckoutSessionAndInvokeCommand::_checkOutSession() {
             }
         });
 
+        if (opCtx->isStartingMultiDocumentTransaction()) {
+            service_entry_point_shard_role_helpers::waitForReadConcern(
+                opCtx, _ecd->getInvocation(), execContext.getRequest());
+        }
+
         txnParticipant.unstashTransactionResources(opCtx, invocation->definition()->getName());
 
         // Unstash success.

src/mongo/db/transaction/transaction_participant.cpp

@@ -1616,7 +1616,10 @@ void TransactionParticipant::Participant::_releaseTransactionResourcesToOpCtx(
 }
 
 void TransactionParticipant::Participant::unstashTransactionResources(
-    OperationContext* opCtx, const std::string& cmdName, bool forRecoveryPreparedTxnApplication) {
+    OperationContext* opCtx,
+    const std::string& cmdName,
+    bool forRecoveryPreparedTxnApplication,
+    bool forUnyield) {
     invariant(!opCtx->getClient()->isInDirectClient());
     invariant(opCtx->getTxnNumber());
 
@@ -1695,6 +1698,11 @@ void TransactionParticipant::Participant::unstashTransactionResources(
         return;
     }
 
+    uassert(9183900,
+            str::stream()
+                << "Expected to have a transaction resource stash when unyielding, but did not.",
+            !forUnyield);
+
     // If we have no transaction resources then we cannot be prepared. If we're not in progress,
     // we don't do anything else.
     invariant(!o().txnState.isPrepared());

src/mongo/db/transaction/transaction_participant.h

@@ -616,7 +616,8 @@ public:
          */
         void unstashTransactionResources(OperationContext* opCtx,
                                          const std::string& cmdName,
-                                         bool forRecoveryPreparedTxnApplication = false);
+                                         bool forRecoveryPreparedTxnApplication = false,
+                                         bool forUnyield = false);
 
         /**
          * Puts a transaction into a prepared state and returns the prepareTimestamp and the list of

src/mongo/db/transaction/transaction_participant_resource_yielder.cpp

@@ -80,7 +80,11 @@ void TransactionParticipantResourceYielder::unyield(OperationContext* opCtx) {
                     opCtx, OperationContextSession::CheckInReason::kYield);
             });
 
-            txnParticipant.unstashTransactionResources(opCtx, _cmdName);
+            txnParticipant.unstashTransactionResources(
+                opCtx,
+                _cmdName,
+                false /* forRecoveryPreparedTxnApplication */,
+                true /* forUnyield */);
             releaseOnError.dismiss();
         }
     }

src/mongo/s/async_requests_sender.cpp

@@ -70,6 +70,7 @@ namespace {
 const int kMaxNumFailedHostRetryAttempts = 3;
 
 MONGO_FAIL_POINT_DEFINE(hangBeforePollResponse);
+MONGO_FAIL_POINT_DEFINE(hangAfterYield);
 
 }  // namespace
 
@@ -149,6 +150,10 @@ AsyncRequestsSender::Response AsyncRequestsSender::next() noexcept {
     try {
         if (_resourceYielder) {
             _resourceYielder->yield(_opCtx);
+
+            if (MONGO_unlikely(hangAfterYield.shouldFail())) {
+                hangAfterYield.pauseWhileSet();
+            }
         }
 
         auto curOp = CurOp::get(_opCtx);

src/mongo/s/multi_statement_transaction_requests_sender_test.cpp

@@ -376,5 +376,102 @@ TEST_F(ShardsvrMultiStatementTransactionRequestsSenderTest,
     SessionCatalog::get(operationContext()->getServiceContext())->reset_forTest();
 }
 
+TEST_F(ShardsvrMultiStatementTransactionRequestsSenderTest, FailUnyieldIfTxnStashDoesNotExist) {
+    auto lsid = makeLogicalSessionIdForTest();
+    operationContext()->setLogicalSessionId(lsid);
+    operationContext()->setTxnNumber(TxnNumber(0));
+    operationContext()->setInMultiDocumentTransaction();
+
+    repl::ReadConcernArgs readConcernArgs{LogicalTime(Timestamp(1, 0)),
+                                          repl::ReadConcernLevel::kMajorityReadConcern};
+    repl::ReadConcernArgs::get(operationContext()) = readConcernArgs;
+
+    // Set up transaction state
+    auto mongoDSessionCatalog = MongoDSessionCatalog::get(operationContext());
+    auto contextSession = mongoDSessionCatalog->checkOutSession(operationContext());
+
+    auto txnParticipant = TransactionParticipant::get(operationContext());
+    txnParticipant.beginOrContinue(operationContext(),
+                                   TxnNumber(0),
+                                   false,
+                                   TransactionParticipant::TransactionActions::kStart);
+    txnParticipant.unstashTransactionResources(operationContext(), "insert");
+
+    // Schedule remote request
+    std::vector<AsyncRequestsSender::Request> requests;
+    requests.emplace_back(kTestRemoteShardId1,
+                          BSON("find"
+                               << "bar"));
+    std::set<ShardId> pendingShardIds{kTestRemoteShardId1};
+
+    auto msars = MultiStatementTransactionRequestsSender(
+        operationContext(),
+        executor(),
+        DatabaseName::createDatabaseName_forTest(boost::none, "db"),
+        requests,
+        ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+        Shard::RetryPolicy::kNoRetry);
+
+    // Force the ARS to hang after it yields
+    setGlobalFailPoint("hangAfterYield",
+                       BSON("mode"
+                            << "alwaysOn"));
+
+    // Call MSARS::next(), which will hang after yielding because of the failpoint above. Once the
+    // failpoint is turned off and the MSARS tries to unyield, it should fail.
+    auto errorOnMsarsNextFuture = launchAsync([&]() {
+        auto response = msars.next();
+
+        // Get the response from the find request. Assert TransactionParticipantFailedUnyield is
+        // thrown even after a successful response, and contains an error with code 9183900.
+        auto status = response.swResponse.getStatus();
+        assertFailedToUnyieldError(status, ErrorCodes::duplicateCodeForTest(9183900));
+    });
+
+    // Mock a concurrent request in the same transaction. This request will block waiting to check
+    // out the session until the MSARS yields. We don't actually run an op here, but instead set up
+    // the TransactionParticipant state in such a way to mimic the request failing without aborting
+    // the transaction. This will mean the request unstashes the txnResources, but does not restash
+    // them nor mark the transaction as aborted.
+    auto mockConcurrentRequestFuture = launchAsync([&]() {
+        auto newClientOwned = getServiceContext()->getService()->makeClient("newClient");
+        AlternativeClientRegion acr(newClientOwned);
+        auto newOpCtx = cc().makeOperationContext();
+
+        newOpCtx->setLogicalSessionId(lsid);
+        newOpCtx->setTxnNumber(TxnNumber(0));
+        newOpCtx->setInMultiDocumentTransaction();
+
+        auto mongoDSessionCatalog = MongoDSessionCatalog::get(newOpCtx.get());
+        auto contextSession = mongoDSessionCatalog->checkOutSession(newOpCtx.get());
+
+        auto txnParticipant = TransactionParticipant::get(newOpCtx.get());
+        txnParticipant.beginOrContinue(newOpCtx.get(),
+                                       TxnNumber(0),
+                                       false,
+                                       TransactionParticipant::TransactionActions::kContinue);
+        txnParticipant.unstashTransactionResources(newOpCtx.get(), "insert");
+    });
+
+    // Wait until the concurrent request has unstashed the txnResources, then turn off the
+    // failpoint.
+    mockConcurrentRequestFuture.default_timed_get();
+    setGlobalFailPoint("hangAfterYield",
+                       BSON("mode"
+                            << "off"));
+
+    // Mock a successful find response from the remote shard, and then wait for the MSARS to fail
+    // when unyieding.
+    onCommand([&](const auto& request) {
+        ASSERT(request.cmdObj["find"]);
+        return CursorResponse(
+                   NamespaceString::createNamespaceString_forTest("db.bar"), 0LL, {BSON("x" << 1)})
+            .toBSON(CursorResponse::ResponseType::InitialResponse);
+    });
+    errorOnMsarsNextFuture.default_timed_get();
+
+    SessionCatalog::get(operationContext()->getServiceContext())->reset_forTest();
+}
+
 }  // namespace
 }  // namespace mongo