path: root/jstests/aggregation/mongos_merge.js

/**
 * Tests that split aggregations whose merge pipelines are eligible to run on mongoS do so, and
 * produce the expected results. Stages which are eligible to merge on mongoS include:
 *
 * - Splitpoints whose merge components are non-blocking, e.g. $skip, $limit, $sort, $sample.
 * - Non-splittable streaming stages, e.g. $match, $project, $unwind.
 * - Blocking stages in cases where 'allowDiskUse' is false, e.g. $group, $bucketAuto.
 *
 * Because wrapping these aggregations in a $facet stage will affect how the pipeline can be merged,
 * and will therefore invalidate the results of the test cases below, we tag this test to prevent it
 * running under the 'aggregation_facet_unwind' passthrough.
 *
 * @tags: [
 *   do_not_wrap_aggregations_in_facets,
 *   requires_spawning_own_processes,
 * ]
 */

(function() {
    load("jstests/libs/profiler.js");         // For profilerHas*OrThrow helper functions.
    load('jstests/libs/geo_near_random.js');  // For GeoNearRandomTest.

    const st = new ShardingTest({shards: 2, mongos: 1, config: 1});

    const mongosDB = st.s0.getDB(jsTestName());
    const mongosColl = mongosDB[jsTestName()];
    const unshardedColl = mongosDB[jsTestName() + "_unsharded"];

    const shard0DB = primaryShardDB = st.shard0.getDB(jsTestName());
    const shard1DB = st.shard1.getDB(jsTestName());

    assert.commandWorked(mongosDB.dropDatabase());

    // Enable profiling on each shard to verify that no $mergeCursors occur.
    assert.commandWorked(shard0DB.setProfilingLevel(2));
    assert.commandWorked(shard1DB.setProfilingLevel(2));

    // Always merge pipelines which cannot merge on mongoS on the primary shard instead, so we know
    // where to check for $mergeCursors.
    assert.commandWorked(
        mongosDB.adminCommand({setParameter: 1, internalQueryAlwaysMergeOnPrimaryShard: true}));

    // Enable sharding on the test DB and ensure its primary is shard0000.
    assert.commandWorked(mongosDB.adminCommand({enableSharding: mongosDB.getName()}));
    st.ensurePrimaryShard(mongosDB.getName(), "shard0000");

    // Shard the test collection on _id.
    assert.commandWorked(
        mongosDB.adminCommand({shardCollection: mongosColl.getFullName(), key: {_id: 1}}));

    // We will need to test $geoNear on this collection, so create a 2dsphere index.
    assert.commandWorked(mongosColl.createIndex({geo: "2dsphere"}));

    // We will test that $textScore metadata is not propagated to the user, so create a text index.
    assert.commandWorked(mongosColl.createIndex({text: "text"}));

    // Split the collection into 4 chunks: [MinKey, -100), [-100, 0), [0, 100), [100, MaxKey).
    assert.commandWorked(
        mongosDB.adminCommand({split: mongosColl.getFullName(), middle: {_id: -100}}));
    assert.commandWorked(
        mongosDB.adminCommand({split: mongosColl.getFullName(), middle: {_id: 0}}));
    assert.commandWorked(
        mongosDB.adminCommand({split: mongosColl.getFullName(), middle: {_id: 100}}));

    // Move the [0, 100) and [100, MaxKey) chunks to shard0001.
    assert.commandWorked(mongosDB.adminCommand(
        {moveChunk: mongosColl.getFullName(), find: {_id: 50}, to: "shard0001"}));
    assert.commandWorked(mongosDB.adminCommand(
        {moveChunk: mongosColl.getFullName(), find: {_id: 150}, to: "shard0001"}));

    // Create a random geo co-ord generator for testing.
    var georng = new GeoNearRandomTest(mongosColl);

    // Write 400 documents across the 4 chunks.
    for (let i = -200; i < 200; i++) {
        assert.writeOK(mongosColl.insert(
            {_id: i, a: [i], b: {redactThisDoc: true}, c: true, geo: georng.mkPt(), text: "txt"}));
        assert.writeOK(unshardedColl.insert({_id: i, x: i}));
    }

    let testNameHistory = new Set();

    /**
     * Runs the aggregation specified by 'pipeline', verifying that:
     * - The number of documents returned by the aggregation matches 'expectedCount'.
     * - The merge was performed on a mongoS if 'mergeType' is 'mongos', and on a shard otherwise.
     */
    function assertMergeBehaviour(
        {testName, pipeline, mergeType, batchSize, allowDiskUse, expectedCount}) {
        // Ensure that this test has a unique name.
        assert(!testNameHistory.has(testName));
        testNameHistory.add(testName);

        // Create the aggregation options from the given arguments.
        const opts = {
            comment: testName,
            cursor: (batchSize ? {batchSize: batchSize} : {}),
        };

        if (allowDiskUse !== undefined) {
            opts.allowDiskUse = allowDiskUse;
        }

        // Verify that the explain() output's 'mergeType' field matches our expectation.
        assert.eq(
            assert.commandWorked(mongosColl.explain().aggregate(pipeline, Object.extend({}, opts)))
                .mergeType,
            mergeType);

        // Verify that the aggregation returns the expected number of results.
        assert.eq(mongosColl.aggregate(pipeline, opts).itcount(), expectedCount);

        // Verify that a $mergeCursors aggregation ran on the primary shard if 'mergeType' is not
        // 'mongos', and that no such aggregation ran otherwise.
        profilerHasNumMatchingEntriesOrThrow({
            profileDB: primaryShardDB,
            numExpectedMatches: (mergeType === "mongos" ? 0 : 1),
            filter: {
                "command.aggregate": mongosColl.getName(),
                "command.comment": testName,
                "command.pipeline.$mergeCursors": {$exists: 1}
            }
        });
    }

    /**
     * Throws an assertion if the aggregation specified by 'pipeline' does not produce
     * 'expectedCount' results, or if the merge phase is not performed on the mongoS.
     */
    function assertMergeOnMongoS({testName, pipeline, batchSize, allowDiskUse, expectedCount}) {
        assertMergeBehaviour({
            testName: testName,
            pipeline: pipeline,
            mergeType: "mongos",
            batchSize: (batchSize || 10),
            allowDiskUse: allowDiskUse,
            expectedCount: expectedCount
        });
    }

    /**
     * Throws an assertion if the aggregation specified by 'pipeline' does not produce
     * 'expectedCount' results, or if the merge phase was not performed on a shard.
     */
    function assertMergeOnMongoD(
        {testName, pipeline, mergeType, batchSize, allowDiskUse, expectedCount}) {
        assertMergeBehaviour({
            testName: testName,
            pipeline: pipeline,
            mergeType: (mergeType || "anyShard"),
            batchSize: (batchSize || 10),
            allowDiskUse: allowDiskUse,
            expectedCount: expectedCount
        });
    }

    /**
     * Runs a series of test cases which will consistently merge on mongoS or mongoD regardless of
     * whether 'allowDiskUse' is true, false or omitted.
     */
    function runTestCasesWhoseMergeLocationIsConsistentRegardlessOfAllowDiskUse(allowDiskUse) {
        // Test that a $match pipeline with an empty merge stage is merged on mongoS.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_match_only",
            pipeline: [{$match: {_id: {$gte: -200, $lte: 200}}}],
            allowDiskUse: allowDiskUse,
            expectedCount: 400
        });

        // Test that a $sort stage which merges pre-sorted streams is run on mongoS.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_sort_presorted",
            pipeline: [{$match: {_id: {$gte: -200, $lte: 200}}}, {$sort: {_id: -1}}],
            allowDiskUse: allowDiskUse,
            expectedCount: 400
        });

        // Test that $skip is merged on mongoS.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_skip",
            pipeline: [{$match: {_id: {$gte: -200, $lte: 200}}}, {$sort: {_id: -1}}, {$skip: 300}],
            allowDiskUse: allowDiskUse,
            expectedCount: 100
        });

        // Test that $limit is merged on mongoS.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_limit",
            pipeline: [{$match: {_id: {$gte: -200, $lte: 200}}}, {$limit: 300}],
            allowDiskUse: allowDiskUse,
            expectedCount: 300
        });

        // Test that $sample is merged on mongoS if it is the splitpoint, since this will result in
        // a merging $sort of presorted streams in the merge pipeline.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_sample_splitpoint",
            pipeline: [{$match: {_id: {$gte: -200, $lte: 200}}}, {$sample: {size: 300}}],
            allowDiskUse: allowDiskUse,
            expectedCount: 300
        });

        // Test that $geoNear is merged on mongoS.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_geo_near",
            pipeline: [
                {$geoNear: {near: [0, 0], distanceField: "distance", spherical: true, limit: 300}}
            ],
            allowDiskUse: allowDiskUse,
            expectedCount: 300
        });

        // Test that a pipeline whose merging half can be run on mongos using only the mongos
        // execution machinery returns the correct results.
        // TODO SERVER-30882 Find a way to assert that all stages get absorbed by mongos.
        assertMergeOnMongoS({
            testName: "agg_mongos_merge_all_mongos_runnable_skip_and_limit_stages",
            pipeline: [
                {$match: {_id: {$gte: -200, $lte: 200}}},
                {$sort: {_id: -1}},
                {$skip: 150},
                {$limit: 150},
                {$skip: 5},
                {$limit: 1},
            ],
            allowDiskUse: allowDiskUse,
            expectedCount: 1
        });

        // Test that a merge pipeline which needs to run on a shard is NOT merged on mongoS
        // regardless of 'allowDiskUse'.
        assertMergeOnMongoD({
            testName: "agg_mongos_merge_primary_shard_disk_use_" + allowDiskUse,
            pipeline: [
                {$match: {_id: {$gte: -200, $lte: 200}}},
                {$_internalSplitPipeline: {mergeType: "anyShard"}}
            ],
            mergeType: "anyShard",
            allowDiskUse: allowDiskUse,
            expectedCount: 400
        });
    }

    /**
     * Runs a series of test cases which will always merge on mongoD when 'allowDiskUse' is true,
     * and on mongoS when 'allowDiskUse' is false or omitted.
     */
    function runTestCasesWhoseMergeLocationDependsOnAllowDiskUse(allowDiskUse) {
        // All test cases should merge on mongoD if allowDiskUse is true, mongoS otherwise.
        const assertMergeOnMongoX = (allowDiskUse ? assertMergeOnMongoD : assertMergeOnMongoS);

        // Test that a blocking $sort is only merged on mongoS if 'allowDiskUse' is not set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_blocking_sort_no_disk_use",
            pipeline:
                [{$match: {_id: {$gte: -200, $lte: 200}}}, {$sort: {_id: -1}}, {$sort: {a: 1}}],
            allowDiskUse: allowDiskUse,
            expectedCount: 400
        });

        // Test that $group is only merged on mongoS if 'allowDiskUse' is not set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_group_allow_disk_use",
            pipeline:
                [{$match: {_id: {$gte: -200, $lte: 200}}}, {$group: {_id: {$mod: ["$_id", 150]}}}],
            allowDiskUse: allowDiskUse,
            expectedCount: 299
        });

        // Test that a blocking $sample is only merged on mongoS if 'allowDiskUse' is not set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_blocking_sample_allow_disk_use",
            pipeline: [
                {$match: {_id: {$gte: -200, $lte: 200}}},
                {$sample: {size: 300}},
                {$sample: {size: 200}}
            ],
            allowDiskUse: allowDiskUse,
            expectedCount: 200
        });

        // Test that $bucketAuto is only merged on mongoS if 'allowDiskUse' is not set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_bucket_auto_allow_disk_use",
            pipeline: [
                {$match: {_id: {$gte: -200, $lte: 200}}},
                {$bucketAuto: {groupBy: "$_id", buckets: 10}}
            ],
            allowDiskUse: allowDiskUse,
            expectedCount: 10
        });

        //
        // Test composite stages.
        //

        // Test that $bucket ($group->$sort) is merged on mongoS iff 'allowDiskUse' is not set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_bucket_allow_disk_use",
            pipeline: [
                {$match: {_id: {$gte: -200, $lte: 200}}},
                {
                  $bucket: {
                      groupBy: "$_id",
                      boundaries: [-200, -150, -100, -50, 0, 50, 100, 150, 200]
                  }
                }
            ],
            allowDiskUse: allowDiskUse,
            expectedCount: 8
        });

        // Test that $sortByCount ($group->$sort) is merged on mongoS iff 'allowDiskUse' isn't set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_sort_by_count_allow_disk_use",
            pipeline:
                [{$match: {_id: {$gte: -200, $lte: 200}}}, {$sortByCount: {$mod: ["$_id", 150]}}],
            allowDiskUse: allowDiskUse,
            expectedCount: 299
        });

        // Test that $count ($group->$project) is merged on mongoS iff 'allowDiskUse' is not set.
        assertMergeOnMongoX({
            testName: "agg_mongos_merge_count_allow_disk_use",
            pipeline: [{$match: {_id: {$gte: -150, $lte: 1500}}}, {$count: "doc_count"}],
            allowDiskUse: allowDiskUse,
            expectedCount: 1
        });
    }

    // Run all test cases for each potential value of 'allowDiskUse'.
    for (let allowDiskUse of[false, undefined, true]) {
        runTestCasesWhoseMergeLocationIsConsistentRegardlessOfAllowDiskUse(allowDiskUse);
        runTestCasesWhoseMergeLocationDependsOnAllowDiskUse(allowDiskUse);
        testNameHistory.clear();
    }

    // Test that merge pipelines containing all mongos-runnable stages produce the expected output.
    assertMergeOnMongoS({
        testName: "agg_mongos_merge_all_mongos_runnable_stages",
        pipeline: [
            {$geoNear: {near: [0, 0], distanceField: "distance", spherical: true, limit: 400}},
            {$sort: {a: 1}},
            {$skip: 150},
            {$limit: 150},
            {$addFields: {d: true}},
            {$unwind: "$a"},
            {$sample: {size: 100}},
            {$project: {c: 0, geo: 0, distance: 0}},
            {$group: {_id: "$_id", doc: {$push: "$$CURRENT"}}},
            {$unwind: "$doc"},
            {$replaceRoot: {newRoot: "$doc"}},
            {
              $redact: {
                  $cond:
                      {if: {$eq: ["$redactThisDoc", true]}, then: "$$PRUNE", else: "$$DESCEND"}
              }
            },
            {
              $match: {
                  _id: {$gte: -50, $lte: 100},
                  a: {$type: "number", $gte: -50, $lte: 100},
                  b: {$exists: false},
                  c: {$exists: false},
                  d: true,
                  geo: {$exists: false},
                  distance: {$exists: false},
                  text: "txt"
              }
            }
        ],
        expectedCount: 100
    });

    // Test that metadata is not propagated to the user when a pipeline which produces metadata
    // fields merges on mongoS.
    const metaDataTests = [
        {pipeline: [{$sort: {_id: -1}}], verifyNoMetaData: (doc) => assert.isnull(doc.$sortKey)},
        {
          pipeline: [{$match: {$text: {$search: "txt"}}}],
          verifyNoMetaData: (doc) => assert.isnull(doc.$textScore)
        },
        {
          pipeline: [{$sample: {size: 300}}],
          verifyNoMetaData: (doc) => assert.isnull(doc.$randVal)
        },
        {
          pipeline: [{$match: {$text: {$search: "txt"}}}, {$sort: {text: 1}}],
          verifyNoMetaData:
              (doc) => assert.docEq([doc.$textScore, doc.$sortKey], [undefined, undefined])
        }
    ];

    for (let metaDataTest of metaDataTests) {
        assert.gte(mongosColl.aggregate(metaDataTest.pipeline).itcount(), 300);
        mongosColl.aggregate(metaDataTest.pipeline).forEach(metaDataTest.verifyNoMetaData);
    }

    st.stop();
})();