SERVER-47209 Account for rollback in change_streams_update_lookup_shard_metadata_missing.js

(cherry picked from commit 8cb02c375ad128bca9bad7edcc29a8c55fbd9269)

2 files changed, 36 insertions, 35 deletions

diff --git a/jstests/multiVersion/libs/multi_cluster.js b/jstests/multiVersion/libs/multi_cluster.js
index 0cea472ea50..7b776c6fe36 100644
--- a/jstests/multiVersion/libs/multi_cluster.js
+++ b/jstests/multiVersion/libs/multi_cluster.js
@@ -2,9 +2,6 @@
 // MultiVersion utility functions for clusters
 //
 
-load('jstests/multiVersion/libs/multi_rs.js');  // For upgradeSet.
-load('jstests/replsets/rslib.js');              // For awaitRSClientHosts.
-
 /**
  * Restarts the specified binaries in options with the specified binVersion.
  * Note: this does not perform any upgrade operations.
@@ -16,8 +13,14 @@ load('jstests/replsets/rslib.js');              // For awaitRSClientHosts.
  *     upgradeShards: <bool>, // defaults to true
  *     upgradeConfigs: <bool>, // defaults to true
  *     upgradeMongos: <bool>, // defaults to true
+ *     waitUntilStable: <bool>, // defaults to false since it provides a more realistic
+ *                                 approximation of real-world upgrade behaviour, even though
+ *                                 certain tests will likely want a stable cluster after upgrading.
  * }
  */
+load("jstests/multiVersion/libs/multi_rs.js");  // Used by upgradeSet.
+load("jstests/replsets/rslib.js");              // For awaitRSClientHosts.
+
 ShardingTest.prototype.upgradeCluster = function(binVersion, options) {
     options = options || {};
     if (options.upgradeShards == undefined)
@@ -93,24 +96,25 @@ ShardingTest.prototype.upgradeCluster = function(binVersion, options) {
     }
 
     if (options.waitUntilStable) {
-        // Wait for the config server and shards to become available.
-        this.configRS.awaitSecondaryNodes();
-        let shardPrimaries = [];
-        for (let rs of this._rs) {
-            rs.test.awaitSecondaryNodes();
-            shardPrimaries.push(rs.test.getPrimary());
-        }
+        this.waitUntilStable();
+    }
+};
 
-        // Wait for the ReplicaSetMonitor on mongoS and each shard to reflect the state of all
-        // shards.
-        for (let client of[...this._mongos, ...shardPrimaries]) {
-            awaitRSClientHosts(client, shardPrimaries, {ok: true, ismaster: true});
-        }
+ShardingTest.prototype.waitUntilStable = function() {
+    // Wait for the config server and shards to become available.
+    this.configRS.awaitSecondaryNodes();
+    let shardPrimaries = [];
+    for (let rs of this._rs) {
+        rs.test.awaitSecondaryNodes();
+        shardPrimaries.push(rs.test.getPrimary());
+    }
+    // Wait for the ReplicaSetMonitor on mongoS and each shard to reflect the state of all shards.
+    for (let client of[...this._mongos, ...shardPrimaries]) {
+        awaitRSClientHosts(client, shardPrimaries, {ok: true, ismaster: true});
     }
 };
 
 ShardingTest.prototype.restartMongoses = function() {
-
     var numMongoses = this._mongos.length;
 
     for (var i = 0; i < numMongoses; i++) {
diff --git a/jstests/sharding/change_streams_update_lookup_shard_metadata_missing.js b/jstests/sharding/change_streams_update_lookup_shard_metadata_missing.js
index 7b37c80d0bf..0602df6ed63 100644
--- a/jstests/sharding/change_streams_update_lookup_shard_metadata_missing.js
+++ b/jstests/sharding/change_streams_update_lookup_shard_metadata_missing.js
@@ -7,6 +7,8 @@
 (function() {
     "use strict";
 
+    load("jstests/multiVersion/libs/multi_cluster.js");  // For ShardingTest.waitUntilStable.
+
     // The UUID consistency check can hit NotMasterNoSlaveOk when it attempts to obtain a list of
     // collections from the shard Primaries through mongoS at the end of this test.
     TestData.skipCheckingUUIDsConsistentAcrossCluster = true;
@@ -18,9 +20,9 @@
         rs: {nodes: 3, setParameter: {writePeriodicNoops: true, periodicNoopIntervalSecs: 1}}
     });
 
-    const mongosDB = st.s.getDB(jsTestName());
-    const mongosColl = mongosDB.test;
-    const shard0 = st.rs0;
+    let mongosDB = st.s.getDB(jsTestName());
+    let mongosColl = mongosDB.test;
+    let shard0 = st.rs0;
 
     // Enable sharding on the the test database and ensure that the primary is shard0.
     assert.commandWorked(mongosDB.adminCommand({enableSharding: mongosDB.getName()}));
@@ -38,26 +40,21 @@
 
     const resumeToken = csCursor.next()._id;
 
-    // get any secondary
-    const newPrimary = st.rs0.getSecondary();
-    let shards = st.s.getDB('config').shards.find().toArray();
+    // Obtain a reference to any secondary.
+    const newPrimary = shard0.getSecondary();
 
     // Step up one of the Secondaries, which will not have any sharding metadata loaded.
-    st.rs0.stepUpNoAwaitReplication(newPrimary);
-
-    // make sure the mongos refreshes it's connections to the shard
-    let primary = {};
-    do {
-        let connPoolStats = st.s0.adminCommand({connPoolStats: 1});
-        primary = connPoolStats.replicaSets[shards[0]._id].hosts.find((host) => {
-            return host.ismaster;
-        }) ||
-            {};
-    } while (newPrimary.host !== primary.addr);
+    shard0.stepUp(newPrimary);
+
+    // Make sure the mongoS and both shards sees shard0's new primary.
+    st.waitUntilStable();
+
+    // Refresh our reference to the test collection.
+    mongosColl = st.s.getDB(mongosDB.getName())[mongosColl.getName()];
 
     // Do a {multi:true} update. This will scatter to all shards and update the document on shard0.
-    // Because no metadata is loaded, the shard will return a StaleShardVersion and fetch it,
-    // the operation will be retried
+    // Because no metadata is loaded, the shard will return a StaleShardVersion and fetch it, and
+    // the operation will be retried until it completes successfully.
     assert.soonNoExcept(() => assert.commandWorked(
                             mongosColl.update({_id: 0}, {$set: {updated: true}}, false, true)));