SERVER-28100 ReplSetTest.awaitNodesAgreeOnPrimary() accepts an optional expectedPrimaryIndex

(cherry-picked from commit 8d146d21661cf755fc15b7eefec147ce85ffde8c)

4 files changed, 48 insertions, 35 deletions

diff --git a/jstests/replsets/catchup.js b/jstests/replsets/catchup.js
index 542ad51c723..7c1cbf41c85 100644
--- a/jstests/replsets/catchup.js
+++ b/jstests/replsets/catchup.js
@@ -34,11 +34,12 @@
         node.adminCommand(verbosity);
     });
 
-    function stepUp(node) {
-        assert.soon(function() {
-            node.adminCommand({replSetStepUp: 1});
+    function stepUpNode(node) {
+        assert.soonNoExcept(function() {
+            assert.commandWorked(node.adminCommand({replSetStepUp: 1}));
+            rst.awaitNodesAgreeOnPrimary(rst.kDefaultTimeoutMS, rst.nodes, rst.getNodeId(node));
             return node.adminCommand('replSetGetStatus').myState == ReplSetTest.State.PRIMARY;
-        });
+        }, 'failed to step up node ' + node.host, rst.kDefaultTimeoutMS);
 
         return node;
     }
@@ -67,8 +68,7 @@
 
     jsTest.log("Case 1: The primary is up-to-date after freshness scan.");
     // Should complete transition to primary immediately.
-    var newPrimary = stepUp(rst.getSecondary());
-    rst.awaitNodesAgreeOnPrimary();
+    var newPrimary = stepUpNode(rst.getSecondary());
     // Should win an election and finish the transition very quickly.
     assert.eq(newPrimary, rst.getPrimary());
     rst.awaitReplication(ReplSetTest.kDefaultTimeoutMS, ReplSetTest.OpTimeType.LAST_DURABLE);
@@ -80,8 +80,7 @@
     doWrites(rst.getPrimary());
     var latestOp = getLatestOp(rst.getPrimary());
     // New primary wins immediately, but needs to catch up.
-    newPrimary = stepUp(rst.getSecondary());
-    rst.awaitNodesAgreeOnPrimary();
+    newPrimary = stepUpNode(rst.getSecondary());
     // Check this node is not writable.
     assert.eq(newPrimary.getDB("test").isMaster().ismaster, false);
     // Disable fail point to allow replication.
@@ -101,8 +100,7 @@
     var oldPrimary = rst.getPrimary();
     originalSecondaries = rst.getSecondaries();
     latestOp = getLatestOp(oldPrimary);
-    newPrimary = stepUp(originalSecondaries[0]);
-    rst.awaitNodesAgreeOnPrimary();
+    newPrimary = stepUpNode(originalSecondaries[0]);
     // Disable fail point on one of the other secondaries.
     // Wait until it catches up with the old primary.
     restartServerReplication(originalSecondaries[1]);
@@ -134,8 +132,7 @@
     latestOp = getLatestOp(rst.getPrimary());
 
     // New primary wins immediately, but needs to catch up.
-    newPrimary = stepUp(originalSecondaries[0]);
-    rst.awaitNodesAgreeOnPrimary();
+    newPrimary = stepUpNode(originalSecondaries[0]);
     var latestOpOnNewPrimary = getLatestOp(newPrimary);
     // Wait until the new primary completes the transition to primary and writes a no-op.
     checkLog.contains(newPrimary, "Cannot catch up oplog after becoming primary");
diff --git a/jstests/replsets/chaining_removal.js b/jstests/replsets/chaining_removal.js
index bab3373bffa..d0cc9b3ed56 100644
--- a/jstests/replsets/chaining_removal.js
+++ b/jstests/replsets/chaining_removal.js
@@ -21,8 +21,7 @@
             {_id: 4, host: nodes[4].host, priority: 0},
         ],
     });
-    replTest.waitForState(nodes[0], ReplSetTest.State.PRIMARY);
-    replTest.awaitNodesAgreeOnPrimary();
+    replTest.awaitNodesAgreeOnPrimary(replTest.kDefaultTimeoutMS, nodes, 0);
     var primary = replTest.getPrimary();
     replTest.awaitReplication();
 
diff --git a/jstests/replsets/libs/tags.js b/jstests/replsets/libs/tags.js
index e27e26fc18e..f21dfbfdc11 100644
--- a/jstests/replsets/libs/tags.js
+++ b/jstests/replsets/libs/tags.js
@@ -137,21 +137,39 @@ var TagsTest = function(options) {
         assert.commandWorked(replTest.getPrimary().getDB('foo').createCollection('bar'));
         replTest.awaitReplication();
 
-        var ensurePrimary = function(nodeId, expectedWritableNodes) {
-            jsTestLog('Node ' + nodeId + ' (' + replTest.nodes[nodeId].host +
+        // nodeId is the index of the node that we expect to see as primary.
+        // expectedNodesAgreeOnPrimary is a set of nodes that should agree that 'nodeId' is the
+        // primary.
+        // expectedWritableNodesCount is the number of nodes we can expect to write to. Defaults to
+        //     expectedNodesAgreeOnPrimary.length.
+        var ensurePrimary = function(
+            nodeId, expectedNodesAgreeOnPrimary, expectedWritableNodesCount) {
+            expectedWritableNodesCount =
+                expectedWritableNodesCount || expectedNodesAgreeOnPrimary.length;
+            jsTestLog('ensurePrimary - Node ' + nodeId + ' (' + replTest.nodes[nodeId].host +
+                      ') should be primary.');
+            replTest.awaitNodesAgreeOnPrimary(
+                replTest.kDefaultTimeoutMS, expectedNodesAgreeOnPrimary, nodeId);
+            jsTestLog('ensurePrimary - Nodes ' + tojson(expectedNodesAgreeOnPrimary) +
+                      ' agree that ' + nodeId + ' (' + replTest.nodes[nodeId].host +
                       ') should be primary.');
-            replTest.waitForState(replTest.nodes[nodeId], ReplSetTest.State.PRIMARY, 60 * 1000);
             primary = replTest.getPrimary();
             if (options.forceWriteMode) {
                 primary.forceWriteMode(options.forceWriteMode);
             }
-            var writeConcern = {writeConcern: {w: expectedWritableNodes, wtimeout: 30 * 1000}};
+            var writeConcern = {
+                writeConcern:
+                    {w: expectedWritableNodesCount, wtimeout: replTest.kDefaultTimeoutMS}
+            };
             assert.writeOK(primary.getDB('foo').bar.insert({x: 100}, writeConcern));
+            jsTestLog('ensurePrimary - Successfully written a document to primary node (' +
+                      replTest.nodes[nodeId].host + ') using a write concern of w:' +
+                      expectedWritableNodesCount);
             return primary;
         };
 
         // 2 should eventually stage a priority takeover from the primary.
-        var primary = ensurePrimary(2, 3);
+        var primary = ensurePrimary(2, replTest.nodes);
 
         jsTestLog('primary is now 2');
         var config = assert.commandWorked(primary.adminCommand({replSetGetConfig: 1})).config;
@@ -183,7 +201,7 @@ var TagsTest = function(options) {
         var timeout = 20 * 1000;
 
         jsTestLog('test1');
-        primary = ensurePrimary(2, 3);
+        primary = ensurePrimary(2, replTest.nodes.slice(0, 3));
 
         jsTestLog('Non-existent write concern should be rejected.');
         options = {writeConcern: {w: 'blahblah', wtimeout: timeout}};
@@ -204,7 +222,7 @@ var TagsTest = function(options) {
         conns[1].reconnect(conns[4]);
         jsTestLog('partitions: [0-1-2] [1-4] [3] ' +
                   '(all nodes besides node 3 can replicate from primary node 2)');
-        primary = ensurePrimary(2, 4);
+        primary = ensurePrimary(2, replTest.nodes.slice(0, 3), 4);
 
         jsTestLog('Write concern "3 or 4" should work - 4 is now connected to the primary ' +
                   primary.host + ' via node 1 ' + replTest.nodes[1].host);
@@ -223,7 +241,7 @@ var TagsTest = function(options) {
         conns[3].reconnect(conns[4]);
         jsTestLog('partitions: [0-1-2] [1-4] [3-4] ' +
                   '(all secondaries can replicate from primary node 2)');
-        primary = ensurePrimary(2, 5);
+        primary = ensurePrimary(2, replTest.nodes.slice(0, 3), replTest.nodes.length);
 
         jsTestLog('Write concern "3 and 4" should work - ' +
                   'nodes 3 and 4 are connected to primary via node 1.');
@@ -242,7 +260,7 @@ var TagsTest = function(options) {
         assert.writeOK(primary.getDB('foo').bar.insert(doc, options));
 
         jsTestLog('Write concern "2 dc and 3 server"');
-        primary = ensurePrimary(2, 5);
+        primary = ensurePrimary(2, replTest.nodes.slice(0, 3), replTest.nodes.length);
         options = {writeConcern: {w: '2 dc and 3 server', wtimeout: timeout}};
         assert.writeOK(primary.getDB('foo').bar.insert(doc));
         assert.writeOK(primary.getDB('foo').bar.insert(doc, options));
@@ -260,7 +278,7 @@ var TagsTest = function(options) {
         // Node 1 with slightly higher priority will take over.
         jsTestLog('1 must become primary here because otherwise the other members will take too ' +
                   'long timing out their old sync threads');
-        primary = ensurePrimary(1, 4);
+        primary = ensurePrimary(1, replTest.nodes.slice(0, 2), 4);
 
         jsTestLog('Write concern "3 and 4" should still work with new primary node 1 ' +
                   primary.host);
diff --git a/jstests/replsets/step_down_during_draining.js b/jstests/replsets/step_down_during_draining.js
index 95a3b2946f8..6c140b04938 100644
--- a/jstests/replsets/step_down_during_draining.js
+++ b/jstests/replsets/step_down_during_draining.js
@@ -50,8 +50,13 @@ load("jstests/replsets/rslib.js");
             node.adminCommand({configureFailPoint: 'rsSyncApplyStop', mode: 'off'}));
     }
 
-    function stepUp(node) {
-        assert.commandWorked(node.adminCommand({replSetStepUp: 1}));
+    function stepUpNode(node) {
+        assert.soonNoExcept(function() {
+            assert.commandWorked(node.adminCommand({replSetStepUp: 1}));
+            replSet.awaitNodesAgreeOnPrimary(
+                replSet.kDefaultTimeoutMS, replSet.nodes, replSet.getNodeId(node));
+            return node.adminCommand('replSetGetStatus').myState == ReplSetTest.State.PRIMARY;
+        }, 'failed to step up node ' + node.host, replSet.kDefaultTimeoutMS);
     }
 
     // Do an initial insert to prevent the secondary from going into recovery
@@ -85,9 +90,7 @@ load("jstests/replsets/rslib.js");
         1000);
 
     reconnect(secondary);
-    stepUp(secondary);
-    replSet.waitForState(secondary, ReplSetTest.State.PRIMARY);
-    replSet.awaitNodesAgreeOnPrimary();
+    stepUpNode(secondary);
 
     // Secondary doesn't allow writes yet.
     var res = secondary.getDB("admin").runCommand({"isMaster": 1});
@@ -103,14 +106,10 @@ load("jstests/replsets/rslib.js");
 
     // Original primary steps up.
     reconnect(primary);
-    stepUp(primary);
-    replSet.waitForState(primary, ReplSetTest.State.PRIMARY);
-    replSet.awaitNodesAgreeOnPrimary();
+    stepUpNode(primary);
 
     reconnect(secondary);
-    stepUp(secondary);
-    replSet.waitForState(secondary, ReplSetTest.State.PRIMARY);
-    replSet.awaitNodesAgreeOnPrimary();
+    stepUpNode(secondary);
 
     // Disable fail point to allow replication.
     secondaries.forEach(disableFailPoint);