SERVER-64628 Add more targeted tests for adding nodes to sharded sets

(cherry picked from commit 181a405baf42fc204e924dcde8f246d6f40a09bb)

4 files changed, 195 insertions, 0 deletions

diff --git a/jstests/sharding/libs/sharding_state_test.js b/jstests/sharding/libs/sharding_state_test.js
new file mode 100644
index 00000000000..86b35432395
--- /dev/null
+++ b/jstests/sharding/libs/sharding_state_test.js
@@ -0,0 +1,114 @@
+"use strict";
+
+/**
+ * This file contains helpers for testing sharding state with various operations in the system.
+ */
+
+var ShardingStateTest = (function() {
+    load("jstests/replsets/rslib.js");
+
+    /**
+     * Adds a node to the given shard or config server replica set. Also waits for the node
+     * to become a steady-state secondary.
+     *
+     * @param {Object} replSet the ReplSetTest instance for the corresponding replica set
+     * @param {bool} elect whether to automatically elect the new node
+     * @param {String} serverTypeFlag startup flag for sharding members ["shardsvr"|"configsvr"]
+     * @param {Object} newNodeParams setParameter-s to start the node with
+     *
+     * @returns a connection to the new node added
+     */
+    function addReplSetNode({replSet, serverTypeFlag, newNodeParams = {}}) {
+        jsTestLog("[ShardingStateTest] Adding new member to replica set.");
+
+        const addParams = {setParameter: newNodeParams};
+
+        if (serverTypeFlag !== undefined) {
+            addParams[serverTypeFlag] = "";
+        }
+
+        const newNode = replSet.add(addParams);
+
+        replSet.reInitiate();
+        replSet.waitForState(newNode, ReplSetTest.State.SECONDARY);
+        replSet.waitForAllNewlyAddedRemovals();
+
+        return newNode;
+    }
+
+    /**
+     * Forces a node to go into startup recovery.
+     *
+     * @param {Object} replSet the ReplSetTest instance for the corresponding replica set
+     * @param {Object} node the ReplSetTest instance for the corresponding replica set
+     * @param {bool} elect whether to automatically elect the new node
+     * @param {Object} startupParams setParameter-s to restart the node with
+     *
+     * @returns a connection to the restarted node
+     */
+    function putNodeInStartupRecovery({replSet, node, startupParams = {}}) {
+        const dbName = "testDB";
+        const collName = "testColl";
+
+        const primary = replSet.getPrimary();
+
+        const testDB = primary.getDB(dbName);
+        const testColl = testDB.getCollection(collName);
+
+        jsTestLog("[ShardingStateTest] Adding a write to pin the stable timestamp at.");
+
+        const ts =
+            assert
+                .commandWorked(testDB.runCommand({insert: testColl.getName(), documents: [{a: 1}]}))
+                .operationTime;
+
+        configureFailPoint(node, 'holdStableTimestampAtSpecificTimestamp', {timestamp: ts});
+
+        jsTestLog("[ShardingStateTest] Adding more data before restarting member.");
+        assert.commandWorked(testColl.insert([{b: 2}, {c: 3}]));
+        replSet.awaitReplication();
+
+        jsTestLog("[ShardingStateTest] Restarting node. It should go into startup recovery.");
+        replSet.restart(node, {setParameter: startupParams});
+        replSet.waitForState(node, ReplSetTest.State.PRIMARY);
+
+        return node;
+    }
+
+    /**
+     * Internal function to failover to a new primary. Waits for the election to complete
+     * and for the rest of the set to acknowledge the new primary.
+     *
+     * @param {Object} replSet replica set to operate on
+     * @param {Object} node an electable target
+     */
+    function failoverToMember(replSet, node) {
+        jsTestLog("[ShardingStateTest] Electing new primary: " + node.host);
+        assert.soonNoExcept(function() {
+            assert.commandWorked(node.adminCommand({replSetStepUp: 1}));
+            return true;
+        });
+
+        replSet.awaitNodesAgreeOnPrimary(undefined /* timesout */, undefined /* nodes */, node);
+    }
+
+    /**
+     * Performs sharding state checks.
+     *
+     * @param {Object} st the sharded cluster to perform checks on
+     */
+    function checkShardingState(st) {
+        jsTestLog("[ShardingStateTest] Performing sharding state checks.");
+
+        assert.commandWorked(st.s.getDB("sstDB").getCollection('sstColl').insert({"sstDoc": 1}));
+
+        jsTestLog("[ShardingStateTest] Sharding state checks succeeded.");
+    }
+
+    return {
+        addReplSetNode: addReplSetNode,
+        putNodeInStartupRecovery: putNodeInStartupRecovery,
+        failoverToMember: failoverToMember,
+        checkShardingState: checkShardingState,
+    };
+})();
diff --git a/jstests/sharding/live_shard_logical_initial_sync.js b/jstests/sharding/live_shard_logical_initial_sync.js
new file mode 100644
index 00000000000..ffecb60d802
--- /dev/null
+++ b/jstests/sharding/live_shard_logical_initial_sync.js
@@ -0,0 +1,27 @@
+/**
+ * Tests that sharding state is properly initialized on new shard members that were added into live
+ * shards using logical initial sync.
+ *
+ * We control our own failovers, and we also need the RSM to react reasonably quickly to those.
+ * @tags: [does_not_support_stepdowns, requires_streamable_rsm]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/sharding/libs/sharding_state_test.js");
+
+const st = new ShardingTest({config: 1, shards: {rs0: {nodes: 1}}});
+const rs = st.rs0;
+
+const newNode = ShardingStateTest.addReplSetNode({replSet: rs, serverTypeFlag: "shardsvr"});
+
+jsTestLog("Checking sharding state before failover.");
+ShardingStateTest.checkShardingState(st);
+
+jsTestLog("Checking sharding state after failover.");
+ShardingStateTest.failoverToMember(rs, newNode);
+ShardingStateTest.checkShardingState(st);
+
+st.stop();
+})();
diff --git a/jstests/sharding/live_shard_logical_initial_sync_config_server.js b/jstests/sharding/live_shard_logical_initial_sync_config_server.js
new file mode 100644
index 00000000000..b9a4803934a
--- /dev/null
+++ b/jstests/sharding/live_shard_logical_initial_sync_config_server.js
@@ -0,0 +1,30 @@
+/**
+ * Tests that sharding state is properly initialized on new config members that were added into live
+ * config server replica sets using logical initial sync.
+ *
+ * We control our own failovers, and we also need the RSM to react reasonably quickly to those.
+ * @tags: [does_not_support_stepdowns, requires_streamable_rsm]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/sharding/libs/sharding_state_test.js");
+
+const st = new ShardingTest({
+    config: 1,
+    shards: {rs0: {nodes: 1}},
+});
+const configRS = st.configRS;
+
+const newNode = ShardingStateTest.addReplSetNode({replSet: configRS, serverTypeFlag: "configsvr"});
+
+jsTestLog("Checking sharding state before failover.");
+ShardingStateTest.checkShardingState(st);
+
+jsTestLog("Checking sharding state after failover.");
+ShardingStateTest.failoverToMember(configRS, newNode);
+ShardingStateTest.checkShardingState(st);
+
+st.stop();
+})();
diff --git a/jstests/sharding/live_shard_startup_recovery.js b/jstests/sharding/live_shard_startup_recovery.js
new file mode 100644
index 00000000000..f40cb38a64a
--- /dev/null
+++ b/jstests/sharding/live_shard_startup_recovery.js
@@ -0,0 +1,24 @@
+/**
+ * Tests that sharding state is properly initialized on shard members that undergo startup recovery.
+ *
+ * We control our own failovers, and we also need the RSM to react reasonably quickly to those.
+ * @tags: [requires_persistence, does_not_support_stepdowns, requires_streamable_rsm]
+ */
+
+(function() {
+"use strict";
+
+load("jstests/sharding/libs/sharding_state_test.js");
+
+const st = new ShardingTest({config: 1, shards: {rs0: {nodes: 1}}});
+const rs = st.rs0;
+let primary = rs.getPrimary();
+
+primary = ShardingStateTest.putNodeInStartupRecovery({replSet: rs, node: primary});
+
+jsTestLog("Ensuring node is up as a primary and checking sharding state");
+ShardingStateTest.failoverToMember(rs, primary);
+ShardingStateTest.checkShardingState(st);
+
+st.stop();
+})();