SERVER-950 Balancer policy now picks up size limits (but balancer is not passing it yet)

SConstruct

@@ -1141,7 +1141,7 @@ mongos = env.Program( "mongos" , commonFiles + coreDbFiles + coreServerFiles + s
 clientLibName = str( env.Library( "mongoclient" , allClientFiles )[0] )
 if GetOption( "sharedclient" ):
     sharedClientLibName = str( env.SharedLibrary( "mongoclient" , allClientFiles )[0] )
-env.Library( "mongotestfiles" , commonFiles + coreDbFiles + coreServerFiles + serverOnlyFiles + ["client/gridfs.cpp"])
+env.Library( "mongotestfiles" , commonFiles + coreDbFiles + coreServerFiles + serverOnlyFiles + shardServerFiles +["client/gridfs.cpp"])
 env.Library( "mongoshellfiles" , allClientFiles + coreServerFiles )
 
 clientTests = []

dbtests/balancer_tests.cpp

@@ -0,0 +1,55 @@
+// histogramtests.cpp : histogram.{h,cpp} unit tests
+
+/**
+ *    Copyright (C) 2010 10gen Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "../pch.h"
+
+#include "dbtests.h"
+#include "../s/balancer_policy.h"
+
+namespace mongo {
+
+    using mongo::BalancerPolicy;
+
+    class MaxSizeForShard{
+    public:
+        void run(){
+            BSONObj shard0 = BSON( "maxSize" << 0LL << "currSize" << 0LL );
+            BSONObj shard1 = BSON( "maxSize" << 100LL << "currSize" << 80LL );
+            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 ) );
+        }
+    };
+
+    class BalancerSuite : public Suite {
+    public:
+        BalancerSuite() : Suite( "balancer" ){}
+
+        void setupTests(){
+            add< MaxSizeForShard >();
+            // TODO: complete the test suite
+        } 
+    } BalancerSuite;
+
+}  // mongo namespace

s/balance.cpp

@@ -247,7 +247,7 @@ namespace mongo {
         map< string, BSONObj > shardLimitsMap; 
         for ( vector<Shard>::const_iterator it = allShards.begin(); it != allShards.end(); ++it ){
             const Shard& s = *it;
-            BSONObj limitsObj = BSON( "maxSize" << 0 << "currSize" << 0 /* TODO */);
+            BSONObj limitsObj = BSON( "maxSize" << 0LL << "currSize" << 0LL /* TODO */);
             shardLimitsMap[s.getName()] = limitsObj;
         }
 

s/balancer_policy.cpp

@@ -36,7 +36,7 @@ namespace mongo {
             const string& shard = i->first;
             unsigned size = i->second.size();
 
-            if ( size < min.second ){
+            if ( isReceiver( shard , shardLimitsMap ) && ( size < min.second ) ){
                 min.first = shard;
                 min.second = size;
             }
@@ -77,6 +77,29 @@ namespace mongo {
             return from[from.size()-1];
 
         return from[0];
-    }    
+    }
+
+    bool BalancerPolicy::isReceiver( const string& shard, const map< string,BSONObj>& shardLimitsMap ){
+
+        // 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() ){
+            return true;
+        }
+
+        BSONObj limits = it->second;
+        long long maxUsage = limits["maxSize"].Long();
+        if ( maxUsage == 0 ){
+            return true;
+        }
+
+        long long currUsage = limits["currSize"].Long();
+        if ( currUsage < maxUsage ){
+            return true;
+        }
+
+        return false;
+    }
 
 }  // namespace mongo

s/balancer_policy.h

@@ -47,6 +47,8 @@ namespace mongo {
         // below exposed for testing purposes only -- treat it as private --
 
         static BSONObj pickChunk( const vector<BSONObj>& from, const vector<BSONObj>& to );
+
+        static bool isReceiver( const string& shard, const map< string,BSONObj>& shardLimitsMap );
     };
 
     struct BalancerPolicy::ChunkInfo {