/*
 * Test that spilling to disk in $setWindowFields works and returns the correct results.
 * @tags: [
 * requires_profiling,
 * assumes_read_concern_unchanged,
 * ]
 */
(function() {
"use strict";

load("jstests/noPassthrough/libs/server_parameter_helpers.js");  // For setParameterOnAllHosts.
load("jstests/libs/discover_topology.js");                       // For findNonConfigNodes.
load("jstests/aggregation/extras/window_function_helpers.js");
load("jstests/libs/analyze_plan.js");         // For getAggPlanStages().
load("jstests/aggregation/extras/utils.js");  // arrayEq.
load("jstests/libs/profiler.js");             // getLatestProfileEntry.

const origParamValue = assert.commandWorked(db.adminCommand({
    getParameter: 1,
    internalDocumentSourceSetWindowFieldsMaxMemoryBytes: 1
}))["internalDocumentSourceSetWindowFieldsMaxMemoryBytes"];
const coll = db[jsTestName()];
coll.drop();
// Doc size was found through logging the size in the SpillableCache. Partition sizes were chosen
// arbitrarily.
let avgDocSize = 274;
let smallPartitionSize = 6;
let largePartitionSize = 21;
setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()),
                       "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
                       avgDocSize * smallPartitionSize + 1);

seedWithTickerData(coll, 10);

// Run $sum test with memory limits that cause spilling to disk.
testAccumAgainstGroup(coll, "$sum", 0);

function checkProfilerForDiskWrite(dbToCheck) {
    if (!FixtureHelpers.isMongos) {
        const profileObj = getLatestProfilerEntry(dbToCheck, {
            $or: [
                {originatingCommand: {pipeline: {$setWindowFields: {$exists: true}}}},
                {command: {pipeline: {$setWindowFields: {$exists: true}}}}
            ]
        });
        assert(profileObj.usedDisk, tojson(profileObj));
    }
}

FixtureHelpers.runCommandOnEachPrimary({db: db, cmdObj: {profile: 2}});
// Test that a query that spills to disk succeeds across getMore requests.
const wfResults =
    coll.aggregate(
            [
                {
                    $setWindowFields: {
                        sortBy: {_id: 1},
                        output: {res: {$sum: "$price", window: {documents: ["unbounded", 5]}}}
                    },
                },
            ],
            {allowDiskUse: true, cursor: {batchSize: 1}})
        .toArray();
assert.eq(wfResults.length, 20);
checkProfilerForDiskWrite(db);

// Test a small, in memory, partition and a larger partition that requires spilling to disk.
coll.drop();
// Create small partition.
for (let i = 0; i < smallPartitionSize; i++) {
    assert.commandWorked(coll.insert({_id: i, val: i, partition: 1}));
}
// Create large partition.
for (let i = 0; i < largePartitionSize; i++) {
    assert.commandWorked(coll.insert({_id: i + smallPartitionSize, val: i, partition: 2}));
}

// Run an aggregation that will keep all documents in the cache for all documents.
let results =
    coll.aggregate(
            [
                {
                    $setWindowFields: {
                        partitionBy: "$partition",
                        sortBy: {partition: 1},
                        output: {
                            sum: {
                                $sum: "$val",
                                window: {documents: [-largePartitionSize, largePartitionSize]}
                            }
                        }
                    }
                },
                {$sort: {_id: 1}}
            ],
            {allowDiskUse: true})
        .toArray();
for (let i = 0; i < results.length; i++) {
    if (results[i].partition === 1) {
        assert.eq(results[i].sum, 15, "Unexepected result in first partition at position " + i);
    } else {
        assert.eq(results[i].sum, 210, "Unexepcted result in second partition at position " + i);
    }
}
checkProfilerForDiskWrite(db);

// We don't execute setWindowFields in a sharded explain.
if (!FixtureHelpers.isMongos(db)) {
    // Test that an explain that executes the query reports usedDisk correctly.
    let explainPipeline = [
        {
            $setWindowFields: {
                partitionBy: "$partition",
                sortBy: {partition: 1},
                output: {arr: {$sum: "$val", window: {documents: [-21, 21]}}}
            }
        },
        {$sort: {_id: 1}}
    ];

    let stages = getAggPlanStages(
        coll.explain("allPlansExecution").aggregate(explainPipeline, {allowDiskUse: true}),
        "$_internalSetWindowFields");
    assert(stages[0]["usedDisk"], stages);
    setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()),
                           "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
                           avgDocSize * largePartitionSize * 2);
    explainPipeline = [
        {
            $setWindowFields: {
                partitionBy: "$partition",
                sortBy: {partition: 1},
                output: {arr: {$sum: "$val", window: {documents: [0, 0]}}}
            }
        },
        {$sort: {_id: 1}}
    ];

    stages = getAggPlanStages(
        coll.explain("allPlansExecution").aggregate(explainPipeline, {allowDiskUse: true}),
        "$_internalSetWindowFields");
    assert(!stages[0]["usedDisk"], stages);
}

// Run an aggregation that will store too many documents in the function and force a spill. Set the
// memory limit to be over the size of the large partition.
setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()),
                       "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
                       largePartitionSize * avgDocSize + 1);
results = coll.aggregate(
                  [
                      {
                          $setWindowFields: {
                              partitionBy: "$partition",
                              sortBy: {partition: 1},
                              output: {arr: {$push: "$val", window: {documents: [-25, 25]}}}
                          }
                      },
                      {$sort: {_id: 1}}
                  ],
                  {allowDiskUse: true})
              .toArray();
checkProfilerForDiskWrite(db);
for (let i = 0; i < results.length; i++) {
    if (results[i].partition === 1) {
        assert(arrayEq(results[i].arr, [0, 1, 2, 3, 4, 5]),
               "Unexepected result in first partition at position " + i);
    } else {
        assert(arrayEq(results[i].arr,
                       [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]),
               "Unexepcted result in second partition at position " + i);
    }
}

// Check that if function memory limit exceeds we fail even though the partition iterator spilled.
// $push uses about ~950 to store all the values in the second partition.
setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()),
                       "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
                       avgDocSize * 2);
assert.commandFailedWithCode(db.runCommand({
    aggregate: coll.getName(),
    pipeline: [
        {
            $setWindowFields: {
                partitionBy: "$partition",
                sortBy: {partition: 1},
                output: {arr: {$push: "$val", window: {documents: [-21, 21]}}}
            }
        },
        {$sort: {_id: 1}}
    ],
    allowDiskUse: true,
    cursor: {}
}),
                             5414201);

coll.drop();
// Test that situations that would require a large spill successfully write to disk.
// Set the limit to spill after ~1000 documents since that is the batch size when we write to disk.
setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()),
                       "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
                       1000 * avgDocSize);
let numDocs = 1111;
let batchArr = [];
for (let docNum = 0; docNum < numDocs; docNum++) {
    batchArr.push({_id: docNum, val: docNum, partition: 1});
}
assert.commandWorked(coll.insert(batchArr));
// Run a document window over the whole collection to keep everything in the cache.
results =
    coll.aggregate(
            [
                {
                    $setWindowFields: {
                        // partitionBy: "$partition",
                        sortBy: {partition: 1},
                        output: {arr: {$sum: "$val", window: {documents: [-numDocs, numDocs]}}}
                    }
                },
                {$sort: {_id: 1}}
            ],
            {allowDiskUse: true})
        .toArray();
checkProfilerForDiskWrite(db);
// Check that the command succeeded.
assert.eq(results.length, numDocs);
for (let i = 0; i < numDocs; i++) {
    assert.eq(results[i].arr, 616605, results);
}
// Reset limit for other tests.
setParameterOnAllHosts(DiscoverTopology.findNonConfigNodes(db.getMongo()),
                       "internalDocumentSourceSetWindowFieldsMaxMemoryBytes",
                       origParamValue);
// Reset profiler.
FixtureHelpers.runCommandOnEachPrimary({db: db, cmdObj: {profile: 0}});
})();
