SERVER-32880 Parse and add autocommit parameter to Session

11 files changed, 177 insertions, 75 deletions

diff --git a/src/mongo/db/initialize_operation_session_info.cpp b/src/mongo/db/initialize_operation_session_info.cpp
index 5751f086691..f6773ce2441 100644
--- a/src/mongo/db/initialize_operation_session_info.cpp
+++ b/src/mongo/db/initialize_operation_session_info.cpp
@@ -38,13 +38,15 @@
 
 namespace mongo {
 
-void initializeOperationSessionInfo(OperationContext* opCtx,
-                                    const BSONObj& requestBody,
-                                    bool requiresAuth,
-                                    bool isReplSetMemberOrMongos,
-                                    bool supportsDocLocking) {
+boost::optional<OperationSessionInfoFromClient> initializeOperationSessionInfo(
+    OperationContext* opCtx,
+    const BSONObj& requestBody,
+    bool requiresAuth,
+    bool isReplSetMemberOrMongos,
+    bool supportsDocLocking) {
+
     if (!requiresAuth) {
-        return;
+        return boost::none;
     }
 
     {
@@ -54,7 +56,7 @@ void initializeOperationSessionInfo(OperationContext* opCtx,
         AuthorizationSession* authSession = AuthorizationSession::get(opCtx->getClient());
         if (authSession && authSession->isUsingLocalhostBypass() &&
             !authSession->getAuthenticatedUserNames().more()) {
-            return;
+            return boost::none;
         }
     }
 
@@ -88,6 +90,14 @@ void initializeOperationSessionInfo(OperationContext* opCtx,
 
         opCtx->setTxnNumber(*osi.getTxnNumber());
     }
+
+    if (osi.getAutocommit()) {
+        uassert(ErrorCodes::IllegalOperation,
+                "Autocommit requires a transaction number to be specified",
+                opCtx->getTxnNumber());
+    }
+
+    return osi;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/initialize_operation_session_info.h b/src/mongo/db/initialize_operation_session_info.h
index 914284f5426..8882c4e7052 100644
--- a/src/mongo/db/initialize_operation_session_info.h
+++ b/src/mongo/db/initialize_operation_session_info.h
@@ -33,22 +33,27 @@
 namespace mongo {
 
 /**
- * Parses the session information from the body of a request and installs it on the current
- * operation context. Must only be called once per operation and should be done right in the
- * beginning.
+ * Parses the session information from the body of a request and stores the sessionId and txnNumber
+ * on the current operation context. Must only be called once per operation and should be done right
+ * in the beginning.
  *
  * Throws if the sessionId/txnNumber combination is not properly formatted.
  *
  * requiresAuth specifies if the command we're initializing operationSessionInfo for requires
  * authorization or not.  This can be determined by invoking ->requiresAuth() on the parsed command.
+ * If it does not require authorization, return boost::none.
  *
  * Both isReplSetMemberOrMongos and supportsDocLocking need to be true if the command contains a
  * transaction number, otherwise this function will throw.
+ *
+ * On success, returns the parsed request information. Returning boost::none implies that the
+ * proper command or session requirements were not met.
  */
-void initializeOperationSessionInfo(OperationContext* opCtx,
-                                    const BSONObj& requestBody,
-                                    bool requiresAuth,
-                                    bool isReplSetMemberOrMongos,
-                                    bool supportsDocLocking);
+boost::optional<OperationSessionInfoFromClient> initializeOperationSessionInfo(
+    OperationContext* opCtx,
+    const BSONObj& requestBody,
+    bool requiresAuth,
+    bool isReplSetMemberOrMongos,
+    bool supportsDocLocking);
 
 }  // namespace mongo
diff --git a/src/mongo/db/logical_session_id.idl b/src/mongo/db/logical_session_id.idl
index 365f06edc0d..173be7b3c4a 100644
--- a/src/mongo/db/logical_session_id.idl
+++ b/src/mongo/db/logical_session_id.idl
@@ -112,6 +112,9 @@ structs:
                       operation executes."
         type: TxnNumber
         optional: true
+      autocommit:
+        type: bool
+        optional: true
 
   SessionsCollectionFetchResultIndividualResult:
     description: "Individual result"
diff --git a/src/mongo/db/s/session_catalog_migration_destination_test.cpp b/src/mongo/db/s/session_catalog_migration_destination_test.cpp
index a03cf0ef458..698f0464029 100644
--- a/src/mongo/db/s/session_catalog_migration_destination_test.cpp
+++ b/src/mongo/db/s/session_catalog_migration_destination_test.cpp
@@ -178,7 +178,7 @@ public:
                                     const LogicalSessionId& sessionId,
                                     const TxnNumber& txnNum) {
         auto scopedSession = SessionCatalog::get(opCtx)->getOrCreateSession(opCtx, sessionId);
-        scopedSession->beginTxn(opCtx, txnNum);
+        scopedSession->beginOrContinueTxnOnMigration(opCtx, txnNum);
         return scopedSession;
     }
 
@@ -249,7 +249,8 @@ public:
             // requests with txnNumbers aren't allowed. To get around this, we have to manually set
             // up the session state and perform the insert.
             initializeOperationSessionInfo(innerOpCtx.get(), insertBuilder.obj(), true, true, true);
-            OperationContextSession sessionTxnState(innerOpCtx.get(), true);
+            OperationContextSession sessionTxnState(innerOpCtx.get(), true, boost::none);
+
             const auto reply = performInserts(innerOpCtx.get(), insertRequest);
             ASSERT(reply.results.size() == 1);
             ASSERT(reply.results[0].isOK());
@@ -357,7 +358,6 @@ TEST_F(SessionCatalogMigrationDestinationTest, OplogEntriesWithSameTxn) {
     returnOplog({oplog1, oplog2, oplog3});
 
     finishSessionExpectSuccess(&sessionMigration);
-
     auto opCtx = operationContext();
     auto session = getSessionWithTxn(opCtx, sessionId, 2);
     TransactionHistoryIterator historyIter(session->getLastWriteOpTime(2));
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index e87c424b80a..14d33980832 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -475,13 +475,30 @@ void execCommandDatabase(OperationContext* opCtx,
         rpc::TrackingMetadata::get(opCtx).initWithOperName(command->getName());
 
         auto const replCoord = repl::ReplicationCoordinator::get(opCtx);
-        initializeOperationSessionInfo(
+        auto sessionOptions = initializeOperationSessionInfo(
             opCtx,
             request.body,
             command->requiresAuth(),
             replCoord->getReplicationMode() == repl::ReplicationCoordinator::modeReplSet,
             opCtx->getServiceContext()->getGlobalStorageEngine()->supportsDocLocking());
 
+        // Session ids are forwarded in requests, so commands that require roundtrips between
+        // servers may result in a deadlock when a server tries to check out a session it is already
+        // using to service an earlier operation in the command's chain. To avoid this, only check
+        // out sessions for commands that require them (i.e. write commands).
+        // Session checkout is also prevented for commands run within DBDirectClient. If checkout is
+        // required, it is expected to be handled by the outermost command.
+        const bool shouldCheckoutSession =
+            sessionCheckoutWhitelist.find(command->getName()) != sessionCheckoutWhitelist.cend() &&
+            !opCtx->getClient()->isInDirectClient();
+
+        boost::optional<bool> autocommitVal = boost::none;
+        if (sessionOptions && sessionOptions->getAutocommit()) {
+            autocommitVal = *sessionOptions->getAutocommit();
+        }
+
+        OperationContextSession sessionTxnState(opCtx, shouldCheckoutSession, autocommitVal);
+
         const auto dbname = request.getDatabase().toString();
         uassert(
             ErrorCodes::InvalidNamespace,
@@ -527,17 +544,6 @@ void execCommandDatabase(OperationContext* opCtx,
             return;
         }
 
-        // Session ids are forwarded in requests, so commands that require roundtrips between
-        // servers may result in a deadlock when a server tries to check out a session it is already
-        // using to service an earlier operation in the command's chain. To avoid this, only check
-        // out sessions for commands that require them (i.e. write commands).
-        // Session checkout is also prevented for commands run within DBDirectClient. If checkout is
-        // required, it is expected to be handled by the outermost command.
-        const bool shouldCheckoutSession =
-            sessionCheckoutWhitelist.find(command->getName()) != sessionCheckoutWhitelist.cend() &&
-            !opCtx->getClient()->isInDirectClient();
-        OperationContextSession sessionTxnState(opCtx, shouldCheckoutSession);
-
         ImpersonationSessionGuard guard(opCtx);
         uassertStatusOK(Command::checkAuthorization(command, opCtx, request));
 
diff --git a/src/mongo/db/session.cpp b/src/mongo/db/session.cpp
index e98bd224922..2b55a2f6b56 100644
--- a/src/mongo/db/session.cpp
+++ b/src/mongo/db/session.cpp
@@ -258,13 +258,23 @@ void Session::refreshFromStorageIfNeeded(OperationContext* opCtx) {
     }
 }
 
-void Session::beginTxn(OperationContext* opCtx, TxnNumber txnNumber) {
+void Session::beginOrContinueTxn(OperationContext* opCtx,
+                                 TxnNumber txnNumber,
+                                 boost::optional<bool> autocommit) {
     invariant(!opCtx->lockState()->isLocked());
 
     stdx::lock_guard<stdx::mutex> lg(_mutex);
-    _beginTxn(lg, txnNumber);
+    _beginOrContinueTxn(lg, txnNumber, autocommit);
 }
 
+void Session::beginOrContinueTxnOnMigration(OperationContext* opCtx, TxnNumber txnNumber) {
+    invariant(!opCtx->lockState()->isLocked());
+
+    stdx::lock_guard<stdx::mutex> lg(_mutex);
+    _beginOrContinueTxnOnMigration(lg, txnNumber);
+}
+
+
 void Session::onWriteOpCompletedOnPrimary(OperationContext* opCtx,
                                           TxnNumber txnNumber,
                                           std::vector<StmtId> stmtIdsWritten,
@@ -296,7 +306,7 @@ void Session::onWriteOpCompletedOnPrimary(OperationContext* opCtx,
 }
 
 bool Session::onMigrateBeginOnPrimary(OperationContext* opCtx, TxnNumber txnNumber, StmtId stmtId) {
-    beginTxn(opCtx, txnNumber);
+    beginOrContinueTxnOnMigration(opCtx, txnNumber);
 
     try {
         if (checkStatementExecuted(opCtx, txnNumber, stmtId)) {
@@ -413,9 +423,28 @@ bool Session::checkStatementExecutedNoOplogEntryFetch(TxnNumber txnNumber, StmtI
     return bool(_checkStatementExecuted(lg, txnNumber, stmtId));
 }
 
-void Session::_beginTxn(WithLock wl, TxnNumber txnNumber) {
+void Session::_beginOrContinueTxn(WithLock wl,
+                                  TxnNumber txnNumber,
+                                  boost::optional<bool> autocommit) {
     _checkValid(wl);
+    _checkTxnValid(wl, txnNumber);
+
+    if (txnNumber == _activeTxnNumber) {
+        // Continuing an existing transaction.
+        uassert(ErrorCodes::IllegalOperation,
+                "Specifying 'autocommit' is only allowed at the beginning of a transaction",
+                autocommit == boost::none);
 
+        return;
+    }
+
+    // Start a new transaction with an autocommit field
+    _setActiveTxn(wl, txnNumber);
+    _autocommit = (autocommit != boost::none) ? *autocommit : true;  // autocommit defaults to true
+    _isSnapshotTxn = false;
+}
+
+void Session::_checkTxnValid(WithLock, TxnNumber txnNumber) const {
     uassert(ErrorCodes::TransactionTooOld,
             str::stream() << "Cannot start transaction " << txnNumber << " on session "
                           << getSessionId()
@@ -423,15 +452,6 @@ void Session::_beginTxn(WithLock wl, TxnNumber txnNumber) {
                           << _activeTxnNumber
                           << " has already started.",
             txnNumber >= _activeTxnNumber);
-
-    // Check for continuing an existing transaction
-    if (txnNumber == _activeTxnNumber)
-        return;
-
-    _activeTxnNumber = txnNumber;
-    _activeTxnCommittedStatements.clear();
-    _hasIncompleteHistory = false;
-    _isSnapshotTxn = false;
 }
 
 void Session::stashTransactionResources(OperationContext* opCtx) {
@@ -500,6 +520,23 @@ void Session::_releaseStashedTransactionResources(WithLock wl, OperationContext*
     _isSnapshotTxn = false;
 }
 
+void Session::_beginOrContinueTxnOnMigration(WithLock wl, TxnNumber txnNumber) {
+    _checkValid(wl);
+    _checkTxnValid(wl, txnNumber);
+    
+    // Check for continuing an existing transaction
+    if (txnNumber == _activeTxnNumber)
+        return;
+
+    _setActiveTxn(wl, txnNumber);
+}
+
+void Session::_setActiveTxn(WithLock, TxnNumber txnNumber) {
+    _activeTxnNumber = txnNumber;
+    _activeTxnCommittedStatements.clear();
+    _hasIncompleteHistory = false;
+}
+
 void Session::_checkValid(WithLock) const {
     uassert(ErrorCodes::ConflictingOperationInProgress,
             str::stream() << "Session " << getSessionId()
@@ -612,7 +649,7 @@ void Session::_registerUpdateCacheOnCommit(OperationContext* opCtx,
             // invalidated and immediately refreshed while there were no writes for newTxnNumber
             // yet. In this case _activeTxnNumber will be less than newTxnNumber and we will fail to
             // update the cache even though the write was successful.
-            _beginTxn(lg, newTxnNumber);
+            _beginOrContinueTxn(lg, newTxnNumber, boost::none);
         }
 
         if (newTxnNumber == _activeTxnNumber) {
diff --git a/src/mongo/db/session.h b/src/mongo/db/session.h
index 097d7a54346..b005dd375ae 100644
--- a/src/mongo/db/session.h
+++ b/src/mongo/db/session.h
@@ -82,13 +82,28 @@ public:
      * been called. If an attempt is made to start a transaction with number less than the latest
      * transaction this session has seen, an exception will be thrown.
      *
+     * Sets the autocommit parameter for this transaction. If it is boost::none, no autocommit 
+     * parameter was passed into the request. If this is the first statement of a transaction,
+     * the autocommit parameter will default to true.
+     *
+     * Autocommit can only be specified on the first statement of a transaction. If otherwise,
+     * this function will throw.
+     *
      * Throws if the session has been invalidated or if an attempt is made to start a transaction
      * older than the active.
      *
      * In order to avoid the possibility of deadlock, this method must not be called while holding a
      * lock.
      */
-    void beginTxn(OperationContext* opCtx, TxnNumber txnNumber);
+    void beginOrContinueTxn(OperationContext* opCtx,
+                            TxnNumber txnNumber,
+                            boost::optional<bool> autocommit);
+    /**
+     * Similar to beginOrContinueTxn except it is used specifically for shard migrations and does
+     * not check or modify the autocommit parameter.
+     */
+    void beginOrContinueTxnOnMigration(OperationContext* opCtx, TxnNumber txnNumber);
+
 
     /**
      * Called after a write under the specified transaction completes while the node is a primary
@@ -190,11 +205,26 @@ public:
      */
     void unstashTransactionResources(OperationContext* opCtx);
 
+    bool getAutocommit() const {
+        return _autocommit;
+    }
+
 private:
-    void _beginTxn(WithLock, TxnNumber txnNumber);
+    void _beginOrContinueTxn(WithLock,
+                             TxnNumber txnNumber,
+                             boost::optional<bool> autocommit);
+
+    void _beginOrContinueTxnOnMigration(WithLock, TxnNumber txnNumber);
 
+    // Checks if there is a conflicting operation on the current Session
     void _checkValid(WithLock) const;
 
+    // Checks that a new txnNumber is higher than the activeTxnNumber so
+    // we don't start a txn that is too old.
+    void _checkTxnValid(WithLock, TxnNumber txnNumber) const;
+
+    void _setActiveTxn(WithLock, TxnNumber txnNumber);
+
     void _checkIsActiveTransaction(WithLock, TxnNumber txnNumber) const;
 
     boost::optional<repl::OpTime> _checkStatementExecuted(WithLock,
@@ -252,6 +282,9 @@ private:
     // opTime. Used for fast retryability check and retrieving the previous write's data without
     // having to scan through the oplog.
     CommittedStatementTimestampMap _activeTxnCommittedStatements;
+
+    // Set in _beginOrContinueTxn and applies to the activeTxn on the session.
+    bool _autocommit{true};
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index 1227030d7a0..352c0f0e6b1 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -240,8 +240,11 @@ void SessionCatalog::_releaseSession(const LogicalSessionId& lsid) {
     sri->availableCondVar.notify_one();
 }
 
-OperationContextSession::OperationContextSession(OperationContext* opCtx, bool checkOutSession)
+OperationContextSession::OperationContextSession(OperationContext* opCtx,
+                                                 bool checkOutSession,
+                                                 boost::optional<bool> autocommit)
     : _opCtx(opCtx) {
+
     if (!opCtx->getLogicalSessionId()) {
         return;
     }
@@ -273,7 +276,8 @@ OperationContextSession::OperationContextSession(OperationContext* opCtx, bool c
     checkedOutSession->scopedSession->refreshFromStorageIfNeeded(opCtx);
 
     if (opCtx->getTxnNumber()) {
-        checkedOutSession->scopedSession->beginTxn(opCtx, *opCtx->getTxnNumber());
+        checkedOutSession->scopedSession->beginOrContinueTxn(
+            opCtx, *opCtx->getTxnNumber(), autocommit);  
     }
 }
 
diff --git a/src/mongo/db/session_catalog.h b/src/mongo/db/session_catalog.h
index 0451147f786..4d4eb28f63e 100644
--- a/src/mongo/db/session_catalog.h
+++ b/src/mongo/db/session_catalog.h
@@ -246,7 +246,10 @@ class OperationContextSession {
     MONGO_DISALLOW_COPYING(OperationContextSession);
 
 public:
-    OperationContextSession(OperationContext* opCtx, bool checkOutSession);
+    OperationContextSession(OperationContext* opCtx,
+                            bool checkOutSession,
+                            boost::optional<bool> autocommit);
+
     ~OperationContextSession();
 
     static Session* get(OperationContext* opCtx);
diff --git a/src/mongo/db/session_catalog_test.cpp b/src/mongo/db/session_catalog_test.cpp
index 721617c8396..5e6c635ce4f 100644
--- a/src/mongo/db/session_catalog_test.cpp
+++ b/src/mongo/db/session_catalog_test.cpp
@@ -70,7 +70,7 @@ TEST_F(SessionCatalogTest, OperationContextSession) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
 
     {
-        OperationContextSession ocs(opCtx(), true);
+        OperationContextSession ocs(opCtx(), true, boost::none);
         auto session = OperationContextSession::get(opCtx());
 
         ASSERT(session);
@@ -82,7 +82,7 @@ TEST_F(SessionCatalogTest, OperationContextSession) {
     }
 
     {
-        OperationContextSession ocs(opCtx(), false);
+        OperationContextSession ocs(opCtx(), false, boost::none);
         auto session = OperationContextSession::get(opCtx());
 
         ASSERT(!session);
@@ -107,7 +107,7 @@ TEST_F(SessionCatalogTest, GetOrCreateSessionAfterCheckOutSession) {
     opCtx()->setLogicalSessionId(lsid);
 
     boost::optional<OperationContextSession> ocs;
-    ocs.emplace(opCtx(), true);
+    ocs.emplace(opCtx(), true, boost::none);
 
     stdx::async(stdx::launch::async, [&] {
         Client::initThreadIfNotAlready();
@@ -136,10 +136,10 @@ TEST_F(SessionCatalogTest, NestedOperationContextSession) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
 
     {
-        OperationContextSession outerScopedSession(opCtx(), true);
+        OperationContextSession outerScopedSession(opCtx(), true, boost::none);
 
         {
-            OperationContextSession innerScopedSession(opCtx(), true);
+            OperationContextSession innerScopedSession(opCtx(), true, boost::none);
 
             auto session = OperationContextSession::get(opCtx());
             ASSERT(session);
@@ -163,10 +163,10 @@ DEATH_TEST_F(SessionCatalogTest,
     opCtx()->setTxnNumber(1);
 
     {
-        OperationContextSession outerScopedSession(opCtx(), true);
+        OperationContextSession outerScopedSession(opCtx(), true, boost::none);
 
         {
-            OperationContextSession innerScopedSession(opCtx(), true);
+            OperationContextSession innerScopedSession(opCtx(), true, boost::none);
             innerScopedSession.stashTransactionResources();
         }
     }
@@ -179,10 +179,10 @@ DEATH_TEST_F(SessionCatalogTest,
     opCtx()->setTxnNumber(1);
 
     {
-        OperationContextSession outerScopedSession(opCtx(), true);
+        OperationContextSession outerScopedSession(opCtx(), true, boost::none);
 
         {
-            OperationContextSession innerScopedSession(opCtx(), true);
+            OperationContextSession innerScopedSession(opCtx(), true, boost::none);
             innerScopedSession.unstashTransactionResources();
         }
     }
@@ -192,12 +192,12 @@ TEST_F(SessionCatalogTest, OnlyCheckOutSessionWithCheckOutSessionTrue) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
 
     {
-        OperationContextSession ocs(opCtx(), true);
+        OperationContextSession ocs(opCtx(), true, boost::none);
         ASSERT(OperationContextSession::get(opCtx()));
     }
 
     {
-        OperationContextSession ocs(opCtx(), false);
+        OperationContextSession ocs(opCtx(), false, boost::none);
         ASSERT(!OperationContextSession::get(opCtx()));
     }
 }
diff --git a/src/mongo/db/session_test.cpp b/src/mongo/db/session_test.cpp
index a1df9d1ec21..6dde3fc0bea 100644
--- a/src/mongo/db/session_test.cpp
+++ b/src/mongo/db/session_test.cpp
@@ -132,7 +132,7 @@ TEST_F(SessionTest, SessionEntryNotWrittenOnBegin) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 20;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     ASSERT_EQ(sessionId, session.getSessionId());
     ASSERT(session.getLastWriteOpTime(txnNum).isNull());
@@ -150,7 +150,7 @@ TEST_F(SessionTest, SessionEntryWrittenAtFirstWrite) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 21;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     const auto opTime = [&] {
         AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
@@ -183,7 +183,7 @@ TEST_F(SessionTest, StartingNewerTransactionUpdatesThePersistedSession) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const auto writeTxnRecordFn = [&](TxnNumber txnNum, StmtId stmtId, repl::OpTime prevOpTime) {
-        session.beginTxn(opCtx(), txnNum);
+        session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
         AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
         WriteUnitOfWork wuow(opCtx());
@@ -222,10 +222,11 @@ TEST_F(SessionTest, StartingOldTxnShouldAssert) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 20;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
-    ASSERT_THROWS_CODE(
-        session.beginTxn(opCtx(), txnNum - 1), AssertionException, ErrorCodes::TransactionTooOld);
+    ASSERT_THROWS_CODE(session.beginOrContinueTxn(opCtx(), txnNum - 1, boost::none),
+                       AssertionException,
+                       ErrorCodes::TransactionTooOld);
     ASSERT(session.getLastWriteOpTime(txnNum).isNull());
 }
 
@@ -241,7 +242,7 @@ TEST_F(SessionTest, SessionTransactionsCollectionNotDefaultCreated) {
     ASSERT(client.runCommand(nss.db().toString(), BSON("drop" << nss.coll()), dropResult));
 
     const TxnNumber txnNum = 21;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
     WriteUnitOfWork wuow(opCtx());
@@ -256,7 +257,7 @@ TEST_F(SessionTest, CheckStatementExecuted) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 100;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     const auto writeTxnRecordFn = [&](StmtId stmtId, repl::OpTime prevOpTime) {
         AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
@@ -297,7 +298,7 @@ TEST_F(SessionTest, CheckStatementExecutedForOldTransactionThrows) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 100;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     ASSERT_THROWS_CODE(session.checkStatementExecuted(opCtx(), txnNum - 1, 0),
                        AssertionException,
@@ -320,7 +321,7 @@ TEST_F(SessionTest, WriteOpCompletedOnPrimaryForOldTransactionThrows) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 100;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     {
         AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
@@ -347,7 +348,7 @@ TEST_F(SessionTest, WriteOpCompletedOnPrimaryForInvalidatedTransactionThrows) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 100;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
     WriteUnitOfWork wuow(opCtx());
@@ -367,7 +368,7 @@ TEST_F(SessionTest, WriteOpCompletedOnPrimaryCommitIgnoresInvalidation) {
     session.refreshFromStorageIfNeeded(opCtx());
 
     const TxnNumber txnNum = 100;
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     {
         AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
@@ -462,7 +463,7 @@ TEST_F(SessionTest, ErrorOnlyWhenStmtIdBeingCheckedIsNotInCache) {
 
     Session session(sessionId);
     session.refreshFromStorageIfNeeded(opCtx());
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     auto firstOpTime = ([&]() {
         AutoGetCollection autoColl(opCtx(), kNss, MODE_IX);
@@ -548,7 +549,7 @@ TEST_F(SessionTest, StashAndUnstashResources) {
     Session session(sessionId);
     session.refreshFromStorageIfNeeded(opCtx());
 
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     repl::ReadConcernArgs readConcernArgs;
     ASSERT_OK(readConcernArgs.initialize(BSON("find"
@@ -593,7 +594,7 @@ TEST_F(SessionTest, StashNotRequired) {
     Session session(sessionId);
     session.refreshFromStorageIfNeeded(opCtx());
 
-    session.beginTxn(opCtx(), txnNum);
+    session.beginOrContinueTxn(opCtx(), txnNum, boost::none);
 
     repl::ReadConcernArgs readConcernArgs;
     ASSERT_OK(readConcernArgs.initialize(BSON("find"