fixes to make profiler work

server/src/main/java/com/metamx/druid/master/CostBalancerStrategy.java

@@ -48,7 +48,7 @@ public class CostBalancerStrategy implements BalancerStrategy
   )
   {
     ServerHolder holder= chooseBestServer(proposalSegment, serverHolders, false).rhs;
-    if (!holder.isServingSegment(proposalSegment))
+    if (holder!=null && !holder.isServingSegment(proposalSegment))
     {
       return holder;
     }

server/src/main/java/com/metamx/druid/master/DruidMasterRuleRunner.java

@@ -68,7 +68,6 @@ public class DruidMasterRuleRunner implements DruidMasterHelper
     DatabaseRuleManager databaseRuleManager = paramsWithReplicationManager.getDatabaseRuleManager();
     for (DataSegment segment : paramsWithReplicationManager.getAvailableSegments()) {
       List<Rule> rules = databaseRuleManager.getRulesWithDefault(segment.getDataSource());
-
       boolean foundMatchingRule = false;
       for (Rule rule : rules) {
         if (rule.appliesTo(segment, now)) {

server/src/test/java/com/metamx/druid/master/DruidMasterBalancerTester.java

@@ -1,10 +1,14 @@
 package com.metamx.druid.master;
 
+import com.google.common.collect.Maps;
 import com.metamx.druid.client.DataSegment;
 import com.metamx.druid.client.DruidServer;
 
+import java.util.Map;
+
 public class DruidMasterBalancerTester extends DruidMasterBalancer
 {
+  private final Map<String,Integer> serverMap = Maps.newHashMap();
   public DruidMasterBalancerTester(DruidMaster master)
   {
     super(master);
@@ -45,7 +49,16 @@ public class DruidMasterBalancerTester extends DruidMasterBalancer
           }
         });
 
-        currentlyMovingSegments.get("normal").put(segmentName, segment);
+        if (serverMap.get(fromServerName)==null)
+        {
+          serverMap.put(fromServerName,segment.getFromServer().getSegments().size());
+        }
+        if (serverMap.get(toServerName)==null)
+        {
+          serverMap.put(toServerName,0);
+        }
+        serverMap.put(fromServerName,serverMap.get(fromServerName)-1);
+        serverMap.put(toServerName, serverMap.get(toServerName)+1);
       }
       catch (Exception e) {
         log.info(e, String.format("[%s] : Moving exception", segmentName));
@@ -54,4 +67,25 @@ public class DruidMasterBalancerTester extends DruidMasterBalancer
       currentlyMovingSegments.get("normal").remove(segment);
     }
   }
+
+  public Map<String,Integer> getServerMap()
+  {
+    return this.serverMap;
+  }
+
+  public boolean isBalanced(int min, int numServers)
+  {
+    if (serverMap.size()==numServers)
+    {
+      for (int numSegments: serverMap.values())
+      {
+        if (numSegments<min)
+        {
+          return false;
+        }
+      }
+      return true;
+    }
+    return false;
+  }
 }

server/src/test/java/com/metamx/druid/utils/DruidMasterBalancerProfiler.java

@@ -28,7 +28,6 @@ import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
 
 public class DruidMasterBalancerProfiler
 {
@@ -61,6 +60,7 @@ public class DruidMasterBalancerProfiler
         EasyMock.<String>anyObject(),
         EasyMock.<LoadPeonCallback>anyObject()
     );
+
     EasyMock.expectLastCall().anyTimes();
     EasyMock.replay(master);
 
@@ -86,7 +86,7 @@ public class DruidMasterBalancerProfiler
       );
     }
 
-    for (int i=0;i<numServers;i++)
+    for (int i=0;i<50;i++)
     {
       DruidServer server =EasyMock.createMock(DruidServer.class);
       EasyMock.expect(server.getMetadata()).andReturn(null).anyTimes();
@@ -132,9 +132,8 @@ public class DruidMasterBalancerProfiler
                                 .build();
 
     DruidMasterBalancerTester tester = new DruidMasterBalancerTester(master);
-
     watch.start();
-    for (int i=0;i<numServers;i++)
+    while (!tester.isBalanced(20,50))
     {
       params = tester.run(params);
     }
@@ -195,8 +194,8 @@ public class DruidMasterBalancerProfiler
                               .withMaxSegmentsToMove(MAX_SEGMENTS_TO_MOVE)
                               .withBalancerReferenceTimestamp(new DateTime("2013-01-01"))
                               .build();
-    params = new DruidMasterBalancerTester(master).run(params);
-    System.out.println(stopwatch.elapsed(TimeUnit.MILLISECONDS));
+    DruidMasterBalancerTester tester = new DruidMasterBalancerTester(master);
+
  }
 
 }