1 files changed, 30 insertions, 9 deletions

diff --git a/jstests/replsets/rollback3.js b/jstests/replsets/rollback3.js
index 5c2f2f1..fa923d8 100755
--- a/jstests/replsets/rollback3.js
+++ b/jstests/replsets/rollback3.js
@@ -30,10 +30,10 @@ function wait(f) {
         if (n % 4 == 0)
             print("rollback3.js waiting " + w);
         if (++n == 4) {
-            print("" + f);

-        }

-        if (n == 200) {

-            print("rollback3.js failing waited too long");

+            print("" + f);
+        }
+        if (n == 200) {
+            print("rollback3.js failing waited too long");
             throw "wait error";
         }
         sleep(1000);
@@ -188,15 +188,20 @@ doTest = function (signal) {
     wait(function () { return b.bar.count() == a.bar.count(); });
 
     A.runCommand({ replSetTest: 1, blind: true });
-    wait(function () { return B.isMaster().ismaster; });
+    reconnect(a,b);
+    wait(function () { try { return B.isMaster().ismaster; } catch(e) { return false; } });
 
     doItemsToRollBack(b);
 
     // a should not have the new data as it was in blind state.
     B.runCommand({ replSetTest: 1, blind: true });
+    reconnect(a,b);
+
     A.runCommand({ replSetTest: 1, blind: false });
-    wait(function () { return !B.isMaster().ismaster; });
-    wait(function () { return A.isMaster().ismaster; });
+    reconnect(a,b);
+    
+    wait(function () { try { return !B.isMaster().ismaster; } catch(e) { return false; } });
+    wait(function () { try { return A.isMaster().ismaster; } catch(e) { return false; } });
 
     assert(a.bar.count() >= 1, "count check");
     doWritesToKeep2(a);
@@ -207,18 +212,34 @@ doTest = function (signal) {
     // bring B back online
     // as A is primary, B will roll back and then catch up
     B.runCommand({ replSetTest: 1, blind: false });
+    reconnect(a,b);
 
     wait(function () { return B.isMaster().ismaster || B.isMaster().secondary; });
 
     // everyone is up here...
     assert(A.isMaster().ismaster || A.isMaster().secondary, "A up");
     assert(B.isMaster().ismaster || B.isMaster().secondary, "B up");
-
+    replTest.awaitReplication();
+    
     assert( dbs_match(a,b), "server data sets do not match after rollback, something is wrong");
 
     pause("rollback3.js SUCCESS");
     replTest.stopSet(signal);
-}
+};
+
+
+var reconnect = function(a,b) {
+  wait(function() { 
+      try {
+        a.bar.stats();
+        b.bar.stats();
+        return true;
+      } catch(e) {
+        print(e);
+        return false;
+      }
+    });
+};
 
 print("rollback3.js");
 doTest( 15 );