/**
 * This file tests that commands that do writes accept a write concern in a sharded cluster. This
 * test defines various database commands and what they expect to be true before and after the fact.
 * It then runs the commands with an invalid writeConcern and a valid writeConcern and
 * ensures that they succeed and fail appropriately. This only tests functions that aren't run
 * on config servers.
 *
 * This test is labeled resource intensive because its total io_write is 58MB compared to a median
 * of 5MB across all sharding tests in wiredTiger.
 * @tags: [
 *   resource_intensive,
 *   requires_scripting
 * ]
 */
import {ShardingTest} from "jstests/libs/shardingtest.js";
import {
    assertWriteConcernError,
    runCommandCheckAdmin,
    shardCollectionWithChunks,
} from "jstests/libs/write_concern_util.js";

let st = new ShardingTest({
    // Set priority of secondaries to zero to prevent spurious elections.
    shards: {
        rs0: {
            nodes: [{}, {rsConfig: {priority: 0}}, {rsConfig: {priority: 0}}],
            settings: {chainingAllowed: false},
        },
        rs1: {
            nodes: [{}, {rsConfig: {priority: 0}}, {rsConfig: {priority: 0}}],
            settings: {chainingAllowed: false},
        },
    },
    configReplSetTestOptions: {settings: {chainingAllowed: false}},
    mongos: 1,
});

let mongos = st.s;
let dbName = "wc-test-shards";
var db = mongos.getDB(dbName);
let collName = "leaves";
let coll = db[collName];

function dropTestDatabase() {
    db.runCommand({dropDatabase: 1});
    db.extra.insert({a: 1});
    coll = db[collName];
    assert.eq(0, coll.find().itcount(), "test collection not empty");
    assert.eq(1, db.extra.find().itcount(), "extra collection should have 1 document");
}

let commands = [];

// Tests a runOnAllShardsCommand against a sharded collection.
commands.push({
    req: {createIndexes: collName, indexes: [{key: {"type": 1}, name: "sharded_type_index"}]},
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
        coll.insert({type: "oak", x: -3});
        coll.insert({type: "maple", x: 23});
        assert.eq(coll.getIndexes().length, 2);
    },
    confirmFunc: function () {
        assert.eq(coll.getIndexes().length, 3);
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

// Tests a runOnAllShardsCommand.
commands.push({
    req: {createIndexes: collName, indexes: [{key: {"type": 1}, name: "type_index"}]},
    setupFunc: function () {
        coll.insert({type: "oak"});
        assert.eq(coll.getIndexes().length, 1);
    },
    confirmFunc: function () {
        assert.eq(coll.getIndexes().length, 2);
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

// Tests a batched write command.
commands.push({
    req: {
        insert: collName,
        documents: [
            {x: -3, type: "maple"},
            {x: 23, type: "maple"},
        ],
    },
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
    },
    confirmFunc: function () {
        assert.eq(coll.count({type: "maple"}), 2);
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

commands.push({
    req: {renameCollection: "renameCollWC.leaves", to: "renameCollWC.pine_needles"},
    setupFunc: function () {
        db = db.getSiblingDB("renameCollWC");
        // Ensure that database is created.
        db.leaves.insert({type: "oak"});
        db.leaves.drop();
        db.pine_needles.drop();
        db.leaves.insert({type: "oak"});
        assert.eq(db.leaves.count(), 1);
        assert.eq(db.pine_needles.count(), 0);
    },
    confirmFunc: function () {
        assert.eq(db.leaves.count(), 0);
        assert.eq(db.pine_needles.count(), 1);
    },
    admin: true,
    isExpectedToWriteOnWriteConcernFailure: true,
});

commands.push({
    req: {
        update: collName,
        updates: [
            {
                q: {type: "oak"},
                u: [{$set: {type: "ginkgo"}}],
            },
        ],
        writeConcern: {w: "majority"},
    },
    setupFunc: function () {
        coll.insert({type: "oak"});
        assert.eq(coll.count({type: "ginkgo"}), 0);
        assert.eq(coll.count({type: "oak"}), 1);
    },
    confirmFunc: function () {
        assert.eq(coll.count({type: "ginkgo"}), 1);
        assert.eq(coll.count({type: "oak"}), 0);
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

commands.push({
    req: {
        findAndModify: collName,
        query: {type: "oak"},
        update: {$set: {type: "ginkgo"}},
        writeConcern: {w: "majority"},
    },
    setupFunc: function () {
        coll.insert({type: "oak"});
        assert.eq(coll.count({type: "ginkgo"}), 0);
        assert.eq(coll.count({type: "oak"}), 1);
    },
    confirmFunc: function () {
        assert.eq(coll.count({type: "ginkgo"}), 1);
        assert.eq(coll.count({type: "oak"}), 0);
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

commands.push({
    req: {
        findAndModify: collName,
        query: {type: "oak"},
        update: [{$set: {type: "ginkgo"}}],
        writeConcern: {w: "majority"},
    },
    setupFunc: function () {
        coll.insert({type: "oak"});
        assert.eq(coll.count({type: "ginkgo"}), 0);
        assert.eq(coll.count({type: "oak"}), 1);
    },
    confirmFunc: function () {
        assert.eq(coll.count({type: "ginkgo"}), 1);
        assert.eq(coll.count({type: "oak"}), 0);
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

// Aggregate run against an unsharded collection.
commands.push({
    req: {aggregate: collName, pipeline: [{$sort: {type: 1}}, {$out: "foo"}], cursor: {}},
    setupFunc: function () {
        coll.insert({_id: 1, x: 3, type: "oak"});
        coll.insert({_id: 2, x: 13, type: "maple"});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.find({type: "oak"}).itcount(), 1);
            assert.eq(db.foo.find({type: "maple"}).itcount(), 1);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    // Aggregation pipelines with a $out stage are expected not to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

// Aggregate that only matches one shard.
commands.push({
    req: {
        aggregate: collName,
        pipeline: [{$match: {x: -3}}, {$match: {type: {$exists: 1}}}, {$out: "foo"}],
        cursor: {},
    },
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
        coll.insert({_id: 1, x: -3, type: "oak"});
        coll.insert({_id: 2, x: -4, type: "maple"});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.find().itcount(), 1);
            assert.eq(db.foo.find({type: "oak"}).itcount(), 1);
            assert.eq(db.foo.find({type: "maple"}).itcount(), 0);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    // Aggregation pipelines with a $out stage are expected not to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

// Aggregate that must go to multiple shards.
commands.push({
    req: {
        aggregate: collName,
        pipeline: [{$match: {type: {$exists: 1}}}, {$sort: {type: 1}}, {$out: "foo"}],
        cursor: {},
    },
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
        coll.insert({_id: 1, x: -3, type: "oak"});
        coll.insert({_id: 2, x: 23, type: "maple"});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.find().itcount(), 2);
            assert.eq(db.foo.find({type: "oak"}).itcount(), 1);
            assert.eq(db.foo.find({type: "maple"}).itcount(), 1);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    // Aggregation pipelines with a $out stage are expected not to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

const map = function () {
    if (!this.tags) {
        return;
    }
    this.tags.forEach(function (z) {
        emit(z, 1);
    });
};

const reduce = function (key, values) {
    let count = 0;
    values.forEach(function (v) {
        count = count + v;
    });
    return count;
};

// MapReduce on an unsharded collection.
commands.push({
    req: {mapReduce: collName, map: map, reduce: reduce, out: "foo"},
    setupFunc: function () {
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.findOne({_id: "a"}).value, 2);
            assert.eq(db.foo.findOne({_id: "b"}).value, 3);
            assert.eq(db.foo.findOne({_id: "c"}).value, 3);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    // MapReduce with a replace output action is not expected to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

// MapReduce on an unsharded collection with an output to a sharded collection.
commands.push({
    req: {mapReduce: collName, map: map, reduce: reduce, out: {merge: "foo", sharded: true}},
    setupFunc: function () {
        db.adminCommand({enablesharding: db.toString()});
        assert.commandWorked(db.foo.createIndex({_id: "hashed"}));
        st.shardColl(db.foo, {_id: "hashed"}, false);
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.findOne({_id: "a"}).value, 2);
            assert.eq(db.foo.findOne({_id: "b"}).value, 3);
            assert.eq(db.foo.findOne({_id: "c"}).value, 3);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

// MapReduce on a sharded collection.
commands.push({
    req: {mapReduce: collName, map: map, reduce: reduce, out: "foo"},
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.findOne({_id: "a"}).value, 2);
            assert.eq(db.foo.findOne({_id: "b"}).value, 3);
            assert.eq(db.foo.findOne({_id: "c"}).value, 3);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    // MapReduce with a replace output action expected not to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

// MapReduce from a sharded collection with merge to a sharded collection.
commands.push({
    req: {mapReduce: collName, map: map, reduce: reduce, out: {merge: "foo", sharded: true}},
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
        assert.commandWorked(db.foo.createIndex({_id: "hashed"}));
        st.shardColl(db.foo, {_id: "hashed"}, false);
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        if (writesExpected) {
            assert.eq(db.foo.findOne({_id: "a"}).value, 2);
            assert.eq(db.foo.findOne({_id: "b"}).value, 3);
            assert.eq(db.foo.findOne({_id: "c"}).value, 3);
            db.foo.drop();
        } else {
            assert.eq(0, db.foo.find().itcount());
        }
    },
    admin: false,
    // When output action is "merge" writes to the output collection are expected to succeed,
    // regardless of writeConcern error.
    isExpectedToWriteOnWriteConcernFailure: true,
});

// MapReduce on an unsharded collection, output to different database.
commands.push({
    req: {mapReduce: collName, map: map, reduce: reduce, out: {replace: "foo", db: "other"}},
    setupFunc: function () {
        db.getSiblingDB("other").createCollection("foo");
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        const otherDB = db.getSiblingDB("other");
        if (writesExpected) {
            assert.eq(otherDB.foo.findOne({_id: "a"}).value, 2);
            assert.eq(otherDB.foo.findOne({_id: "b"}).value, 3);
            assert.eq(otherDB.foo.findOne({_id: "c"}).value, 3);
            otherDB.foo.drop();
        } else {
            assert.eq(0, otherDB.foo.find().itcount());
        }
    },
    admin: false,
    // MapReduce with a replace output action is not expected to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

// MapReduce on a sharded collection, output to a different database.
commands.push({
    req: {mapReduce: collName, map: map, reduce: reduce, out: {replace: "foo", db: "other"}},
    setupFunc: function () {
        db.getSiblingDB("other").createCollection("foo");
        shardCollectionWithChunks(st, coll);
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        const otherDB = db.getSiblingDB("other");
        if (writesExpected) {
            assert.eq(otherDB.foo.findOne({_id: "a"}).value, 2);
            assert.eq(otherDB.foo.findOne({_id: "b"}).value, 3);
            assert.eq(otherDB.foo.findOne({_id: "c"}).value, 3);
            otherDB.foo.drop();
        } else {
            assert.eq(0, otherDB.foo.find().itcount());
        }
    },
    admin: false,
    // MapReduce with a replace output action expected not to write on write concern error.
    isExpectedToWriteOnWriteConcernFailure: false,
});

// MapReduce on an unsharded collection with merge to a sharded collection in a different db.
commands.push({
    req: {
        mapReduce: collName,
        map: map,
        reduce: reduce,
        out: {merge: "foo", db: "other", sharded: true},
    },
    setupFunc: function () {
        const otherDB = db.getSiblingDB("other");
        otherDB.createCollection("foo");
        otherDB.adminCommand({enablesharding: otherDB.toString()});
        assert.commandWorked(otherDB.foo.createIndex({_id: "hashed"}));
        st.shardColl(otherDB.foo, {_id: "hashed"}, false);
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        const otherDB = db.getSiblingDB("other");
        if (writesExpected) {
            assert.eq(otherDB.foo.findOne({_id: "a"}).value, 2);
            assert.eq(otherDB.foo.findOne({_id: "b"}).value, 3);
            assert.eq(otherDB.foo.findOne({_id: "c"}).value, 3);
            otherDB.foo.drop();
        } else {
            assert.eq(0, otherDB.foo.find().itcount());
        }
    },
    admin: false,
    isExpectedToWriteOnWriteConcernFailure: true,
});

// MapReduce from a sharded collection with merge to a sharded collection in different db.
commands.push({
    req: {
        mapReduce: collName,
        map: map,
        reduce: reduce,
        out: {merge: "foo", db: "other", sharded: true},
    },
    setupFunc: function () {
        shardCollectionWithChunks(st, coll);
        const otherDB = db.getSiblingDB("other");
        otherDB.createCollection("foo");
        assert.commandWorked(otherDB.foo.createIndex({_id: "hashed"}));
        st.shardColl(otherDB.foo, {_id: "hashed"}, false);
        coll.insert({x: -3, tags: ["a", "b"]});
        coll.insert({x: -7, tags: ["b", "c"]});
        coll.insert({x: 23, tags: ["c", "a"]});
        coll.insert({x: 27, tags: ["b", "c"]});
    },
    confirmFunc: function (writesExpected = true) {
        const otherDB = db.getSiblingDB("other");
        if (writesExpected) {
            assert.eq(otherDB.foo.findOne({_id: "a"}).value, 2);
            assert.eq(otherDB.foo.findOne({_id: "b"}).value, 3);
            assert.eq(otherDB.foo.findOne({_id: "c"}).value, 3);
            otherDB.foo.drop();
        } else {
            assert.eq(0, otherDB.foo.find().itcount());
        }
    },
    admin: false,
    // When output action is "merge" writes to the output collection are expected to succeed,
    // regardless of writeConcern error.
    isExpectedToWriteOnWriteConcernFailure: true,
});

function testValidWriteConcern(cmd) {
    cmd.req.writeConcern = {w: "majority", wtimeout: 5 * 60 * 1000};
    jsTest.log("Testing " + tojson(cmd.req));

    dropTestDatabase();
    cmd.setupFunc();
    let res = runCommandCheckAdmin(db, cmd);
    assert.commandWorked(res);
    assert(!res.writeConcernError, "command on a full cluster had writeConcernError: " + tojson(res));
    cmd.confirmFunc();
}

function testInvalidWriteConcern(cmd) {
    cmd.req.writeConcern = {w: "invalid"};
    jsTest.log("Testing " + tojson(cmd.req));

    dropTestDatabase();
    cmd.setupFunc();
    let res = runCommandCheckAdmin(db, cmd);

    // When a cluster aggregate command with a write stage fails due to write concern error it will
    // return a regular error response with "ok: 0" rather than "ok: 1" with a `writeConcernError`
    // sub-object detailing the failure. This is motivated by the fact that the number of documents
    // written is bounded only by the number of documents produced by the pipeline, which could lead
    // a writeConcernError object to exceed maximum BSON size.
    if (cmd.req.aggregate || cmd.req.mapReduce) {
        assert.commandFailedWithCode(res, [ErrorCodes.WriteConcernTimeout, ErrorCodes.UnknownReplWriteConcern]);

        cmd.confirmFunc(cmd.isExpectedToWriteOnWriteConcernFailure);
    } else if (cmd.req.renameCollection) {
        // The renameCollection spans multiple nodes and potentially performs writes to the config
        // server, so the user-specified write concern has no effect.
        assert.commandWorked(res);
        assert(!res.writeConcernError, "command on a full cluster had writeConcernError: " + tojson(res));
        cmd.confirmFunc();
    } else {
        assert.commandWorkedIgnoringWriteConcernErrors(res);
        assertWriteConcernError(res, ErrorCodes.UnknownReplWriteConcern);
        const writesExpected = true;
        cmd.confirmFunc(writesExpected);
    }
}

commands.forEach(function (cmd) {
    testValidWriteConcern(cmd);
    testInvalidWriteConcern(cmd);
});

st.stop();