Integrated Hungarian Algo implementation by Wenting

src/mongo/SConscript

@@ -155,6 +155,7 @@ commonFiles = [ "pch.cpp",
                 "util/net/listen.cpp",
                 "util/startup_test.cpp",
                 "util/version.cpp",
+                "util/HungarianAlgo.cpp",
                 "client/connpool.cpp",
                 "client/dbclient.cpp",
                 "client/dbclient_rs.cpp",

src/mongo/s/commands_admin.cpp

@@ -48,6 +48,7 @@
 #include "mongo/util/stringutils.h"
 #include "mongo/util/timer.h"
 #include "mongo/util/version.h"
+#include "mongo/util/HungarianAlgo.h"
 
 namespace mongo {
 
@@ -1072,8 +1073,21 @@ namespace mongo {
 
 				// 5. Run the algorithm
 				log() << "[MYCODE_TIME] Running the algorithm" << endl;
+                bool loadBalance = cmdObj["loadBalance"].trueValue();
 				int assignment[numChunk];
-				runAlgorithm(splitPoints, ns, removedReplicas, numChunk, numShards, proposedKey, assignment);
+				long long **datainkr;
+                datainkr = new long long*[numChunk];
+				for (int i = 0; i < numChunk; i++)
+                    datainkr[i] = new long long[numShards];
+
+				for (int i = 0; i < numChunk; i++)
+					for (int j = 0; j < numShards; j++)
+						datainkr[i][j] = 0;
+
+                if (loadBalance)
+                    runLBAlgorithm(splitPoints, ns, removedReplicas, numChunk, numShards, proposedKey, assignment, datainkr);
+                else
+				    runAlgorithm(splitPoints, ns, removedReplicas, numChunk, numShards, proposedKey, assignment, datainkr);
 
 				log() << "[MYCODE_TIME] End of Algorithm Phase:\tmillis:" << t.millis() << endl;
 
@@ -1484,17 +1498,54 @@ namespace mongo {
 				}
 			}
 
-			void runAlgorithm(BSONObjSet splitPoints, string ns, string replicas[], int numChunk, int numShards, BSONObj proposedKey, int assignment[])
+            void runLBAlgorithm(BSONObjSet splitPoints, string ns, string replicas[], int numChunk, int numShards, BSONObj proposedKey, int assignment[],long long **datainkr)
+            {
+                printf("[MYCODE] RUN-LOADBALANCE-ALGORITHM\n");
+
+                collectData(splitPoints, ns, replicas, numChunk, numShards, proposedKey, datainkr);
+                int chunkpershard = (int)ceil((double)numChunk/numShards);
+                long long **newdatainkr = new long long*[numChunk];
+                for (int i = 0; i < numChunk; i++)
+                    newdatainkr[i] = new long long[numShards*chunkpershard];
+
+                for (int i = 0; i < numChunk; i++)
+                    for (int j = 0; j < numShards; j++)
+                        for (int k = 0; k < chunkpershard; k++) 
+                            newdatainkr[i][j*chunkpershard+k] = datainkr[i][j];
+
+                cout << "[MYCODE] NEWDATAINKR: with "<< numChunk<< " chunks and "<< chunkpershard<< " ChunkPerShard" << endl;
+                for (int i = 0; i < numChunk; i++)
+                {
+                    cout << "[MYCODE] ";
+                    for (int j = 0; j < numShards*chunkpershard; j++)
+                        cout << newdatainkr[i][j] << "\t";
+                    cout << endl;
+                }
+
+                HungarianAlgo* algo = new HungarianAlgo();
+                algo->max_cost_assignment(newdatainkr,numChunk,numShards*chunkpershard,assignment);
+                cout << "[MYCODE] LOAD BAlANCE ASSIGNMENT:\n [MYCODE] ";
+                for (int i = 0; i < numChunk; i++)
+                    cout << assignment[i] << "\t";
+                cout << "\n";
+
+                for (int i = 0; i < numChunk; i++)
+                {
+                    assignment[i] = (int)floor(assignment[i]/chunkpershard);
+                }
+
+                cout << "[MYCODE] ASSIGNMENT:\n [MYCODE] ";
+                for (int i = 0; i < numChunk; i++)
+                    cout << assignment[i] << "\t";
+                cout << "\n";
+
+                delete[] newdatainkr;
+                delete algo;
+          }
+
+		    void runAlgorithm(BSONObjSet splitPoints, string ns, string replicas[], int numChunk, int numShards, BSONObj proposedKey, int assignment[], long long **datainkr)
 			{
 				printf("[MYCODE] RUNALGORITHM\n");
-				long long **datainkr;
-                datainkr = new long long*[numChunk];
-				for (int i = 0; i < numChunk; i++)
-                    datainkr[i] = new long long[numShards];
-
-				for (int i = 0; i < numChunk; i++)
-					for (int j = 0; j < numShards; j++)
-						datainkr[i][j] = 0;
 
                 collectData(splitPoints, ns, replicas, numChunk, numShards, proposedKey, datainkr);
 

src/mongo/shell/utils_sh.js

@@ -35,7 +35,7 @@ sh.help = function() {
     print( "\tsh.addShard( host )                       server:port OR setname/server:port" )
     print( "\tsh.enableSharding(dbname)                 enables sharding on the database dbname" )
     print( "\tsh.shardCollection(fullName,key,unique)   shards the collection" );
-    print( "\tsh.reShardCollection(fullName,key,unique) reshards the collection" );
+    print( "\tsh.reShardCollection(fullName,key,loadbalance,unique) reshards the collection" );
 
     print( "\tsh.splitFind(fullName,find)               splits the chunk that find is in at the median" );
     print( "\tsh.splitAt(fullName,middle)               splits the chunk that middle is in at middle" );
@@ -78,7 +78,7 @@ sh.shardCollection = function( fullName , key , unique ) {
     return sh._adminCommand( cmd );
 }
 
-sh.reShardCollection = function( fullName , key , unique ) {
+sh.reShardCollection = function( fullName , key , loadbalance, unique ) {
     sh._checkFullName( fullName )
     assert( key , "need a key" )
     assert( typeof( key ) == "object" , "key needs to be an object" )
@@ -87,6 +87,9 @@ sh.reShardCollection = function( fullName , key , unique ) {
     if ( unique ) 
         cmd.unique = true;
 
+    if (loadbalance)
+        cmd.loadBalance = true;
+
     return sh._adminCommand( cmd );
 }