SERVER-30912 Only check out sessions for write commands

5 files changed, 67 insertions, 7 deletions

diff --git a/jstests/sharding/shard_collection_basic.js b/jstests/sharding/shard_collection_basic.js
index 6dd195cffbe..26ed855f393 100644
--- a/jstests/sharding/shard_collection_basic.js
+++ b/jstests/sharding/shard_collection_basic.js
@@ -173,6 +173,21 @@
         mongos.adminCommand({shardCollection: kDbName + '.foo', key: {aKey: 1}, unique: true}));
 
     //
+    // Session-related tests
+    //
+
+    assert.commandWorked(mongos.getDB(kDbName).dropDatabase());
+    assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
+
+    // shardCollection can be called under a session.
+    const sessionDb = mongos.startSession().getDatabase(kDbName);
+    assert.commandWorked(
+        sessionDb.adminCommand({shardCollection: kDbName + '.foo', key: {_id: 'hashed'}}));
+    sessionDb.getSession().endSession();
+
+    assert.commandWorked(mongos.getDB(kDbName).dropDatabase());
+
+    //
     // Collation-related tests
     //
 
diff --git a/src/mongo/db/service_entry_point_mongod.cpp b/src/mongo/db/service_entry_point_mongod.cpp
index ad74a5be211..5e4bd07e762 100644
--- a/src/mongo/db/service_entry_point_mongod.cpp
+++ b/src/mongo/db/service_entry_point_mongod.cpp
@@ -92,6 +92,21 @@ MONGO_FP_DECLARE(rsStopGetMore);
 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> cmdWhitelist = {{"delete", 1},
+                                     {"eval", 1},
+                                     {"$eval", 1},
+                                     {"findandmodify", 1},
+                                     {"findAndModify", 1},
+                                     {"insert", 1},
+                                     {"refreshLogicalSessionCacheNow", 1},
+                                     {"update", 1}};
+
 inline void opread(const Message& m) {
     if (_diaglog.getLevel() & 2) {
         _diaglog.readop(m.singleData().view2ptr(), m.header().getLen());
@@ -581,7 +596,12 @@ void execCommandDatabase(OperationContext* opCtx,
             return;
         }
 
-        OperationContextSession sessionTxnState(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).
+        OperationContextSession sessionTxnState(
+            opCtx, cmdWhitelist.find(command->getName()) != cmdWhitelist.cend());
 
         ImpersonationSessionGuard guard(opCtx);
         uassertStatusOK(Command::checkAuthorization(command, opCtx, request));
diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index 48a7025ed61..ba8199ec1ab 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -202,11 +202,21 @@ void SessionCatalog::_releaseSession(const LogicalSessionId& lsid) {
     sri->availableCondVar.notify_one();
 }
 
-OperationContextSession::OperationContextSession(OperationContext* opCtx) : _opCtx(opCtx) {
+OperationContextSession::OperationContextSession(OperationContext* opCtx, bool checkOutSession)
+    : _opCtx(opCtx) {
     if (!opCtx->getLogicalSessionId()) {
         return;
     }
 
+    if (!checkOutSession) {
+        // The session may have already been checked out by this operation, so bump the nesting
+        // level if necessary to avoid resetting the session when this command completes.
+        if (auto& checkedOutSession = operationSessionDecoration(opCtx)) {
+            checkedOutSession->checkOutNestingLevel++;
+        }
+        return;
+    }
+
     auto& checkedOutSession = operationSessionDecoration(opCtx);
     if (!checkedOutSession) {
         auto sessionTransactionTable = SessionCatalog::get(opCtx);
diff --git a/src/mongo/db/session_catalog.h b/src/mongo/db/session_catalog.h
index 94c3dba97af..3cb341697da 100644
--- a/src/mongo/db/session_catalog.h
+++ b/src/mongo/db/session_catalog.h
@@ -250,7 +250,7 @@ class OperationContextSession {
     MONGO_DISALLOW_COPYING(OperationContextSession);
 
 public:
-    OperationContextSession(OperationContext* opCtx);
+    OperationContextSession(OperationContext* opCtx, bool checkOutSession);
     ~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 67d2cc61154..11170d02c78 100644
--- a/src/mongo/db/session_catalog_test.cpp
+++ b/src/mongo/db/session_catalog_test.cpp
@@ -69,7 +69,7 @@ TEST_F(SessionCatalogTest, OperationContextSession) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
 
     {
-        OperationContextSession ocs(opCtx());
+        OperationContextSession ocs(opCtx(), true);
         auto session = OperationContextSession::get(opCtx());
 
         ASSERT(session);
@@ -91,7 +91,8 @@ TEST_F(SessionCatalogTest, GetOrCreateSessionAfterCheckOutSession) {
     const auto lsid = makeLogicalSessionIdForTest();
     opCtx()->setLogicalSessionId(lsid);
 
-    boost::optional<OperationContextSession> ocs(opCtx());
+    boost::optional<OperationContextSession> ocs;
+    ocs.emplace(opCtx(), true);
 
     stdx::async(stdx::launch::async, [&] {
         Client::initThreadIfNotAlready();
@@ -120,10 +121,10 @@ TEST_F(SessionCatalogTest, NestedOperationContextSession) {
     opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
 
     {
-        OperationContextSession outerScopedSession(opCtx());
+        OperationContextSession outerScopedSession(opCtx(), true);
 
         {
-            OperationContextSession innerScopedSession(opCtx());
+            OperationContextSession innerScopedSession(opCtx(), true);
 
             auto session = OperationContextSession::get(opCtx());
             ASSERT(session);
@@ -140,5 +141,19 @@ TEST_F(SessionCatalogTest, NestedOperationContextSession) {
     ASSERT(!OperationContextSession::get(opCtx()));
 }
 
+TEST_F(SessionCatalogTest, OnlyCheckOutSessionWithCheckOutSessionTrue) {
+    opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
+
+    {
+        OperationContextSession ocs(opCtx(), true);
+        ASSERT(OperationContextSession::get(opCtx()));
+    }
+
+    {
+        OperationContextSession ocs(opCtx(), false);
+        ASSERT(!OperationContextSession::get(opCtx()));
+    }
+}
+
 }  // namespace
 }  // namespace mongo