From ba82fe746a0044ffcc435a2b3bbf57071d0fff96 Mon Sep 17 00:00:00 2001
From: Nelson Ray <ncray86@gmail.com>
Date: Mon, 7 Jan 2013 16:28:27 -0800
Subject: [PATCH] put MAX_SEGMENTS_TO_MOVE in the config

---
 .../metamx/druid/master/BalancerCostAnalyzer.java |  5 +++--
 .../java/com/metamx/druid/master/DruidMaster.java |  7 ++++++-
 .../metamx/druid/master/DruidMasterConfig.java    |  4 ++++
 .../com/metamx/druid/master/rules/LoadRule.java   | 15 ++++++++++++---
 4 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/com/metamx/druid/master/BalancerCostAnalyzer.java b/server/src/main/java/com/metamx/druid/master/BalancerCostAnalyzer.java
index 3375e4b4c91..4a0f2b90bcd 100644
--- a/server/src/main/java/com/metamx/druid/master/BalancerCostAnalyzer.java
+++ b/server/src/main/java/com/metamx/druid/master/BalancerCostAnalyzer.java
@@ -44,7 +44,7 @@ import java.util.Set;
 public class BalancerCostAnalyzer
 {
   private static final Logger log = new Logger(BalancerCostAnalyzer.class);
-  private static final int MAX_SEGMENTS_TO_MOVE = 5;
+  private final int MAX_SEGMENTS_TO_MOVE;
   private static final int DAY_IN_MILLIS = 1000 * 60 * 60 * 24;
   private static final int SEVEN_DAYS_IN_MILLIS = 7 * DAY_IN_MILLIS;
   private static final int THIRTY_DAYS_IN_MILLIS = 30 * DAY_IN_MILLIS;
@@ -58,9 +58,10 @@ public class BalancerCostAnalyzer
   private double totalCostChange;
   private int totalSegments;
 
-  public BalancerCostAnalyzer(DateTime referenceTimestamp)
+  public BalancerCostAnalyzer(DateTime referenceTimestamp, int MAX_SEGMENTS_TO_MOVE)
   {
     this.referenceTimestamp = referenceTimestamp;
+    this.MAX_SEGMENTS_TO_MOVE = MAX_SEGMENTS_TO_MOVE;
     rand = new Random(0);
     totalCostChange = 0;
   }
diff --git a/server/src/main/java/com/metamx/druid/master/DruidMaster.java b/server/src/main/java/com/metamx/druid/master/DruidMaster.java
index ff01957d102..d5833b203fe 100644
--- a/server/src/main/java/com/metamx/druid/master/DruidMaster.java
+++ b/server/src/main/java/com/metamx/druid/master/DruidMaster.java
@@ -132,6 +132,11 @@ public class DruidMaster
     return master;
   }
 
+  public DruidMasterConfig getConfig()
+  {
+    return config;
+  }
+
   public Map<String, Double> getLoadStatus()
   {
     Map<String, Integer> availableSegmentMap = Maps.newHashMap();
@@ -675,7 +680,7 @@ public class DruidMaster
               },
               new DruidMasterRuleRunner(DruidMaster.this),
               new DruidMasterCleanup(DruidMaster.this),
-              new DruidMasterBalancer(DruidMaster.this, new BalancerCostAnalyzer(DateTime.now())),
+              new DruidMasterBalancer(DruidMaster.this, new BalancerCostAnalyzer(DateTime.now(), config.getMaxSegmentsToMove())),
               new DruidMasterLogger()
           )
       );
diff --git a/server/src/main/java/com/metamx/druid/master/DruidMasterConfig.java b/server/src/main/java/com/metamx/druid/master/DruidMasterConfig.java
index 87df78d4f5e..274b3765164 100644
--- a/server/src/main/java/com/metamx/druid/master/DruidMasterConfig.java
+++ b/server/src/main/java/com/metamx/druid/master/DruidMasterConfig.java
@@ -80,4 +80,8 @@ public abstract class DruidMasterConfig
   {
     return Integer.MAX_VALUE;
   }
+
+  @Config("druid.master.balancer.maxSegmentsToMove")
+  @Default("5")
+  public abstract int getMaxSegmentsToMove();
 }
diff --git a/server/src/main/java/com/metamx/druid/master/rules/LoadRule.java b/server/src/main/java/com/metamx/druid/master/rules/LoadRule.java
index 748f66f135d..cd63fa69a7e 100644
--- a/server/src/main/java/com/metamx/druid/master/rules/LoadRule.java
+++ b/server/src/main/java/com/metamx/druid/master/rules/LoadRule.java
@@ -59,7 +59,15 @@ public abstract class LoadRule implements Rule
       return stats;
     }
 
-    stats.accumulate(assign(expectedReplicants, totalReplicants, serverQueue, segment));
+    stats.accumulate(
+        assign(
+            expectedReplicants,
+            totalReplicants,
+            serverQueue,
+            segment,
+            master.getConfig().getMaxSegmentsToMove()
+        )
+    );
     stats.accumulate(drop(expectedReplicants, clusterReplicants, segment, params));
 
     return stats;
@@ -69,7 +77,8 @@ public abstract class LoadRule implements Rule
       int expectedReplicants,
       int totalReplicants,
       MinMaxPriorityQueue<ServerHolder> serverQueue,
-      DataSegment segment
+      DataSegment segment,
+      int MAX_SEGMENTS_TO_MOVE
   )
   {
     MasterStats stats = new MasterStats();
@@ -78,7 +87,7 @@ public abstract class LoadRule implements Rule
 
     List<ServerHolder> assignedServers = Lists.newArrayList();
     while (totalReplicants < expectedReplicants) {
-      BalancerCostAnalyzer analyzer = new BalancerCostAnalyzer(DateTime.now());
+      BalancerCostAnalyzer analyzer = new BalancerCostAnalyzer(DateTime.now(), MAX_SEGMENTS_TO_MOVE);
       BalancerCostAnalyzer.BalancerCostAnalyzerHelper helper = analyzer.new BalancerCostAnalyzerHelper(serverHolderList, segment);
       helper.computeAllCosts();
       Pair<Double, ServerHolder> minPair = helper.getCostsServerHolderPairs().pollFirst();