SERVER-6006 import code and fixes from 2.1.x into test

The addition of a call to s.stopBalancer() failed because that
function didn't exist in 2.0.  Bring in all the code that is
needed to make it work in 2.0.  Also pick up assert changes from
2.1 that make the test more reliable.

1 files changed, 152 insertions, 3 deletions

diff --git a/jstests/slowNightly/sharding_rs2.js b/jstests/slowNightly/sharding_rs2.js
index e8649f3ad3b..11a9a06fe25 100644
--- a/jstests/slowNightly/sharding_rs2.js
+++ b/jstests/slowNightly/sharding_rs2.js
@@ -84,7 +84,7 @@ after = rs.test.getMaster().adminCommand( "serverStatus" ).opcounters
 printjson( before )
 printjson( after )
 
-assert.eq( before.query + 10 , after.query , "B3" )
+assert.lte( before.query + 10 , after.query , "B3" )
 
 // --- add more data ----
 
@@ -110,6 +110,155 @@ assert.eq( 100 , db.foo.count() , "C1" )
 s.adminCommand( { enablesharding : "test" } );
 s.adminCommand( { shardcollection : "test.foo" , key : { x : 1 } } );
 
+// Changes in this test were backported from 2.1.x into 2.0.x, but the supporting infrastructure did
+// not exists in 2.0, so ... create the functions we need for this test
+//
+if ( ! sh.waitForPingChange ) {
+    sh.waitForPingChange = function( activePings, timeout, interval ){
+    
+        var isPingChanged = function( activePing ){
+            var newPing = db.getSisterDB( "config" ).mongos.findOne({ _id : activePing._id })
+            return ! newPing || newPing.ping + "" != activePing.ping + ""
+        }
+    
+        // First wait for all active pings to change, so we're sure a settings reload
+        // happened
+    
+        // Timeout all pings on the same clock
+        var start = new Date()
+    
+        var remainingPings = []
+        for( var i = 0; i < activePings.length; i++ ){
+        
+            var activePing = activePings[ i ]
+            print( "Waiting for active host " + activePing._id + " to recognize new settings... (ping : " + activePing.ping + ")" )
+       
+            // Do a manual timeout here, avoid scary assert.soon errors
+            var timeout = timeout || 30000;
+            var interval = interval || 200;
+            while( isPingChanged( activePing ) != true ){
+                if( ( new Date() ).getTime() - start.getTime() > timeout ){
+                    print( "Waited for active ping to change for host " + activePing._id + 
+                           ", a migration may be in progress or the host may be down." )
+                    remainingPings.push( activePing )
+                    break
+                }
+                sleep( interval )   
+            }
+    
+        }
+    
+        return remainingPings
+    }
+}
+
+if ( ! sh.waitForBalancerOff ) {
+    sh.waitForBalancerOff = function( timeout, interval ){
+    
+        var pings = db.getSisterDB( "config" ).mongos.find().toArray()
+        var activePings = []
+        for( var i = 0; i < pings.length; i++ ){
+            if( ! pings[i].waiting ) activePings.push( pings[i] )
+        }
+    
+        print( "Waiting for active hosts..." )
+    
+        activePings = sh.waitForPingChange( activePings, 60 * 1000 )
+    
+        // After 1min, we assume that all hosts with unchanged pings are either 
+        // offline (this is enough time for a full errored balance round, if a network
+        // issue, which would reload settings) or balancing, which we wait for next
+        // Legacy hosts we always have to wait for
+    
+        print( "Waiting for the balancer lock..." )
+    
+        // Wait for the balancer lock to become inactive
+        // We can guess this is stale after 15 mins, but need to double-check manually
+        try{ 
+            sh.waitForDLock( "balancer", false, 15 * 60 * 1000 )
+        }
+        catch( e ){
+            print( "Balancer still may be active, you must manually verify this is not the case using the config.changelog collection." )
+            throw e
+        }
+        
+        print( "Waiting again for active hosts after balancer is off..." )
+    
+        // Wait a short time afterwards, to catch the host which was balancing earlier
+        activePings = sh.waitForPingChange( activePings, 5 * 1000 )
+    
+        // Warn about all the stale host pings remaining
+        for( var i = 0; i < activePings.length; i++ ){
+            print( "Warning : host " + activePings[i]._id + " seems to have been offline since " + activePings[i].ping )
+        }
+    
+    }
+}
+
+if ( ! sh.waitForDLock ) {
+    sh.waitForDLock = function( lockId, onOrNot, timeout, interval ){
+    
+        // Wait for balancer to be on or off
+        // Can also wait for particular balancer state
+        var state = onOrNot
+    
+        var beginTS = undefined
+        if( state == undefined ){
+            var currLock = db.getSisterDB( "config" ).locks.findOne({ _id : lockId })
+            if( currLock != null ) beginTS = currLock.ts
+        }
+        
+        var lockStateOk = function(){
+            var lock = db.getSisterDB( "config" ).locks.findOne({ _id : lockId })
+
+            if( state == false ) return ! lock || lock.state == 0
+            if( state == true ) return lock && lock.state == 2
+            if( state == undefined ) return (beginTS == undefined && lock) || 
+                                            (beginTS != undefined && ( !lock || lock.ts + "" != beginTS + "" ) )
+            else return lock && lock.state == state
+        }
+    
+        assert.soon( lockStateOk,
+                     "Waited too long for lock " + lockId + " to " + 
+                          (state == true ? "lock" : ( state == false ? "unlock" : 
+                                           "change to state " + state ) ),
+                     timeout,
+                     interval
+        )
+    }
+}
+
+if ( ! sh.waitForBalancer ) {
+    sh.waitForBalancer = function( onOrNot, timeout, interval ){
+    
+        // If we're waiting for the balancer to turn on or switch state or
+        // go to a particular state
+        if( onOrNot ){
+            // Just wait for the balancer lock to change, can't ensure we'll ever see it
+            // actually locked
+            sh.waitForDLock( "balancer", undefined, timeout, interval )
+        }
+        else {
+            // Otherwise we need to wait until we're sure balancing stops
+            sh.waitForBalancerOff( timeout, interval )
+        }
+    
+    }
+}
+
+if ( ! s.stopBalancer ) {
+    s.stopBalancer = function( timeout, interval ) {
+        this.setBalancer( false )
+    
+        if( typeof db == "undefined" ) db = undefined
+        var oldDB = db
+    
+        db = this.config
+        sh.waitForBalancer( false, timeout, interval )
+        db = oldDB
+    }
+}
+
 // We're doing some manual chunk stuff, so stop the balancer first
 s.stopBalancer()
 
@@ -143,7 +292,7 @@ for ( i=0; i<10; i++ )
 
 after = rs.test.getMaster().adminCommand( "serverStatus" ).opcounters
 
-assert.eq( before.query + 10 , after.query , "D3" )
+assert.lte( before.query + 10 , after.query , "D3" )
 
 // by shard key
 
@@ -161,7 +310,7 @@ for ( i=0; i<10; i++ )
 
 after = rs.test.getMaster().adminCommand( "serverStatus" ).opcounters
 
-assert.eq( before.query + 10 , after.query , "E3" )
+assert.lte( before.query + 10 , after.query , "E3" )
 
 assert.eq( 100 , ts.count() , "E4" )
 assert.eq( 100 , ts.find().itcount() , "E5" )