From a26a5583323f8aa059cf930ac7840e3d51a1b62e Mon Sep 17 00:00:00 2001
From: Jay J Wylie <jay.j.wylie@gmail.com>
Date: Sun, 23 Jun 2013 09:11:25 -0700
Subject: [PATCH] Tweaks to PartitionBalance and RebalancePlan

PartitionBalance
- calculate partition-stores per zone.
- This measure provides more context to evaluate the size of any plans to rebalance the cluster.

RebalancePlanCLI
- fix typo in verbose usage message
---
 .../voldemort/tools/PartitionBalance.java     | 21 +++++++++++++++++++
 .../voldemort/tools/RebalancePlanCLI.java     |  2 +-
 .../voldemort/tools/PartitionBalanceTest.java | 17 ++++++++++-----
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/java/voldemort/tools/PartitionBalance.java b/src/java/voldemort/tools/PartitionBalance.java
index b47537ca61..5315a161ab 100644
--- a/src/java/voldemort/tools/PartitionBalance.java
+++ b/src/java/voldemort/tools/PartitionBalance.java
@@ -60,6 +60,7 @@ public class PartitionBalance {
     private final Map<Integer, Integer> primaryAggNodeIdToPartitionCount;
     private final Map<Integer, Integer> aggNodeIdToZonePrimaryCount;
     private final Map<Integer, Integer> allAggNodeIdToPartitionCount;
+    private final Map<Integer, Integer> zoneIdToPartitionStoreCount;
 
     public PartitionBalance(Cluster cluster, List<StoreDefinition> storeDefs) {
         this.cluster = cluster;
@@ -69,6 +70,7 @@ public PartitionBalance(Cluster cluster, List<StoreDefinition> storeDefs) {
 
         HashMap<StoreDefinition, Integer> uniqueStores = KeyDistributionGenerator.getUniqueStoreDefinitionsWithCounts(storeDefs);
         Set<Integer> nodeIds = cluster.getNodeIds();
+        Set<Integer> zoneIds = cluster.getZoneIds();
 
         builder.append("PARTITION DUMP\n");
         this.primaryAggNodeIdToPartitionCount = Maps.newHashMap();
@@ -86,6 +88,11 @@ public PartitionBalance(Cluster cluster, List<StoreDefinition> storeDefs) {
             allAggNodeIdToPartitionCount.put(nodeId, 0);
         }
 
+        this.zoneIdToPartitionStoreCount = Maps.newHashMap();
+        for(Integer zoneId: zoneIds) {
+            zoneIdToPartitionStoreCount.put(zoneId, 0);
+        }
+
         for(StoreDefinition storeDefinition: uniqueStores.keySet()) {
             StoreRoutingPlan storeRoutingPlan = new StoreRoutingPlan(cluster, storeDefinition);
 
@@ -130,11 +137,25 @@ public PartitionBalance(Cluster cluster, List<StoreDefinition> storeDefs) {
                 allAggNodeIdToPartitionCount.put(nodeId,
                                                  allAggNodeIdToPartitionCount.get(nodeId)
                                                          + (nodeIdToNaryCount.get(nodeId) * uniqueStores.get(storeDefinition)));
+
+                // Count partition-stores per-zone
+                int zoneId = cluster.getNodeById(nodeId).getZoneId();
+                zoneIdToPartitionStoreCount.put(zoneId,
+                                                zoneIdToPartitionStoreCount.get(zoneId)
+                                                        + (nodeIdToNaryCount.get(nodeId) * uniqueStores.get(storeDefinition)));
             }
         }
 
         builder.append(Utils.NEWLINE).append(Utils.NEWLINE);
 
+        builder.append("PARTITION-STORES PER-ZONE:").append(Utils.NEWLINE);
+        for(Integer zoneId: zoneIds) {
+            builder.append("\tZone ID: " + zoneId + " : " + zoneIdToPartitionStoreCount.get(zoneId))
+                   .append(Utils.NEWLINE);
+        }
+
+        builder.append(Utils.NEWLINE).append(Utils.NEWLINE);
+
         Pair<Double, String> summary = summarizeBalance(primaryAggNodeIdToPartitionCount,
                                                         "AGGREGATE PRIMARY-PARTITION COUNT (across all stores)");
         builder.append(summary.getSecond());
diff --git a/src/java/voldemort/tools/RebalancePlanCLI.java b/src/java/voldemort/tools/RebalancePlanCLI.java
index 29d9c27bde..2a78491d4b 100644
--- a/src/java/voldemort/tools/RebalancePlanCLI.java
+++ b/src/java/voldemort/tools/RebalancePlanCLI.java
@@ -91,7 +91,7 @@ private static void printUsage() {
         help.append("    --final-cluster <clusterXML>\n");
         help.append("  Optional:\n");
         help.append("    --final-stores <storesXML> [ Needed for zone expansion ]\n");
-        help.append("    --batch <batch> [ Number of primary partitions to move in each rebalancing batch. ]\n");
+        help.append("    --batch-size <batchSize> [ Number of primary partitions to move in each rebalancing batch. ]\n");
         help.append("    --output-dir <outputDir> [ Directory in which cluster metadata is dumped for each batch of the plan. ]\n");
 
         try {
diff --git a/test/unit/voldemort/tools/PartitionBalanceTest.java b/test/unit/voldemort/tools/PartitionBalanceTest.java
index dd73a21d0f..db86f7166a 100644
--- a/test/unit/voldemort/tools/PartitionBalanceTest.java
+++ b/test/unit/voldemort/tools/PartitionBalanceTest.java
@@ -46,11 +46,17 @@ public class PartitionBalanceTest {
 
     @Test
     public void testBasicThingsThatShouldWork() {
-        new PartitionBalance(ClusterTestUtils.getZZCluster(),
-                             ClusterTestUtils.getZZStoreDefsInMemory());
-
-        new PartitionBalance(ClusterTestUtils.getZZZCluster(),
-                             ClusterTestUtils.getZZZStoreDefsInMemory());
+        PartitionBalance pb = new PartitionBalance(ClusterTestUtils.getZZCluster(),
+                                                   ClusterTestUtils.getZZStoreDefsInMemory());
+        // Print out results so there is a test case that demonstrates toString
+        // method output for 2 zones
+        System.out.println(pb);
+
+        pb = new PartitionBalance(ClusterTestUtils.getZZZCluster(),
+                                  ClusterTestUtils.getZZZStoreDefsInMemory());
+        // Print out results so there is a test case that demonstrates toString
+        // method output for 3 zones
+        System.out.println(pb);
     }
 
     @Test
@@ -133,4 +139,5 @@ public void testNonContiguousZonesThatShouldWorkButDoNot() {
         }
         assertTrue(veCaught);
     }
+
 }