SERVER-108341 Add missing detach/reattach calls to DocumentSourceLookup::_resolvedIntrospectionPipeline (#39280) (#39356)

GitOrigin-RevId: 8bcc5b675614ba61eea8482f478403599adc200f

Author: jsteemann

etc/backports_required_for_multiversion_tests.yml

@@ -615,6 +615,8 @@ last-continuous:
     ticket: SERVER-94731
   - test_file: jstests/sharding/sharded_data_distribution.js
     ticket: SERVER-92419
+  - test_file: jstests/sharding/query/union_with_doubly_nested_lookup.js
+    ticket: SERVER-108341
   suites: null
 last-lts:
   all:
@@ -1282,4 +1284,6 @@ last-lts:
     ticket: SERVER-92419
   - test_file: jstests/replsets/initial_sync_source_has_capped_ttl.js
     ticket: SERVER-104771
+  - test_file: jstests/sharding/query/union_with_doubly_nested_lookup.js
+    ticket: SERVER-108341
   suites: null

jstests/sharding/query/union_with_doubly_nested_lookup.js

@@ -0,0 +1,76 @@
+/**
+ * Test that a $unionWith with recursively nested $lookup pipelines that target other shards behave
+ * correctly when the query requires multiple getMore calls to be completed.
+ *
+ * @tags: [
+ *   requires_sharding,
+ *   multiversion_incompatible,
+ * ]
+ */
+
+(function() {
+"use strict";
+
+const name = jsTestName();
+
+const st = new ShardingTest({shards: 2, mongos: 1});
+
+function runTest(conn) {
+    const db = conn.getDB(name);
+    const shardedCollName = `${jsTestName()}_sharded`;
+    const unshardedCollName = `${jsTestName()}_unsharded`;
+
+    assert.commandWorked(db.adminCommand({enableSharding: db.getName()}));
+    assert.commandWorked(
+        db.adminCommand({shardCollection: `${db.getName()}.${shardedCollName}`, key: {_id: 1}}));
+    assert.commandWorked(
+        db.adminCommand({split: `${db.getName()}.${shardedCollName}`, middle: {_id: 1}}));
+
+    const docs = Array.from(
+        {length: 10}, (_, i) => ({_id: 0.2 * i, key: i % 5, arr: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]}));
+    assert.commandWorked(db[shardedCollName].insert(docs));
+    assert.commandWorked(db[unshardedCollName].insert(docs));
+
+    let result = db[unshardedCollName]
+        .aggregate(
+            [
+                {$match: {_id: 0}},
+                {$unwind: "$arr"},
+                {
+                    $unionWith: {
+                        coll: shardedCollName,
+                        pipeline: [{
+                            $lookup: {
+                                from: unshardedCollName,
+                                "let": {localKey: "$key"},
+                                pipeline: [
+                                    {$match: {$expr: {$eq: ["$key", "$$localKey"]}}},
+                                    {
+                                        $lookup: {
+                                            from: unshardedCollName,
+                                            "let": {localKey: "$key"},
+                                            pipeline: [
+                                                {$match: {$expr: {$eq: ["$key", "$$localKey"]}}},
+                                                {$addFields: {computed: {$add: ["$$localKey", 1]}}},
+                                            ],
+                                            as: "nestedMatches",
+                                        }
+                                    },
+                                ],
+                                as: "matches",
+                            },
+                        }]
+                    }
+                }
+            ],
+            // Note: a small batchSize is required here so that we have multiple getMore calls.
+            {batchSize: 2})
+        .toArray();
+
+    assert.eq(20, result.length);
+}
+
+runTest(st.s0);
+
+st.stop();
+}());

src/mongo/db/pipeline/document_source_lookup.cpp

@@ -1246,9 +1246,13 @@ void DocumentSourceLookUp::detachFromOperationContext() {
         // use Pipeline::detachFromOperationContext() to take care of updating '_fromExpCtx->opCtx'.
         _pipeline->detachFromOperationContext();
         invariant(_fromExpCtx->opCtx == nullptr);
-    } else if (_fromExpCtx) {
+    }
+    if (_fromExpCtx) {
         _fromExpCtx->opCtx = nullptr;
     }
+    if (_resolvedIntrospectionPipeline) {
+        _resolvedIntrospectionPipeline->detachFromOperationContext();
+    }
 }
 
 void DocumentSourceLookUp::reattachToOperationContext(OperationContext* opCtx) {
@@ -1257,9 +1261,13 @@ void DocumentSourceLookUp::reattachToOperationContext(OperationContext* opCtx) {
         // use Pipeline::reattachToOperationContext() to take care of updating '_fromExpCtx->opCtx'.
         _pipeline->reattachToOperationContext(opCtx);
         invariant(_fromExpCtx->opCtx == opCtx);
-    } else if (_fromExpCtx) {
+    }
+    if (_fromExpCtx) {
         _fromExpCtx->opCtx = opCtx;
     }
+    if (_resolvedIntrospectionPipeline) {
+        _resolvedIntrospectionPipeline->reattachToOperationContext(opCtx);
+    }
 }
 
 bool DocumentSourceLookUp::validateOperationContext(const OperationContext* opCtx) const {
@@ -1269,9 +1277,15 @@ bool DocumentSourceLookUp::validateOperationContext(const OperationContext* opCt
 
     if (_pipeline) {
         const auto& sources = _pipeline->getSources();
-        return std::all_of(sources.begin(), sources.end(), [opCtx](const auto& s) {
-            return s->validateOperationContext(opCtx);
-        });
+        if (!std::all_of(sources.begin(), sources.end(), [opCtx](const auto& s) {
+                return s->validateOperationContext(opCtx);
+            })) {
+            return false;
+        }
+    }
+    if (_resolvedIntrospectionPipeline &&
+        !_resolvedIntrospectionPipeline->validateOperationContext(opCtx)) {
+        return false;
     }
 
     return true;