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

(cherry picked from commit 8d146d21661cf755fc15b7eefec147ce85ffde8c)

3 files changed, 57 insertions, 43 deletions

diff --git a/jstests/replsets/chaining_removal.js b/jstests/replsets/chaining_removal.js
index 744b6bb8de5..2c42b3f83e4 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, 60 * 1000);
-    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 076b621c523..80ea4264fc3 100644
--- a/jstests/replsets/libs/tags.js
+++ b/jstests/replsets/libs/tags.js
@@ -138,23 +138,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 (forceWriteMode) {
                 primary.forceWriteMode(forceWriteMode);
             }
             var writeConcern = {
-                writeConcern: {w: expectedWritableNodes, wtimeout: 30 * 1000}
+                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;
@@ -188,7 +204,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 = {
@@ -213,7 +229,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);
@@ -236,7 +252,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.');
@@ -261,7 +277,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}
         };
@@ -281,7 +297,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/src/mongo/shell/replsettest.js b/src/mongo/shell/replsettest.js
index 0b1419caa07..040fa531021 100644
--- a/src/mongo/shell/replsettest.js
+++ b/src/mongo/shell/replsettest.js
@@ -445,45 +445,44 @@ var ReplSetTest = function(opts) {
 
     /**
      * Blocks until all nodes agree on who the primary is.
+     * If 'expectedPrimaryNodeId' is provided, ensure that every node is seeing this node as the
+     * primary. Otherwise, ensure that all the nodes in the set agree with the first node on the
+     * identity of the primary.
      */
-    this.awaitNodesAgreeOnPrimary = function(timeout) {
+    this.awaitNodesAgreeOnPrimary = function(timeout, nodes, expectedPrimaryNodeId) {
         timeout = timeout || self.kDefaultTimeoutMS;
+        nodes = nodes || self.nodes;
+        expectedPrimaryNodeId = expectedPrimaryNodeId || -1;
 
-        assert.soon(function() {
-            try {
-                var primary = -1;
-
-                for (var i = 0; i < self.nodes.length; i++) {
-                    var replSetGetStatus =
-                        self.nodes[i].getDB("admin").runCommand({replSetGetStatus: 1});
-                    var nodesPrimary = -1;
-                    for (var j = 0; j < replSetGetStatus.members.length; j++) {
-                        if (replSetGetStatus.members[j].state === ReplSetTest.State.PRIMARY) {
-                            // Node sees two primaries.
-                            if (nodesPrimary !== -1) {
-                                return false;
-                            }
-                            nodesPrimary = j;
-                        }
-                    }
-                    // Node doesn't see a primary.
-                    if (nodesPrimary < 0) {
-                        return false;
-                    }
+        assert.soonNoExcept(function() {
+            var primary = expectedPrimaryNodeId;
 
-                    if (primary < 0) {
-                        // If we haven't seen a primary yet, set it to this.
-                        primary = nodesPrimary;
-                    } else if (primary !== nodesPrimary) {
-                        return false;
+            for (var i = 0; i < nodes.length; i++) {
+                var replSetGetStatus = nodes[i].getDB("admin").runCommand({replSetGetStatus: 1});
+                var nodesPrimary = -1;
+                for (var j = 0; j < replSetGetStatus.members.length; j++) {
+                    if (replSetGetStatus.members[j].state === ReplSetTest.State.PRIMARY) {
+                        // Node sees two primaries.
+                        if (nodesPrimary !== -1) {
+                            return false;
+                        }
+                        nodesPrimary = j;
                     }
                 }
+                // Node doesn't see a primary.
+                if (nodesPrimary < 0) {
+                    return false;
+                }
 
-                return true;
-            } catch (e) {
-                print("caught exception " + e);
-                return false;
+                if (primary < 0) {
+                    // If we haven't seen a primary yet, set it to this.
+                    primary = nodesPrimary;
+                } else if (primary !== nodesPrimary) {
+                    return false;
+                }
             }
+
+            return true;
         }, "Awaiting nodes to agree on primary", timeout);
     };