SERVER-48170 Multiversion tests assume primary stability when using upgradeCluster() with 2-node replica set shards

1 files changed, 29 insertions, 18 deletions

diff --git a/jstests/multiVersion/libs/multi_rs.js b/jstests/multiVersion/libs/multi_rs.js
index 6c19c08d371..f1c49769825 100644
--- a/jstests/multiVersion/libs/multi_rs.js
+++ b/jstests/multiVersion/libs/multi_rs.js
@@ -2,6 +2,13 @@
 // Utility functions for multi-version replica sets
 //
 
+ReplSetTest.prototype._stablePrimaryOnRestarts = function() {
+    // In a 2-node replica set the secondary can step up after a restart. In fact while the
+    // secondary is being restarted, the primary may end up stepping down (due to heartbeats not
+    // being received) and for the restarted node to run for and win the election.
+    return this.nodes.length > 2;
+};
+
 /**
  * Upgrade or downgrade replica sets.
  *
@@ -10,12 +17,22 @@
  * @param pwd {string} optional, password for authentication. Must be set if user is set.
  */
 ReplSetTest.prototype.upgradeSet = function(options, user, pwd) {
+    this.awaitNodesAgreeOnPrimary();
     let primary = this.getPrimary();
 
-    this.upgradeSecondaries(primary, Object.assign({}, options), user, pwd);
-    this.upgradeArbiters(primary, Object.assign({}, options), user, pwd);
+    this.upgradeSecondaries(Object.assign({}, options), user, pwd);
+    this.upgradeArbiters(Object.assign({}, options), user, pwd);
+
+    if (!this._stablePrimaryOnRestarts()) {
+        this.awaitNodesAgreeOnPrimary();
+        if (this.getPrimary() != primary) {
+            this.upgradeMembers([primary], Object.assign({}, options), user, pwd);
+            return;
+        }
+    }
 
-    // Upgrade the primary after stepping down.
+    assert.eq(
+        this.getPrimary(), primary, "Primary changed unexpectedly after upgrading secondaries");
     this.upgradePrimary(primary, Object.assign({}, options), user, pwd);
 };
 
@@ -26,17 +43,12 @@ function mergeNodeOptions(nodeOptions, options) {
     return nodeOptions;
 }
 
-ReplSetTest.prototype.upgradeMembers = function(primary, members, options, user, pwd) {
-    const noDowntimePossible = this.nodes.length > 2;
-
+ReplSetTest.prototype.upgradeMembers = function(members, options, user, pwd) {
     // Merge new options into node settings.
     this.nodeOptions = mergeNodeOptions(this.nodeOptions, options);
 
     for (let member of members) {
         this.upgradeNode(member, options, user, pwd);
-
-        if (noDowntimePossible)
-            assert.eq(this.getPrimary(), primary);
     }
 };
 
@@ -47,8 +59,8 @@ ReplSetTest.prototype.getNonArbiterSecondaries = function() {
     return nonArbiters;
 };
 
-ReplSetTest.prototype.upgradeSecondaries = function(primary, options, user, pwd) {
-    this.upgradeMembers(primary, this.getNonArbiterSecondaries(), options, user, pwd);
+ReplSetTest.prototype.upgradeSecondaries = function(options, user, pwd) {
+    this.upgradeMembers(this.getNonArbiterSecondaries(), options, user, pwd);
 };
 
 ReplSetTest.prototype.upgradeArbiters = function(primary, options, user, pwd) {
@@ -57,14 +69,12 @@ ReplSetTest.prototype.upgradeArbiters = function(primary, options, user, pwd) {
     if (options && options.binVersion == "last-lts") {
         options["startClean"] = true;
     }
-    this.upgradeMembers(primary, this.getArbiters(), options, user, pwd);
+    this.upgradeMembers(this.getArbiters(), options, user, pwd);
     // Make sure we don't set {startClean:true} on other nodes unless the user explicitly requested.
     this.nodeOptions = mergeNodeOptions(this.nodeOptions, oldStartClean);
 };
 
 ReplSetTest.prototype.upgradePrimary = function(primary, options, user, pwd) {
-    const noDowntimePossible = this.nodes.length > 2;
-
     // Merge new options into node settings.
     this.nodeOptions = mergeNodeOptions(this.nodeOptions, options);
 
@@ -85,9 +95,10 @@ ReplSetTest.prototype.upgradePrimary = function(primary, options, user, pwd) {
 
     let newPrimary = this.getPrimary();
 
-    if (noDowntimePossible)
-        assert.eq(newPrimary, primary);
-
+    if (this._stablePrimaryOnRestarts()) {
+        assert.eq(
+            newPrimary, primary, "Primary changed unexpectedly after upgrading old primary node");
+    }
     return newPrimary;
 };
 
@@ -124,7 +135,7 @@ ReplSetTest.prototype.upgradeNode = function(node, opts = {}, user, pwd) {
 
 ReplSetTest.prototype.stepdown = function(nodeId) {
     nodeId = this.getNodeId(nodeId);
-    assert.eq(this.getNodeId(this.getPrimary()), nodeId);
+    assert.eq(this.getNodeId(this.getPrimary()), nodeId, "Trying to stepdown a non primary node");
     var node = this.nodes[nodeId];
 
     assert.soonNoExcept(function() {