/**  * Repro for the IFR-flag-pinning-vs-getMore bug: a $$SEARCH_META $search sub-pipeline inside  * $unionWith / $lookup against a SHARDED collection fails when the sub-pipeline is first pulled on a  * getMore rather than inside the original aggregate operation.  *  * Mechanism: with featureFlagSearchExtension on and featureFlagExtensionsInsideHybridSearch off, the  * extension $search stage inside a $unionWith/$lookup sub-pipeline throws an IFR kickback  * (document_source_extension_optimizable.h), so the router retries the aggregate with  * featureFlagSearchExtension pinned to false (cluster_aggregate.cpp). That pinning lives on the  * per-operation IFR context, so a getMore -- a new operation -- runs with the flag back at true . Work  * planned under flag= false then executes under flag= true and the query fails.  *  * Every case below is correct behavior and should pass. Cases that currently fail are marked. This  * is a repro script, not suite coverage: it is expected to fail until the bug is fixed, so do not  * commit it into a suite-selected directory as-is.  *  * Required ingredients: the search collection sharded across >= 2 shards (so the sub-pipeline is  * split and a metadata cursor exists) AND the sub-pipeline first pulled on a getMore.  *  * @tags: [  *   requires_sharding,  *   assumes_unsharded_collection,  *   requires_getmore,  * ]  */ import {getShardNames} from "jstests/libs/cluster_helpers/sharded_cluster_fixture_helpers.js" ; import {after, before, describe, it} from "jstests/libs/mochalite.js" ; import {createSearchIndex, dropSearchIndex} from "jstests/libs/query_integration_search/search.js" ; const kNumDocs = 200; const kIndexName = "subpipeline_getmore_ifr_index" ; const kNumLookups = 2; // 'count' makes $$SEARCH_META resolve to something we can assert on. const kSearchQuery = {     index: kIndexName,     text: {query: "hello" , path: "t" },     count: {type: "total" }, }; // A blocking, memory-tracked stage, so that the outer pipeline must be split and the $unionWith // sub-pipeline cannot be reached until the outer side is drained. const kGroup = {$group: {_id: "$k" , padding: {$max: "$padding" }}}; const kSearchSubPipeline = [{$search: kSearchQuery}, kGroup, {$addFields: {meta: "$$SEARCH_META" }}];describe( "$$SEARCH_META in a sharded sub-pipeline first pulled on a getMore" , function () {     let testDB;     let searchColl;     let outerColl;    const unionPipeline = () => [         kGroup,         {$unionWith: {coll: searchColl.getName(), pipeline: kSearchSubPipeline}},     ];    const lookupPipeline = () => [         kGroup,         {$limit: kNumLookups},         {$lookup: {from: searchColl.getName(), pipeline: kSearchSubPipeline, as: "searchMeta" }},     ];    /**      * Runs 'pipeline' and logs the outcome before asserting, so that a failing run reports the error      * code of each case rather than only the first failure's stack.      */     function run(label, pipeline, options) {         jsTest.log.info( "running case " , {label, options});         try {             const results = outerColl.aggregate(pipeline, options).toArray();             jsTest.log.info( " case succeeded" , {label, numResults: results.length});             return results;         } catch (e) {             jsTest.log.info( " case FAILED" , {label, code: e.code, errmsg: e.message});             throw e;         }     }    // Every document produced by the sub-pipeline carries the metadata of a search matching     // everything, so a query that succeeds but silently loses $$SEARCH_META is caught too.     function assertSearchMetaResolved(docs) {         assert .eq(kNumDocs, docs.length, "expected one group per distinct 'k' " );         for ( const doc of docs) {             assert .eq( Number (doc.meta.count.total), kNumDocs, "unexpected $$SEARCH_META" , {doc});         }     }    before(function () {         testDB = db.getSiblingDB(jsTestName());         searchColl = testDB.getCollection(jsTestName() + "_search" );         outerColl = testDB.getCollection(jsTestName() + "_outer" );         searchColl.drop();         outerColl.drop();        const shardNames = getShardNames(testDB.getMongo());         assert .gte(shardNames.length, 2, "Test requires at least 2 shards" );         assert .commandWorked(             testDB.adminCommand({enableSharding: testDB.getName(), primaryShard: shardNames[0]}),         );        const docs = [];         for (let i = 0; i < kNumDocs; i++) {             docs.push({_id: i, k: i, t: "hello" , padding: "x" .repeat(100)});         }        // Populate 'coll' and put half of it on each shard, so both shards take part in every query.         function populateAndShard(coll) {             assert .commandWorked(coll.insert(docs));             assert .commandWorked(                 testDB.adminCommand({shardCollection: coll.getFullName(), key: {_id: 1}}),             );             assert .commandWorked(                 testDB.adminCommand({split: coll.getFullName(), middle: {_id: kNumDocs / 2}}),             );             assert .commandWorked(                 testDB.adminCommand({                     moveChunk: coll.getFullName(),                     find: {_id: kNumDocs / 2},                     to: shardNames[1],                     _waitForDelete: true ,                 }),             );         }         populateAndShard(searchColl);         populateAndShard(outerColl);        createSearchIndex(searchColl, {name: kIndexName, definition: {mappings: {dynamic: true }}});     });    after(function () {         dropSearchIndex(searchColl, {name: kIndexName});         searchColl.drop();         outerColl.drop();     });    function assertUnionResults(results) {         // One group per distinct 'k' from each collection; only the union side has $$SEARCH_META.         assert .eq(2 * kNumDocs, results.length, "unexpected result count" , {             numResults: results.length,         });         assertSearchMetaResolved(results.filter((doc) => doc.hasOwnProperty( "meta" )));     }    function assertLookupResults(results) {         assert .eq(kNumLookups, results.length, "unexpected result count" , {             numResults: results.length,         });         for ( const doc of results) {             assertSearchMetaResolved(doc.searchMeta);         }     }    // ---------------------------------------------------------------------------------------------     // Control cases: the sub-pipeline is reached inside the original aggregate operation, while the     // kickback retry's pinned flag value is still in effect. These pass.     // ---------------------------------------------------------------------------------------------    it( "$unionWith with a batchSize large enough to finish inside the aggregate" , function () {         // 400 results fit in one batch, so no getMore is ever issued.         assertUnionResults(run( "unionWith batchSize=1000" , unionPipeline(), {batchSize: 1000}));     });    it( "$lookup with the default batchSize" , function () {         // $lookup pulls its sub-pipeline as soon as the first outer document is produced, which         // happens while building the first batch.         assertLookupResults(run( "lookup default batchSize" , lookupPipeline(), {}));     });    // ---------------------------------------------------------------------------------------------     // Failing cases: the sub-pipeline is first pulled on a getMore.     // ---------------------------------------------------------------------------------------------    it( "$unionWith with the default batchSize" , function () {         // The outer $group must be drained before the union side is touched, so the sub-pipeline is         // first pulled on a getMore.         //         // Currently fails: tassert 6448002 "Expected to have already attached a cursor source to the         // pipeline". The sub-pipeline's shard halves are dispatched under flag= true , so they do not         // return a SearchMetaResult cursor, and injectMetaCursor() never attaches a source to the         // $setVariableFromSubPipeline that the flag= false plan put in the merge half.         assertUnionResults(run( "unionWith default batchSize" , unionPipeline(), {}));     });    it( "$unionWith with batchSize 0" , function () {         // batchSize 0 returns a cursor without running anything, so all execution -- including the         // outer $group -- happens on getMores.         assertUnionResults(run( "unionWith batchSize=0" , unionPipeline(), {batchSize: 0}));     });    it( "$lookup with batchSize 0" , function () {         // Same deferral for $lookup: nothing executes until the first getMore.         //         // Currently fails with the raw kickback error surfacing to the client: IFRFlagRetry (479)         // "The $search/$searchMeta extension stage is not supported in a $lookup" . A getMore has no         // retry handler, so the kickback is returned instead of being retried with the flag off.         assertLookupResults(run( "lookup batchSize=0" , lookupPipeline(), {batchSize: 0}));     }); }); Run with buildscripts/resmoke.py run --suites=with_mongot_extension_hybridSearch_disabled_sharded_cluster