1 files changed, 20 insertions, 2 deletions

diff --git a/src/mongo/s/transaction_router.h b/src/mongo/s/transaction_router.h
index 276a473aead..a7b49c5eb12 100644
--- a/src/mongo/s/transaction_router.h
+++ b/src/mongo/s/transaction_router.h
@@ -71,6 +71,9 @@ public:
         kRecoverWithToken,
     };
 
+    // The reason why TransactionRouter::Router::stash() is called.
+    enum class StashReason { kDone, kYield };
+
     // The default value to use as the statement id of the first command in the transaction if none
     // was sent.
     static const StmtId kDefaultFirstStmtId = 0;
@@ -326,6 +329,15 @@ public:
             return o().txnNumberAndRetryCounter.getTxnNumber() != kUninitializedTxnNumber;
         }
 
+        /**
+         * Returns if this TransactionRouter instance can be reaped. Always true unless an operation
+         * has yielded the router and has not unyielded yet. We cannot reap the instance in that
+         * case or the unyield would check out a different TransactionRouter than it yielded.
+         */
+        auto canBeReaped() const {
+            return o().activeYields == 0;
+        }
+
     protected:
         explicit Observer(TransactionRouter* tr) : _tr(tr) {}
 
@@ -372,9 +384,10 @@ public:
                                 TransactionActions action);
 
         /**
-         * Updates transaction diagnostics when the transaction's session is checked in.
+         * Updates transaction diagnostics and, if necessary, the number of active yielders when the
+         * transaction's session is checked in.
          */
-        void stash(OperationContext* opCtx);
+        void stash(OperationContext* opCtx, StashReason reason);
 
         /**
          * Validates transaction state is still compatible after a yield.
@@ -775,6 +788,11 @@ private:
         // certain transaction events. Unset until the transaction router has processed at least one
         // transaction command.
         boost::optional<MetricsTracker> metricsTracker;
+
+        // How many operations that checked out the router's session have currently yielded it. The
+        // transaction number cannot be changed until this returns to 0, otherwise we cannot
+        // guarantee that unyielding the session cannot fail.
+        int32_t activeYields{0};
     } _o;
 
     /**