SERVER-33440 DBDirectClient operations should be able to perform session checkout

3 files changed, 80 insertions, 38 deletions

diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index 09bf4b29962..76de224614c 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -96,25 +96,28 @@ MONGO_FP_DECLARE(skipCheckingForNotMasterInCommandDispatch);
 namespace {
 using logger::LogComponent;
 
-// The command names for which to check out a session.
-//
-// Note: Eval should check out a session because it defaults to running under a global write lock,
-// so if it didn't, and the function it was given contains any of these whitelisted commands, they
-// would try to check out a session under a lock, which is not allowed.  Similarly,
-// refreshLogicalSessionCacheNow triggers a bulk update under a lock on the sessions collection.
-const StringMap<int> sessionCheckoutWhitelist = {{"delete", 1},
+// The command names for which to check out a session. These are commands that support retryable
+// writes, readConcern snapshot, or multi-statement transactions. We additionally check out the
+// session for commands that can take a lock and then run another whitelisted command in
+// DBDirectClient. Otherwise, the nested command would try to check out a session under a lock,
+// which is not allowed.
+const StringMap<int> sessionCheckoutWhitelist = {{"applyOps", 1},
+                                                 {"count", 1},
+                                                 {"delete", 1},
                                                  {"eval", 1},
                                                  {"$eval", 1},
+                                                 {"explain", 1},
+                                                 {"find", 1},
                                                  {"findandmodify", 1},
                                                  {"findAndModify", 1},
+                                                 {"geoSearch", 1},
+                                                 {"getMore", 1},
+                                                 {"group", 1},
                                                  {"insert", 1},
+                                                 {"mapReduce", 1},
+                                                 {"parallelCollectionScan", 1},
                                                  {"refreshLogicalSessionCacheNow", 1},
-                                                 {"update", 1},
-                                                 {"find", 1},
-                                                 {"getMore", 1},
-                                                 {"count", 1},
-                                                 {"geoSearch", 1},
-                                                 {"parallelCollectionScan", 1}};
+                                                 {"update", 1}};
 
 // The command names for which readConcern level snapshot is allowed. The getMore command is
 // implicitly allowed to operate on a cursor which was opened under readConcern level snapshot.
@@ -486,12 +489,9 @@ void execCommandDatabase(OperationContext* opCtx,
         // 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.
+        // out sessions for commands that require them.
         const bool shouldCheckoutSession =
-            sessionCheckoutWhitelist.find(command->getName()) != sessionCheckoutWhitelist.cend() &&
-            !opCtx->getClient()->isInDirectClient();
+            sessionCheckoutWhitelist.find(command->getName()) != sessionCheckoutWhitelist.cend();
 
         boost::optional<bool> autocommitVal = boost::none;
         if (sessionOptions && sessionOptions->getAutocommit()) {
diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index f651b7121f9..dcccae45804 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -292,14 +292,11 @@ OperationContextSession::~OperationContextSession() {
 }
 
 void OperationContextSession::stashTransactionResources() {
-    if (_opCtx->getClient()->isInDirectClient()) {
-        return;
-    }
-
     if (auto& checkedOutSession = operationSessionDecoration(_opCtx)) {
-        invariant(checkedOutSession->checkOutNestingLevel == 1);
-        if (auto session = checkedOutSession->scopedSession.get()) {
-            session->stashTransactionResources(_opCtx);
+        if (checkedOutSession->checkOutNestingLevel == 1) {
+            if (auto session = checkedOutSession->scopedSession.get()) {
+                session->stashTransactionResources(_opCtx);
+            }
         }
     }
 }
@@ -309,14 +306,11 @@ void OperationContextSession::unstashTransactionResources() {
         return;
     }
 
-    if (_opCtx->getClient()->isInDirectClient()) {
-        return;
-    }
-
     if (auto& checkedOutSession = operationSessionDecoration(_opCtx)) {
-        invariant(checkedOutSession->checkOutNestingLevel == 1);
-        if (auto session = checkedOutSession->scopedSession.get()) {
-            session->unstashTransactionResources(_opCtx);
+        if (checkedOutSession->checkOutNestingLevel == 1) {
+            if (auto session = checkedOutSession->scopedSession.get()) {
+                session->unstashTransactionResources(_opCtx);
+            }
         }
     }
 }
diff --git a/src/mongo/db/session_catalog_test.cpp b/src/mongo/db/session_catalog_test.cpp
index 5e6c635ce4f..03a00304ea5 100644
--- a/src/mongo/db/session_catalog_test.cpp
+++ b/src/mongo/db/session_catalog_test.cpp
@@ -31,6 +31,7 @@
 #include "mongo/db/client.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/repl/mock_repl_coord_server_fixture.h"
+#include "mongo/db/repl/read_concern_args.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/session_catalog.h"
 #include "mongo/stdx/future.h"
@@ -156,34 +157,81 @@ TEST_F(SessionCatalogTest, NestedOperationContextSession) {
     ASSERT(!OperationContextSession::get(opCtx()));
 }
 
-DEATH_TEST_F(SessionCatalogTest,
-             CannotStashInNestedOperationContext,
-             "Invariant failure checkedOutSession->checkOutNestingLevel == 1") {
+TEST_F(SessionCatalogTest, StashInNestedSessionIsANoop) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
     opCtx()->setTxnNumber(1);
 
     {
         OperationContextSession outerScopedSession(opCtx(), true, boost::none);
 
+        Locker* originalLocker = opCtx()->lockState();
+        RecoveryUnit* originalRecoveryUnit = opCtx()->recoveryUnit();
+        ASSERT(originalLocker);
+        ASSERT(originalRecoveryUnit);
+
+        // Set the readConcern on the OperationContext.
+        repl::ReadConcernArgs readConcernArgs;
+        ASSERT_OK(readConcernArgs.initialize(BSON("find"
+                                                  << "test"
+                                                  << repl::ReadConcernArgs::kReadConcernFieldName
+                                                  << BSON(repl::ReadConcernArgs::kLevelFieldName
+                                                          << "snapshot"))));
+        repl::ReadConcernArgs::get(opCtx()) = readConcernArgs;
+
+        // Perform initial unstash, which sets up a WriteUnitOfWork.
+        outerScopedSession.unstashTransactionResources();
+        ASSERT_EQUALS(originalLocker, opCtx()->lockState());
+        ASSERT_EQUALS(originalRecoveryUnit, opCtx()->recoveryUnit());
+        ASSERT(opCtx()->getWriteUnitOfWork());
+
         {
             OperationContextSession innerScopedSession(opCtx(), true, boost::none);
+
+            // Indicate that there is a stashed cursor. If we were not in a nested session, this
+            // would ensure that stashing is not a noop.
+            opCtx()->setStashedCursor();
+
             innerScopedSession.stashTransactionResources();
+
+            // The stash was a noop, so the locker, RecoveryUnit, and WriteUnitOfWork on the
+            // OperationContext are unaffected.
+            ASSERT_EQUALS(originalLocker, opCtx()->lockState());
+            ASSERT_EQUALS(originalRecoveryUnit, opCtx()->recoveryUnit());
+            ASSERT(opCtx()->getWriteUnitOfWork());
         }
     }
 }
 
-DEATH_TEST_F(SessionCatalogTest,
-             CannotUnstashInNestedOperationContext,
-             "Invariant failure checkedOutSession->checkOutNestingLevel == 1") {
+TEST_F(SessionCatalogTest, UnstashInNestedSessionIsANoop) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
     opCtx()->setTxnNumber(1);
 
     {
         OperationContextSession outerScopedSession(opCtx(), true, boost::none);
 
+        Locker* originalLocker = opCtx()->lockState();
+        RecoveryUnit* originalRecoveryUnit = opCtx()->recoveryUnit();
+        ASSERT(originalLocker);
+        ASSERT(originalRecoveryUnit);
+
+        // Set the readConcern on the OperationContext.
+        repl::ReadConcernArgs readConcernArgs;
+        ASSERT_OK(readConcernArgs.initialize(BSON("find"
+                                                  << "test"
+                                                  << repl::ReadConcernArgs::kReadConcernFieldName
+                                                  << BSON(repl::ReadConcernArgs::kLevelFieldName
+                                                          << "snapshot"))));
+        repl::ReadConcernArgs::get(opCtx()) = readConcernArgs;
+
         {
             OperationContextSession innerScopedSession(opCtx(), true, boost::none);
+
             innerScopedSession.unstashTransactionResources();
+
+            // The unstash was a noop, so the OperationContext did not get a WriteUnitOfWork.
+            ASSERT_EQUALS(originalLocker, opCtx()->lockState());
+            ASSERT_EQUALS(originalRecoveryUnit, opCtx()->recoveryUnit());
+            ASSERT_FALSE(opCtx()->getWriteUnitOfWork());
         }
     }
 }