SERVER-66851 Internal session pool should use parent sessions as keys for child session map

(cherry picked from commit 37e28d05ca12d2b3835bbe2145cbe258872e7ed6)

3 files changed, 144 insertions, 14 deletions

diff --git a/jstests/sharding/internal_txns/session_pooling.js b/jstests/sharding/internal_txns/session_pooling.js
new file mode 100644
index 00000000000..38d3a165776
--- /dev/null
+++ b/jstests/sharding/internal_txns/session_pooling.js
@@ -0,0 +1,95 @@
+/*
+ * Tests that sessions used for non-retryable internal transactions are pooled.
+ *
+ * @tags: [requires_fcv_60, uses_transactions]
+ */
+(function() {
+"use strict";
+
+const st = new ShardingTest({shards: 1, config: 1});
+const shard0Primary = st.rs0.getPrimary();
+
+const kDbName = "testDb";
+const kCollName = "testColl";
+const mongosTestColl = st.s.getCollection(kDbName + "." + kCollName);
+assert.commandWorked(mongosTestColl.insert({x: 1}));  // Set up the collection.
+
+const sessionsColl = st.s.getCollection("config.system.sessions");
+const transactionsCollOnShard = shard0Primary.getCollection("config.transactions");
+
+function assertNumEntries({sessionUUID, numSessionsCollEntries, numTransactionsCollEntries}) {
+    const filter = {"_id.id": sessionUUID};
+    assert.eq(numSessionsCollEntries,
+              sessionsColl.find(filter).itcount(),
+              tojson(sessionsColl.find().toArray()));
+    assert.eq(numTransactionsCollEntries,
+              transactionsCollOnShard.find(filter).itcount(),
+              tojson(transactionsCollOnShard.find().toArray()));
+}
+
+function runInternalTxn(conn, lsid) {
+    const testInternalTxnCmdObj = {
+        testInternalTransactions: 1,
+        commandInfos:
+            [{dbName: kDbName, command: {insert: kCollName, documents: [{x: -10}, {x: 10}]}}],
+        lsid: lsid,
+    };
+    assert.commandWorked(conn.adminCommand(testInternalTxnCmdObj));
+}
+
+function runTest(conn) {
+    const parentLsid = {id: UUID()};
+
+    runInternalTxn(conn, parentLsid);
+
+    assert.commandWorked(conn.adminCommand({refreshLogicalSessionCacheNow: 1}));
+    assert.commandWorked(shard0Primary.adminCommand({refreshLogicalSessionCacheNow: 1}));
+    assertNumEntries(
+        {sessionUUID: parentLsid.id, numSessionsCollEntries: 1, numTransactionsCollEntries: 1});
+
+    // Run more transactions and verify the number of sessions used doesn't change.
+
+    runInternalTxn(conn, parentLsid);
+    runInternalTxn(conn, parentLsid);
+    runInternalTxn(conn, parentLsid);
+
+    assert.commandWorked(conn.adminCommand({refreshLogicalSessionCacheNow: 1}));
+    assert.commandWorked(shard0Primary.adminCommand({refreshLogicalSessionCacheNow: 1}));
+    assertNumEntries(
+        {sessionUUID: parentLsid.id, numSessionsCollEntries: 1, numTransactionsCollEntries: 1});
+
+    // Verify other sessions can be pooled concurrently.
+
+    const otherParentLsid1 = {id: UUID()};
+    const otherParentLsid2 = {id: UUID()};
+
+    runInternalTxn(conn, otherParentLsid1);
+    runInternalTxn(conn, otherParentLsid2);
+    runInternalTxn(conn, otherParentLsid2);
+    runInternalTxn(conn, otherParentLsid1);
+    runInternalTxn(conn, otherParentLsid2);
+    runInternalTxn(conn, otherParentLsid2);
+    runInternalTxn(conn, otherParentLsid1);
+
+    runInternalTxn(conn, parentLsid);
+
+    assert.commandWorked(conn.adminCommand({refreshLogicalSessionCacheNow: 1}));
+    assert.commandWorked(shard0Primary.adminCommand({refreshLogicalSessionCacheNow: 1}));
+    assertNumEntries(
+        {sessionUUID: parentLsid.id, numSessionsCollEntries: 1, numTransactionsCollEntries: 1});
+    assertNumEntries({
+        sessionUUID: otherParentLsid1.id,
+        numSessionsCollEntries: 1,
+        numTransactionsCollEntries: 1
+    });
+    assertNumEntries({
+        sessionUUID: otherParentLsid2.id,
+        numSessionsCollEntries: 1,
+        numTransactionsCollEntries: 1
+    });
+}
+
+runTest(st.s);
+
+st.stop();
+})();
diff --git a/src/mongo/db/internal_session_pool.cpp b/src/mongo/db/internal_session_pool.cpp
index ab94050f22b..b4bd608cbdb 100644
--- a/src/mongo/db/internal_session_pool.cpp
+++ b/src/mongo/db/internal_session_pool.cpp
@@ -98,7 +98,7 @@ InternalSessionPool::Session InternalSessionPool::acquireSystemSession() {
     const InternalSessionPool::Session session = [&] {
         stdx::lock_guard<Latch> lock(_mutex);
 
-        auto const& systemUserDigest = (*internalSecurity.getUser())->getDigest();
+        const auto& systemUserDigest = makeSystemLogicalSessionId().getUid();
         auto session = _acquireSession(systemUserDigest, lock);
         return session ? *session
                        : InternalSessionPool::Session(makeSystemLogicalSessionId(), TxnNumber(0));
@@ -118,7 +118,7 @@ InternalSessionPool::Session InternalSessionPool::acquireStandaloneSession(
     const InternalSessionPool::Session session = [&] {
         stdx::lock_guard<Latch> lock(_mutex);
 
-        auto const& userDigest = getLogicalSessionUserDigestForLoggedInUser(opCtx);
+        const auto& userDigest = getLogicalSessionUserDigestForLoggedInUser(opCtx);
         auto session = _acquireSession(userDigest, lock);
         return session ? *session
                        : InternalSessionPool::Session(makeLogicalSessionId(opCtx), TxnNumber(0));
@@ -141,11 +141,11 @@ InternalSessionPool::Session InternalSessionPool::acquireChildSession(
         auto it = _childSessions.find(parentLsid);
         if (it != _childSessions.end()) {
             auto session = std::move(it->second);
+            invariant(parentLsid == getParentSessionId(session.getSessionId()));
             _childSessions.erase(it);
             return session;
         } else {
-            auto lsid = LogicalSessionId{parentLsid.getId(), parentLsid.getUid()};
-            lsid.setTxnUUID(UUID::gen());
+            auto lsid = makeLogicalSessionIdWithTxnUUID(parentLsid);
             auto session = InternalSessionPool::Session(lsid, TxnNumber(0));
             return session;
         }
@@ -172,7 +172,7 @@ void InternalSessionPool::release(Session session) {
     const auto& lsid = session.getSessionId();
     if (lsid.getTxnUUID()) {
         stdx::lock_guard<Latch> lock(_mutex);
-        _childSessions.insert({std::move(lsid), std::move(session)});
+        _childSessions.insert({*getParentSessionId(lsid), std::move(session)});
     } else {
         stdx::lock_guard<Latch> lock(_mutex);
 
diff --git a/src/mongo/db/internal_session_pool_test.cpp b/src/mongo/db/internal_session_pool_test.cpp
index 48967f24a58..983fbb7e429 100644
--- a/src/mongo/db/internal_session_pool_test.cpp
+++ b/src/mongo/db/internal_session_pool_test.cpp
@@ -103,24 +103,59 @@ TEST_F(InternalSessionPoolTest, AcquireWithParentSessionFromPoolWithoutParentEnt
     ASSERT_EQ(TxnNumber(0), session.getTxnNumber());
 }
 
+TEST_F(InternalSessionPoolTest, AcquireChildForSameParentWithoutIntermediateRelease) {
+    LogicalSessionId parentLsid = makeLogicalSessionIdForTest();
+
+    auto childLsid1 = makeLogicalSessionIdWithTxnUUID(parentLsid);
+    _pool->release(InternalSessionPool::Session(childLsid1, TxnNumber(0)));
+
+    auto childSession1 = _pool->acquireChildSession(opCtx(), parentLsid);
+
+    ASSERT_NOT_EQUALS(parentLsid, childSession1.getSessionId());
+    ASSERT_EQ(childLsid1, childSession1.getSessionId());
+    ASSERT_EQ(TxnNumber(1), childSession1.getTxnNumber());
+
+    // If we acquire a child for the same parent without first releasing the checked out child, it
+    // shouldn't block or prevent future child session reuse.
+
+    auto childSession2 = _pool->acquireChildSession(opCtx(), parentLsid);
+    ASSERT_NOT_EQUALS(parentLsid, childSession2.getSessionId());
+    ASSERT_NOT_EQUALS(childLsid1, childSession2.getSessionId());
+    ASSERT_EQ(TxnNumber(0), childSession2.getTxnNumber());
+
+    _pool->release(childSession2);
+
+    auto childSession3 = _pool->acquireChildSession(opCtx(), parentLsid);
+    ASSERT_EQ(childSession2.getSessionId(), childSession3.getSessionId());
+    ASSERT_EQ(TxnNumber(1), childSession3.getTxnNumber());
+
+    // Releasing the first child session back into the pool should overwrite the previous session
+    // and still allow for reuse.
+
+    _pool->release(childSession1);
+
+    auto childSession4 = _pool->acquireChildSession(opCtx(), parentLsid);
+    ASSERT_EQ(childSession1.getSessionId(), childSession4.getSessionId());
+    ASSERT_EQ(TxnNumber(2), childSession4.getTxnNumber());
+}
+
 TEST_F(InternalSessionPoolTest, AcquireWithParentSessionFromPoolWithParentEntry) {
     LogicalSessionId parentLsid1 = makeLogicalSessionIdForTest();
     LogicalSessionId parentLsid2 = makeLogicalSessionIdForTest();
 
-    // Set txnUUID for parentLsids.
-    parentLsid1.setTxnUUID(UUID::gen());
-    parentLsid2.setTxnUUID(UUID::gen());
-
-    auto parentSession1 = InternalSessionPool::Session(parentLsid1, TxnNumber(1));
-    _pool->release(parentSession1);
+    // Create child sessions for each parent session and release them in the pool.
+    auto childLsid1 = makeLogicalSessionIdWithTxnUUID(parentLsid1);
+    auto childLsid2 = makeLogicalSessionIdWithTxnUUID(parentLsid2);
 
-    auto parentSession2 = InternalSessionPool::Session(parentLsid2, TxnNumber(2));
-    _pool->release(parentSession2);
+    _pool->release(InternalSessionPool::Session(childLsid1, TxnNumber(1)));
+    _pool->release(InternalSessionPool::Session(childLsid2, TxnNumber(2)));
 
     auto childSession2 = _pool->acquireChildSession(opCtx(), parentLsid2);
 
     ASSERT_NOT_EQUALS(parentLsid1, childSession2.getSessionId());
-    ASSERT_EQ(parentLsid2, childSession2.getSessionId());
+    ASSERT_NOT_EQUALS(childLsid1, childSession2.getSessionId());
+    ASSERT_NOT_EQUALS(parentLsid2, childSession2.getSessionId());
+    ASSERT_EQ(childLsid2, childSession2.getSessionId());
 
     // txnNumber should be 1 larger than the released parent session.
     ASSERT_EQ(TxnNumber(3), childSession2.getTxnNumber());