SERVER-29093 double_rollback.js should ensure stable connections before running database commands

2 files changed, 49 insertions, 15 deletions

diff --git a/jstests/replsets/double_rollback.js b/jstests/replsets/double_rollback.js
index c7adc95415d..1b819cd2452 100644
--- a/jstests/replsets/double_rollback.js
+++ b/jstests/replsets/double_rollback.js
@@ -2,7 +2,7 @@
  * This test causes node 2 to enter rollback, reach the common point, and exit rollback, but before
  * it can apply operations to bring it back to a consistent state, switch sync sources to the node
  * that originally gave it the ops it is now rolling back (node 0).  This test then verifies that
- * node 2 refuses to use node0 as a sync source because it doesn't contain the minValid document
+ * node 2 refuses to use node 0 as a sync source because it doesn't contain the minValid document
  * it needs to reach consistency.  Node 2 is then allowed to reconnect to the node it was
  * originally rolling back against (node 1) and finish its rollback.  This is a regression test
  * against the case where we *did* allow node 2 to sync from node 0 which gave it the very ops
@@ -106,6 +106,10 @@
     checkLog.contains(
         nodes[2], "remote oplog does not contain entry with optime matching our required optime");
 
+    // Ensure our connection to node 0 is re-established, since our
+    // original connection should have gotten killed after node 0 stepped down.
+    reconnect(nodes[0]);
+
     var node0RBID = assert.commandWorked(nodes[0].adminCommand('replSetGetRBID')).rbid;
     var node1RBID = assert.commandWorked(nodes[1].adminCommand('replSetGetRBID')).rbid;
 
@@ -122,6 +126,10 @@
     waitForState(nodes[2], ReplSetTest.State.SECONDARY);
     rst.awaitReplication();
 
+    // Ensure our connection to node 0 is re-established, since our connection should have gotten
+    // killed during node 0's transition to ROLLBACK.
+    reconnect(nodes[0]);
+
     // Check that rollback happened on node 0, but not on node 2 since it had already rolled back
     // and just needed to finish applying ops to reach minValid.
     assert.neq(node0RBID, assert.commandWorked(nodes[0].adminCommand('replSetGetRBID')).rbid);
diff --git a/jstests/replsets/rslib.js b/jstests/replsets/rslib.js
index 5911723d717..5403b129b3c 100644
--- a/jstests/replsets/rslib.js
+++ b/jstests/replsets/rslib.js
@@ -52,19 +52,32 @@ var getLastOpTime;
         rst.awaitSyncSource(syncingNode, desiredSyncSource);
     };
 
-    wait = function(f, msg) {
+    /**
+     * Calls a function 'f' once a second until it returns true. Throws an exception once 'f' has
+     * been called more than 'retries' times without returning true. If 'retries' is not given,
+     * it defaults to 200. 'retries' must be an integer greater than or equal to zero.
+     */
+    wait = function(f, msg, retries) {
         w++;
         var n = 0;
+        var default_retries = 200;
+        var delay_interval_ms = 1000;
+
+        // Set default value if 'retries' was not given.
+        if (retries === undefined) {
+            retries = default_retries;
+        }
         while (!f()) {
-            if (n % 4 == 0)
-                print("waiting " + w);
+            if (n % 4 == 0) {
+                print("Waiting " + w);
+            }
             if (++n == 4) {
                 print("" + f);
             }
-            if (n >= 200) {
-                throw new Error('tried 200 times, giving up on ' + msg);
+            if (n >= retries) {
+                throw new Error('Tried ' + retries + ' times, giving up on ' + msg);
             }
-            sleep(1000);
+            sleep(delay_interval_ms);
         }
     };
 
@@ -85,19 +98,32 @@ var getLastOpTime;
         count++;
     };
 
-    reconnect = function(a) {
+    /**
+     * Attempt to re-establish and re-authenticate a Mongo connection if it was dropped, with
+     * multiple retries.
+     *
+     * Returns upon successful re-connnection. If connection cannot be established after 200
+     * retries, throws an exception.
+     *
+     * @param conn - a Mongo connection object or DB object.
+     */
+    reconnect = function(conn) {
+        var retries = 200;
         wait(function() {
             var db;
             try {
-                // make this work with either dbs or connections
-                if (typeof(a.getDB) == "function") {
-                    db = a.getDB('foo');
+                // Make this work with either dbs or connections.
+                if (typeof(conn.getDB) == "function") {
+                    db = conn.getDB('foo');
                 } else {
-                    db = a;
+                    db = conn;
                 }
+
+                // Run a simple command to re-establish connection.
                 db.bar.stats();
-                if (jsTest.options().keyFile) {  // SERVER-4241: Shell connections don't
-                                                 // re-authenticate on reconnect
+
+                // SERVER-4241: Shell connections don't re-authenticate on reconnect.
+                if (jsTest.options().keyFile) {
                     return jsTest.authenticate(db.getMongo());
                 }
                 return true;
@@ -105,7 +131,7 @@ var getLastOpTime;
                 print(e);
                 return false;
             }
-        });
+        }, retries);
     };
 
     getLatestOp = function(server) {