2 files changed, 55 insertions, 17 deletions

diff --git a/s/balancer_policy.cpp b/s/balancer_policy.cpp
index d20f2023e73..f92cb41aaf6 100644
--- a/s/balancer_policy.cpp
+++ b/s/balancer_policy.cpp
@@ -30,7 +30,6 @@ namespace mongo {
                                                         const map< string, BSONObj>& shardLimitsMap,  
                                                         const map< string,vector<BSONObj> >& shardToChunksMap, 
                                                         int balancedLastTime ){
-        
         pair<string,unsigned> min("",9999999);
         pair<string,unsigned> max("",0);
         
@@ -38,8 +37,14 @@ namespace mongo {
             const string& shard = i->first;
             unsigned size = i->second.size();
 
+            BSONObj shardLimits;
+            map< string,BSONObj >::const_iterator it = shardLimitsMap.find( shard );
+            if ( it != shardLimitsMap.end() ) {
+                shardLimits = it->second;
+            }
+
             if ( size < min.second ){
-                if ( isReceiver( shard , shardLimitsMap ) ){
+                if ( isReceiver( shardLimits ) ){
                     min.first = shard;
                     min.second = size;
                 } else {
@@ -47,7 +52,7 @@ namespace mongo {
                 }
             }
             
-            if ( size > max.second ){
+            if ( isDraining( shardLimits ) || ( size > max.second )){
                 max.first = shard;
                 max.second = size;
             }
@@ -85,16 +90,19 @@ namespace mongo {
         return from[0];
     }
 
-    bool BalancerPolicy::isReceiver( const string& shard, const map< string,BSONObj>& shardLimitsMap ){
+    bool BalancerPolicy::isReceiver( BSONObj limits ){
+
+        // A draining shard can never be a receiver
+        if ( isDraining( limits ) ){
+            return false;
+        }
 
         // If there's no limit information for the shard, assume it can be a chunk receiver 
         // (i.e., there's not bound on space utilization)
-        map< string,BSONObj >::const_iterator it = shardLimitsMap.find( shard );
-        if ( it == shardLimitsMap.end() ){
+        if ( limits.isEmpty() ){
             return true;
         }
 
-        BSONObj limits = it->second;
         long long maxUsage = limits["maxSize"].Long();
         if ( maxUsage == 0 ){
             return true;
@@ -108,20 +116,41 @@ namespace mongo {
         return false;
     }
 
+    bool BalancerPolicy::isDraining( BSONObj limits ){
+
+        // If there's no entry saying it is draining, it isn't.
+        if ( limits.isEmpty() || limits[ "draining" ].eoo() ){
+            return false;
+        }
+
+        return true;
+    }
+
     class PolicyObjUnitTest : public UnitTest {
     public:
         void maxSizeForShard(){
             BSONObj shard0 = BSON( "maxSize" << 0LL << "currSize" << 0LL );
+            assert( BalancerPolicy::isReceiver( shard0 ) );
+
             BSONObj shard1 = BSON( "maxSize" << 100LL << "currSize" << 80LL );
+            assert( BalancerPolicy::isReceiver( shard1 ) );
+
             BSONObj shard2 = BSON( "maxSize" << 100LL << "currSize" << 110LL );
-            map< string, BSONObj > shardLimitsMap; 
-            shardLimitsMap["shard0"] = shard0;
-            shardLimitsMap["shard1"] = shard1;
-            shardLimitsMap["shard2"] = shard2;
-
-            assert( BalancerPolicy::isReceiver( "shard0", shardLimitsMap ) );
-            assert( BalancerPolicy::isReceiver( "shard1", shardLimitsMap ) );
-            assert( ! BalancerPolicy::isReceiver( "shard2", shardLimitsMap ) );
+            assert( ! BalancerPolicy::isReceiver( shard2 ) );
+
+            BSONObj shard3 = BSON( "draining" << true );
+            assert( ! BalancerPolicy::isReceiver( shard3 ) );
+
+            BSONObj empty;
+            assert( BalancerPolicy::isReceiver( empty ) );
+        }
+
+        void drainingShard(){
+            BSONObj shard0 = BSON( "draining" << true );
+            assert( BalancerPolicy::isDraining( shard0 ) );
+
+            BSONObj empty;
+            assert( ! BalancerPolicy::isDraining( empty ) );
         }
 
         void run(){
diff --git a/s/balancer_policy.h b/s/balancer_policy.h
index d5333d3e658..b78e7da54ac 100644
--- a/s/balancer_policy.h
+++ b/s/balancer_policy.h
@@ -39,7 +39,6 @@ namespace mongo {
          * is { }. 
          * @param balancedLastTime is the number of chunks effectively moved in the last round.
          * @returns NULL or ChunkInfo of the best move to make towards balacing the collection.
-         *
          */
         static ChunkInfo* balance( const string& ns, const map< string, BSONObj>& shardLimitsMap,  
                                    const map< string,vector<BSONObj> >& shardToChunksMap, int balancedLastTime );
@@ -48,7 +47,17 @@ namespace mongo {
 
         static BSONObj pickChunk( const vector<BSONObj>& from, const vector<BSONObj>& to );
 
-        static bool isReceiver( const string& shard, const map< string,BSONObj>& shardLimitsMap );
+        /**
+         * Returns true if a shard can receive new chunks based on that shard's 'shardLimits' information.
+         * Expects the optional fields "maxSize", can in size in MB, and "usedSize", currently used size in MB,
+         * on 'shardLimits'.
+         */
+        static bool isReceiver( BSONObj shardLimits );
+
+        /**
+         * Returns true if 'shardLimist' contains a field "draining".
+         */
+        static bool isDraining( BSONObj shardLimits );
     };
 
     struct BalancerPolicy::ChunkInfo {