SERVER-69802 Support sampling write queries on sharded clusters

3 files changed, 475 insertions, 18 deletions

diff --git a/jstests/sharding/analyze_shard_key/libs/query_sampling_util.js b/jstests/sharding/analyze_shard_key/libs/query_sampling_util.js
index 9f3e07d8e50..7a627bf6a0e 100644
--- a/jstests/sharding/analyze_shard_key/libs/query_sampling_util.js
+++ b/jstests/sharding/analyze_shard_key/libs/query_sampling_util.js
@@ -27,13 +27,44 @@ var QuerySamplingUtil = (function() {
     }
 
     /**
-     * Waits for the config.sampledQueries collection to have 'expectedSampledQueryDocs.length'
-     * number of documents for the collection 'ns'. For every (sampleId, cmdName, cmdObj) in
+     * Waits for the given mongos to have one active collection for query sampling.
+     */
+    function waitForActiveSampling(mongosConn) {
+        assert.soon(() => {
+            const res = assert.commandWorked(mongosConn.adminCommand({serverStatus: 1}));
+            return res.queryAnalyzers.activeCollections == 1;
+        });
+    }
+
+    /**
+     * Returns true if 'subsetObj' is a sub object of 'supersetObj'. That is, every key that exists
+     * in 'subsetObj' also exists in 'supersetObj' and the values of that key in the two objects are
+     * equal.
+     */
+    function assertSubObject(supersetObj, subsetObj) {
+        for (let key in subsetObj) {
+            const value = subsetObj[key];
+            if (typeof value === 'object') {
+                assertSubObject(supersetObj[key], subsetObj[key]);
+            } else {
+                assert.eq(supersetObj[key],
+                          subsetObj[key],
+                          {key, actual: supersetObj, expected: subsetObj});
+            }
+        }
+    }
+
+    const kSampledQueriesNs = "config.sampledQueries";
+    const kSampledQueriesDiffNs = "config.sampledQueriesDiff";
+
+    /**
+     * Waits for the number of the config.sampledQueries documents for the collection 'ns' to be
+     * equal to 'expectedSampledQueryDocs.length'. Then, for every (sampleId, cmdName, cmdObj) in
      * 'expectedSampledQueryDocs', asserts that there is a config.sampledQueries document with _id
-     * equal to sampleId and that it has the given fields.
+     * equal to 'sampleId' and that the document has the expected fields.
      */
     function assertSoonSampledQueryDocuments(conn, ns, collectionUuid, expectedSampledQueryDocs) {
-        const coll = conn.getCollection("config.sampledQueries");
+        const coll = conn.getCollection(kSampledQueriesNs);
 
         let actualSampledQueryDocs;
         assert.soon(() => {
@@ -51,19 +82,71 @@ var QuerySamplingUtil = (function() {
             assert.eq(doc.ns, ns, doc);
             assert.eq(doc.collectionUuid, collectionUuid, doc);
             assert.eq(doc.cmdName, cmdName, doc);
+            assertSubObject(doc.cmd, cmdObj);
+        }
+    }
+
+    /**
+     * Waits for the total number of the config.sampledQueries documents for the collection 'ns' and
+     * commands 'cmdNames' across all shards to be equal to 'expectedSampledQueryDocs.length'. Then,
+     * for every (filter, shardNames, cmdName, cmdObj, diff) in 'expectedSampledQueryDocs', asserts
+     * that:
+     * - There is exactly one shard that has the config.sampledQueries document that 'filter'
+     *   matches against, and that shard is one of the shards in 'shardNames'.
+     * - The document has the expected fields. If 'diff' is not null, the query has a corresponding
+     *   config.sampledQueriesDiff document with the expected diff on that same shard.
+     */
+    function assertSoonSampledQueryDocumentsAcrossShards(
+        st, ns, collectionUuid, cmdNames, expectedSampledQueryDocs) {
+        let actualSampledQueryDocs, actualCount;
+        assert.soon(() => {
+            actualSampledQueryDocs = {};
+            actualCount = 0;
+
+            st._rs.forEach((rs) => {
+                const docs = rs.test.getPrimary()
+                                 .getCollection(kSampledQueriesNs)
+                                 .find({cmdName: {$in: cmdNames}})
+                                 .toArray();
+                actualSampledQueryDocs[[rs.test.name]] = docs;
+                actualCount += docs.length;
+            });
+            return actualCount >= expectedSampledQueryDocs.length;
+        }, "timed out waiting for sampled query documents");
+        assert.eq(actualCount,
+                  expectedSampledQueryDocs.length,
+                  {actualSampledQueryDocs, expectedSampledQueryDocs});
+
+        for (let {filter, shardNames, cmdName, cmdObj, diff} of expectedSampledQueryDocs) {
+            let shardName = null;
+            for (let rs of st._rs) {
+                const primary = rs.test.getPrimary();
+                const queryDoc = primary.getCollection(kSampledQueriesNs).findOne(filter);
+
+                if (shardName) {
+                    assert.eq(queryDoc,
+                              null,
+                              "Found a sampled query on more than one shard " +
+                                  tojson({shardNames: [shardName, rs.test.name], cmdName, cmdObj}));
+                    continue;
+                } else if (queryDoc) {
+                    shardName = rs.test.name;
+                    assert(shardNames.includes(shardName),
+                           "Found a sampled query on an unexpected shard " +
+                               tojson({actual: shardName, expected: shardNames, cmdName, cmdObj}));
+
+                    assert.eq(queryDoc.ns, ns, queryDoc);
+                    assert.eq(queryDoc.collectionUuid, collectionUuid, queryDoc);
+                    assert.eq(queryDoc.cmdName, cmdName, queryDoc);
+                    assertSubObject(queryDoc.cmd, cmdObj);
 
-            for (let key in cmdObj) {
-                const value = cmdObj[key];
-                if (typeof value === 'object') {
-                    for (let subKey in value) {
-                        assert.eq(doc.cmd[key][subKey],
-                                  cmdObj[key][subKey],
-                                  {subKey, actual: doc.cmd, expected: cmdObj});
+                    if (diff) {
+                        assertSoonSingleSampledDiffDocument(
+                            primary, queryDoc._id, ns, collectionUuid, [diff]);
                     }
-                } else {
-                    assert.eq(doc.cmd[key], cmdObj[key], {key, actual: doc.cmd, expected: cmdObj});
                 }
             }
+            assert(shardName, "Failed to find the sampled query " + tojson({cmdName, cmdObj}));
         }
     }
 
@@ -74,12 +157,12 @@ var QuerySamplingUtil = (function() {
 
     /**
      * Waits for the config.sampledQueriesDiff collection to have a document with _id equal to
-     * sampleId, and then asserts that the diff in that document matches one of the diffs in
-     * 'expectedSampledDiffs'.
+     * 'sampleId' for the collection 'ns', and then asserts that the diff in that document matches
+     * one of the diffs in 'expectedSampledDiffs'.
      */
     function assertSoonSingleSampledDiffDocument(
         conn, sampleId, ns, collectionUuid, expectedSampledDiffs) {
-        const coll = conn.getCollection("config.sampledQueriesDiff");
+        const coll = conn.getCollection(kSampledQueriesDiffNs);
 
         assert.soon(() => {
             const doc = coll.findOne({_id: sampleId});
@@ -97,12 +180,12 @@ var QuerySamplingUtil = (function() {
     }
 
     function assertNoSampledDiffDocuments(conn, ns) {
-        const coll = conn.getCollection("config.sampledQueriesDiff");
+        const coll = conn.getCollection(kSampledQueriesDiffNs);
         assert.eq(coll.find({ns: ns}).itcount(), 0);
     }
 
     function clearSampledDiffCollection(primary) {
-        const coll = primary.getCollection("config.sampledQueriesDiff");
+        const coll = primary.getCollection(kSampledQueriesDiffNs);
         assert.commandWorked(coll.remove({}));
     }
 
@@ -111,7 +194,9 @@ var QuerySamplingUtil = (function() {
         generateRandomString,
         generateRandomCollation,
         makeCmdObjIgnoreSessionInfo,
+        waitForActiveSampling,
         assertSoonSampledQueryDocuments,
+        assertSoonSampledQueryDocumentsAcrossShards,
         assertNoSampledQueryDocuments,
         assertSoonSingleSampledDiffDocument,
         assertNoSampledDiffDocuments,
diff --git a/jstests/sharding/analyze_shard_key/sample_write_queries_sharded.js b/jstests/sharding/analyze_shard_key/sample_write_queries_sharded.js
new file mode 100644
index 00000000000..920020c4d7b
--- /dev/null
+++ b/jstests/sharding/analyze_shard_key/sample_write_queries_sharded.js
@@ -0,0 +1,207 @@
+/**
+ * Tests basic support for sampling write queries against a sharded collection on a sharded cluster.
+ *
+ * @tags: [requires_fcv_62, featureFlagAnalyzeShardKey]
+ */
+(function() {
+"use strict";
+
+load("jstests/sharding/analyze_shard_key/libs/query_sampling_util.js");
+
+// Make the periodic jobs for refreshing sample rates and writing sampled queries and diffs have a
+// period of 1 second to speed up the test.
+const st = new ShardingTest({
+    shards: 3,
+    rs: {nodes: 2, setParameter: {queryAnalysisWriterIntervalSecs: 1}},
+    mongosOptions: {setParameter: {queryAnalysisSamplerConfigurationRefreshSecs: 1}}
+});
+
+const dbName = "testDb";
+const collName = "testColl";
+const ns = dbName + "." + collName;
+const mongosDB = st.s.getDB(dbName);
+const mongosColl = mongosDB.getCollection(collName);
+
+// Make the collection have two chunks:
+// shard0: [MinKey, 0]
+// shard1: [0, 1000]
+// shard1: [1000, MaxKey]
+assert.commandWorked(st.s.adminCommand({enableSharding: dbName}));
+st.ensurePrimaryShard(dbName, st.shard0.name);
+// TODO (SERVER-69237): Use a regular collection once pre-images are always available in the
+// OpObserver. Currently, the pre-images are not available in the test cases involving array
+// updates.
+assert.commandWorked(
+    mongosDB.createCollection(collName, {changeStreamPreAndPostImages: {enabled: true}}));
+assert.commandWorked(mongosColl.createIndex({x: 1}));
+assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {x: 1}}));
+assert.commandWorked(st.s.adminCommand({split: ns, middle: {x: 0}}));
+assert.commandWorked(st.s.adminCommand({split: ns, middle: {x: 1000}}));
+assert.commandWorked(st.s.adminCommand({moveChunk: ns, find: {x: 0}, to: st.shard1.shardName}));
+assert.commandWorked(st.s.adminCommand({moveChunk: ns, find: {x: 1000}, to: st.shard2.shardName}));
+const collectionUuid = QuerySamplingUtil.getCollectionUuid(mongosDB, collName);
+
+assert.commandWorked(
+    st.s.adminCommand({configureQueryAnalyzer: ns, mode: "full", sampleRate: 1000}));
+QuerySamplingUtil.waitForActiveSampling(st.s);
+
+const expectedSampledQueryDocs = [];
+
+// Make each write below have a unique filter and use that to look up the corresponding
+// config.sampledQueries document later.
+
+{
+    // Perform some updates.
+    assert.commandWorked(mongosColl.insert([
+        // The docs below are on shard1.
+        {x: 1, y: 1, z: [1, 0, 1]},
+        {x: 2, y: 2, z: [2]},
+        // The doc below is on shard2.
+        {x: 1002, y: 2, z: [2]}
+    ]));
+
+    const cmdName = "update";
+
+    const updateOp0 = {
+        q: {x: 1},
+        u: {$mul: {y: 10}, $set: {"z.$[element]": 10}},
+        arrayFilters: [{"element": {$gte: 1}}],
+        multi: false,
+        upsert: false,
+        collation: QuerySamplingUtil.generateRandomCollation(),
+    };
+    const diff0 = {y: 'u', z: 'u'};
+    const shardNames0 = [st.rs1.name];
+
+    const updateOp1 = {
+        q: {x: {$gte: 2}},
+        u: [{$set: {y: 20, w: 200}}],
+        c: {var0: 1},
+        multi: true,
+    };
+    const diff1 = {y: 'u', w: 'i'};
+    const shardNames1 = [st.rs1.name, st.rs2.name];
+
+    const originalCmdObj = {
+        update: collName,
+        updates: [updateOp0, updateOp1],
+        let : {var1: 1},
+    };
+
+    // Use a transaction, otherwise updateOp1 would get routed to all shards.
+    const lsid = {id: UUID()};
+    const txnNumber = NumberLong(1);
+    assert.commandWorked(mongosDB.runCommand(Object.assign(
+        {}, originalCmdObj, {lsid, txnNumber, startTransaction: true, autocommit: false})));
+    assert.commandWorked(
+        mongosDB.adminCommand({commitTransaction: 1, lsid, txnNumber, autocommit: false}));
+    assert.neq(mongosColl.findOne({x: 1, y: 10, z: [10, 0, 10]}), null);
+    assert.neq(mongosColl.findOne({x: 2, y: 20, z: [2], w: 200}), null);
+    assert.neq(mongosColl.findOne({x: 1002, y: 20, z: [2], w: 200}), null);
+
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.updates.0.q": updateOp0.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {updates: [updateOp0]}),
+        diff: diff0,
+        shardNames: shardNames0
+    });
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.updates.0.q": updateOp1.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {updates: [updateOp1]}),
+        diff: diff1,
+        shardNames: shardNames1
+    });
+}
+
+{
+    // Perform some deletes.
+    assert.commandWorked(mongosColl.insert([
+        // The docs below are on shard1.
+        {x: 3},
+        {x: 4},
+        // The docs below are on shard2.
+        {x: 1004}
+    ]));
+
+    const cmdName = "delete";
+
+    const deleteOp0 = {
+        q: {x: 3},
+        limit: 1,
+        collation: QuerySamplingUtil.generateRandomCollation(),
+    };
+    const shardNames0 = [st.rs1.name];
+
+    const deleteOp1 = {q: {x: {$gte: 4}}, limit: 0};
+    const shardNames1 = [st.rs1.name, st.rs2.name];
+
+    const originalCmdObj = {delete: collName, deletes: [deleteOp0, deleteOp1]};
+
+    // Use a transaction, otherwise deleteOp1 would get routed to all shards.
+    const lsid = {id: UUID()};
+    const txnNumber = NumberLong(1);
+    assert.commandWorked(mongosDB.runCommand(Object.assign(
+        {}, originalCmdObj, {lsid, txnNumber, startTransaction: true, autocommit: false})));
+    assert.commandWorked(
+        mongosDB.adminCommand({commitTransaction: 1, lsid, txnNumber, autocommit: false}));
+    assert.eq(mongosColl.findOne({x: 3}), null);
+    assert.eq(mongosColl.findOne({x: 4}), null);
+
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.deletes.0.q": deleteOp0.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {deletes: [deleteOp0]}),
+        shardNames: shardNames0
+    });
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.deletes.0.q": deleteOp1.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {deletes: [deleteOp1]}),
+        shardNames: shardNames1
+    });
+}
+
+{
+    // Perform some findAndModify.
+    assert.commandWorked(mongosColl.insert([
+        // The doc below is on shard0.
+        {x: -5, y: -5, z: [-5, 0, -5]}
+    ]));
+
+    const cmdName = "findAndModify";
+    const originalCmdObj = {
+        findAndModify: collName,
+        query: {x: -5},
+        update: {$mul: {y: 10}, $set: {"z.$[element]": -50}},
+        arrayFilters: [{"element": {$lte: -5}}],
+        sort: {_id: 1},
+        collation: QuerySamplingUtil.generateRandomCollation(),
+        new: true,
+        upsert: false,
+        let : {var0: 1}
+    };
+    const diff = {y: 'u', z: 'u'};
+    const shardNames = [st.rs0.name];
+
+    assert.commandWorked(mongosDB.runCommand(originalCmdObj));
+    assert.neq(mongosColl.findOne({x: -5, y: -50, z: [-50, 0, -50]}), null);
+
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.query": originalCmdObj.query},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj),
+        diff,
+        shardNames
+    });
+}
+
+const cmdNames = ["update", "delete", "findAndModify"];
+QuerySamplingUtil.assertSoonSampledQueryDocumentsAcrossShards(
+    st, ns, collectionUuid, cmdNames, expectedSampledQueryDocs);
+
+assert.commandWorked(st.s.adminCommand({configureQueryAnalyzer: ns, mode: "off"}));
+
+st.stop();
+})();
diff --git a/jstests/sharding/analyze_shard_key/sample_write_queries_unsharded.js b/jstests/sharding/analyze_shard_key/sample_write_queries_unsharded.js
new file mode 100644
index 00000000000..c935d819374
--- /dev/null
+++ b/jstests/sharding/analyze_shard_key/sample_write_queries_unsharded.js
@@ -0,0 +1,165 @@
+/**
+ * Tests basic support for sampling write queries against an unsharded collection on a sharded
+ * cluster.
+ *
+ * @tags: [requires_fcv_62, featureFlagAnalyzeShardKey]
+ */
+(function() {
+"use strict";
+
+load("jstests/sharding/analyze_shard_key/libs/query_sampling_util.js");
+
+// Make the periodic jobs for refreshing sample rates and writing sampled queries and diffs have a
+// period of 1 second to speed up the test.
+const st = new ShardingTest({
+    shards: 2,
+    rs: {nodes: 2, setParameter: {queryAnalysisWriterIntervalSecs: 1}},
+    mongosOptions: {setParameter: {queryAnalysisSamplerConfigurationRefreshSecs: 1}}
+});
+
+const dbName = "testDb";
+const collName = "testColl";
+const ns = dbName + "." + collName;
+const mongosDB = st.s.getDB(dbName);
+const mongosColl = mongosDB.getCollection(collName);
+
+assert.commandWorked(st.s.adminCommand({enableSharding: dbName}));
+st.ensurePrimaryShard(dbName, st.shard0.name);
+// TODO (SERVER-69237): Use a regular collection once pre-images are always available in the
+// OpObserver. Currently, the pre-images are not available in the test cases involving array
+// updates.
+assert.commandWorked(
+    mongosDB.createCollection(collName, {changeStreamPreAndPostImages: {enabled: true}}));
+const collectionUuid = QuerySamplingUtil.getCollectionUuid(mongosDB, collName);
+
+assert.commandWorked(
+    st.s.adminCommand({configureQueryAnalyzer: ns, mode: "full", sampleRate: 1000}));
+QuerySamplingUtil.waitForActiveSampling(st.s);
+
+const expectedSampledQueryDocs = [];
+// This is an unsharded collection so all documents are on the primary shard.
+const shardNames = [st.rs0.name];
+
+// Make each write below have a unique filter and use that to look up the corresponding
+// config.sampledQueries document later.
+
+{
+    // Perform some updates.
+    assert.commandWorked(mongosColl.insert([{x: 1, y: 1, z: [1, 0, 1]}, {x: 2, y: 2, z: [2]}]));
+
+    const cmdName = "update";
+    const updateOp0 = {
+        q: {x: 1},
+        u: {$mul: {y: 10}, $set: {"z.$[element]": 10}},
+        arrayFilters: [{"element": {$gte: 1}}],
+        multi: false,
+        upsert: false,
+        collation: QuerySamplingUtil.generateRandomCollation(),
+    };
+    const diff0 = {y: 'u', z: 'u'};
+    const updateOp1 = {
+        q: {x: 2},
+        u: [{$set: {y: 20, w: 200}}],
+        c: {var0: 1},
+        multi: true,
+    };
+    const diff1 = {y: 'u', w: 'i'};
+    const originalCmdObj = {
+        update: collName,
+        updates: [updateOp0, updateOp1],
+        let : {var1: 1},
+    };
+
+    assert.commandWorked(mongosDB.runCommand(originalCmdObj));
+    assert.neq(mongosColl.findOne({x: 1, y: 10, z: [10, 0, 10]}), null);
+    assert.neq(mongosColl.findOne({x: 2, y: 20, z: [2], w: 200}), null);
+
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.updates.0.q": updateOp0.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {updates: [updateOp0]}),
+        diff: diff0,
+        shardNames
+    });
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.updates.0.q": updateOp1.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {updates: [updateOp1]}),
+        diff: diff1,
+        shardNames
+    });
+}
+
+{
+    // Perform some deletes.
+    assert.commandWorked(mongosColl.insert([{x: 3}, {x: 4}]));
+
+    const cmdName = "delete";
+    const deleteOp0 = {
+        q: {x: 3},
+        limit: 1,
+        collation: QuerySamplingUtil.generateRandomCollation(),
+    };
+    const deleteOp1 = {q: {x: 4}, limit: 0};
+    const originalCmdObj = {
+        delete: collName,
+        deletes: [deleteOp0, deleteOp1],
+
+    };
+
+    assert.commandWorked(mongosDB.runCommand(originalCmdObj));
+    assert.eq(mongosColl.findOne({x: 3}), null);
+    assert.eq(mongosColl.findOne({x: 4}), null);
+
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.deletes.0.q": deleteOp0.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {deletes: [deleteOp0]}),
+        shardNames
+    });
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.deletes.0.q": deleteOp1.q},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj, {deletes: [deleteOp1]}),
+        shardNames
+    });
+}
+
+{
+    // Perform some findAndModify.
+    assert.commandWorked(mongosColl.insert([{x: 5, y: 5, z: [5, 0, 5]}]));
+
+    const cmdName = "findAndModify";
+    const originalCmdObj = {
+        findAndModify: collName,
+        query: {x: 5},
+        update: {$mul: {y: 10}, $set: {"z.$[element]": 50}},
+        arrayFilters: [{"element": {$gte: 5}}],
+        sort: {_id: 1},
+        collation: QuerySamplingUtil.generateRandomCollation(),
+        new: true,
+        upsert: false,
+        let : {var0: 1}
+    };
+    const diff = {y: 'u', z: 'u'};
+
+    assert.commandWorked(mongosDB.runCommand(originalCmdObj));
+    assert.neq(mongosColl.findOne({x: 5, y: 50, z: [50, 0, 50]}), null);
+
+    expectedSampledQueryDocs.push({
+        filter: {"cmd.query": originalCmdObj.query},
+        cmdName: cmdName,
+        cmdObj: Object.assign({}, originalCmdObj),
+        diff,
+        shardNames
+    });
+}
+
+const cmdNames = ["update", "delete", "findAndModify"];
+QuerySamplingUtil.assertSoonSampledQueryDocumentsAcrossShards(
+    st, ns, collectionUuid, cmdNames, expectedSampledQueryDocs);
+
+assert.commandWorked(st.s.adminCommand({configureQueryAnalyzer: ns, mode: "off"}));
+
+st.stop();
+})();