@Test public void testRebalanceTableWithHelixAutoRebalanceStrategy() { int numSegments = 100; List<String> segments = SegmentAssignmentTestUtils.getNameList(SEGMENT_NAME_PREFIX, numSegments); int numInstances = 10; List<String> instances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, numInstances); Map<String, Map<String, String>> currentAssignment = new TreeMap<>(); int assignedInstanceId = 0; for (String segmentName : segments) { List<String> instancesAssigned = new ArrayList<>(NUM_REPLICAS); for (int replicaId = 0; replicaId < NUM_REPLICAS; replicaId++) { instancesAssigned.add(instances.get(assignedInstanceId)); assignedInstanceId = (assignedInstanceId + 1) % numInstances; } currentAssignment .put(segmentName, SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned, SegmentStateModel.ONLINE)); } assertEquals(currentAssignment.size(), numSegments); for (Map<String, String> instanceStateMap : currentAssignment.values()) { assertEquals(instanceStateMap.size(), NUM_REPLICAS); } int[] numSegmentsAssignedPerInstance = SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(currentAssignment, instances); int[] expectedNumSegmentsAssignedPerInstance = new int[numInstances]; int numSegmentsPerInstance = numSegments * NUM_REPLICAS / numInstances; Arrays.fill(expectedNumSegmentsAssignedPerInstance, numSegmentsPerInstance); assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance); assertEquals( SegmentAssignmentUtils.rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, instances, NUM_REPLICAS), currentAssignment); List<String> newInstances = new ArrayList<>(instances); String newInstanceName = INSTANCE_NAME_PREFIX + 10; newInstances.set(0, newInstanceName); Map<String, Map<String, String>> newAssignment = SegmentAssignmentUtils .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS); assertEquals(currentAssignment.size(), numSegments); for (Map<String, String> instanceStateMap : newAssignment.values()) { assertEquals(instanceStateMap.size(), NUM_REPLICAS); } numSegmentsAssignedPerInstance = SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances); assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance); Map<String, Integer> numSegmentsToBeMovedPerInstance = SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment); assertEquals(numSegmentsToBeMovedPerInstance.size(), 1); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstanceName), numSegmentsPerInstance); String oldInstanceName = INSTANCE_NAME_PREFIX + 0; for (String segmentName : segments) { if (currentAssignment.get(segmentName).containsKey(oldInstanceName)) { assertTrue(newAssignment.get(segmentName).containsKey(newInstanceName)); } } int newNumInstances = numInstances - 5; newInstances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, newNumInstances); newAssignment = SegmentAssignmentUtils .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS); assertEquals(newAssignment.size(), numSegments); for (Map<String, String> instanceStateMap : newAssignment.values()) { assertEquals(instanceStateMap.size(), NUM_REPLICAS); } numSegmentsAssignedPerInstance = SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances); assertEquals(numSegmentsAssignedPerInstance[0], 56); assertEquals(numSegmentsAssignedPerInstance[1], 60); assertEquals(numSegmentsAssignedPerInstance[2], 60); assertEquals(numSegmentsAssignedPerInstance[3], 60); assertEquals(numSegmentsAssignedPerInstance[4], 64); numSegmentsToBeMovedPerInstance = SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment); assertEquals(numSegmentsToBeMovedPerInstance.size(), newNumInstances); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(0)), 26); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(1)), 30); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(2)), 30); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(3)), 30); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(4)), 34); newNumInstances = numInstances + 5; newInstances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, newNumInstances); newAssignment = SegmentAssignmentUtils .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS); assertEquals(newAssignment.size(), numSegments); for (Map<String, String> instanceStateMap : newAssignment.values()) { assertEquals(instanceStateMap.size(), NUM_REPLICAS); } numSegmentsAssignedPerInstance = SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances); expectedNumSegmentsAssignedPerInstance = new int[newNumInstances]; int newNumSegmentsPerInstance = numSegments * NUM_REPLICAS / newNumInstances; Arrays.fill(expectedNumSegmentsAssignedPerInstance, newNumSegmentsPerInstance); assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance); numSegmentsToBeMovedPerInstance = SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment); assertEquals(numSegmentsToBeMovedPerInstance.size(), 5); for (int instanceId = numInstances; instanceId < newNumInstances; instanceId++) { assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(instanceId)), newNumSegmentsPerInstance); } String newInstanceNamePrefix = "i_"; newInstances = SegmentAssignmentTestUtils.getNameList(newInstanceNamePrefix, numInstances); newAssignment = SegmentAssignmentUtils .rebalanceTableWithHelixAutoRebalanceStrategy(currentAssignment, newInstances, NUM_REPLICAS); assertEquals(newAssignment.size(), numSegments); for (Map<String, String> instanceStateMap : newAssignment.values()) { assertEquals(instanceStateMap.size(), NUM_REPLICAS); } numSegmentsAssignedPerInstance = SegmentAssignmentUtils.getNumSegmentsAssignedPerInstance(newAssignment, newInstances); expectedNumSegmentsAssignedPerInstance = new int[numInstances]; Arrays.fill(expectedNumSegmentsAssignedPerInstance, numSegmentsPerInstance); assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance); numSegmentsToBeMovedPerInstance = SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment); assertEquals(numSegmentsToBeMovedPerInstance.size(), numInstances); for (String instanceName : newInstances) { assertEquals((int) numSegmentsToBeMovedPerInstance.get(instanceName), numSegmentsPerInstance); } }
static Map<String, Map<String, String>> rebalanceTableWithHelixAutoRebalanceStrategy( Map<String, Map<String, String>> currentAssignment, List<String> instances, int replication) { LinkedHashMap<String, Integer> states = new LinkedHashMap<>(); states.put(SegmentStateModel.ONLINE, replication); AutoRebalanceStrategy autoRebalanceStrategy = new AutoRebalanceStrategy(null, new ArrayList<>(currentAssignment.keySet()), states); Map<String, Map<String, String>> currentAssignmentCopy = new TreeMap<>(); for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) { String segmentName = entry.getKey(); Map<String, String> instanceStateMap = entry.getValue(); currentAssignmentCopy.put(segmentName, new TreeMap<>(instanceStateMap)); } return autoRebalanceStrategy.computePartitionAssignment(instances, instances, currentAssignmentCopy, null) .getMapFields(); }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceTableWithHelixAutoRebalanceStrategy( Map<String, Map<String, String>> currentAssignment, List<String> instances, int replication) { LinkedHashMap<String, Integer> states = new LinkedHashMap<>(); states.put(SegmentStateModel.ONLINE, replication); AutoRebalanceStrategy autoRebalanceStrategy = new AutoRebalanceStrategy(null, new ArrayList<>(currentAssignment.keySet()), states); Map<String, Map<String, String>> currentAssignmentCopy = new TreeMap<>(); for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) { String segmentName = entry.getKey(); Map<String, String> instanceStateMap = entry.getValue(); currentAssignmentCopy.put(segmentName, new TreeMap<>(instanceStateMap)); } return autoRebalanceStrategy.computePartitionAssignment(instances, instances, currentAssignmentCopy, null) .getMapFields(); } }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceTableWithHelixAutoRebalanceStrategy( Map<String, Map<String, String>> currentAssignment, List<String> instances, int replication) { LinkedHashMap<String, Integer> states = new LinkedHashMap<>(); states.put(SegmentStateModel.ONLINE, replication); AutoRebalanceStrategy autoRebalanceStrategy = new AutoRebalanceStrategy(null, new ArrayList<>(currentAssignment.keySet()), states); Map<String, Map<String, String>> currentAssignmentCopy = new TreeMap<>(); for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) { String segmentName = entry.getKey(); Map<String, String> instanceStateMap = entry.getValue(); currentAssignmentCopy.put(segmentName, new TreeMap<>(instanceStateMap)); } return autoRebalanceStrategy.computePartitionAssignment(instances, instances, currentAssignmentCopy, null) .getMapFields(); } private  SegmentAssignmentUtils(); }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceTableWithHelixAutoRebalanceStrategy( Map<String, Map<String, String>> currentAssignment, List<String> instances, int replication) { LinkedHashMap<String, Integer> states = new LinkedHashMap<>(); states.put(SegmentStateModel.ONLINE, replication); AutoRebalanceStrategy autoRebalanceStrategy = new AutoRebalanceStrategy(null, new ArrayList<>(currentAssignment.keySet()), states); Map<String, Map<String, String>> currentAssignmentCopy = new TreeMap<>(); for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) { String segmentName = entry.getKey(); Map<String, String> instanceStateMap = entry.getValue(); currentAssignmentCopy.put(segmentName, new TreeMap<>(instanceStateMap)); } return autoRebalanceStrategy.computePartitionAssignment(instances, instances, currentAssignmentCopy, null) .getMapFields(); } private  SegmentAssignmentUtils(); static Map<String, String> getInstanceStateMap(Collection<String> instances, String state); static Map<String, Integer> getNumSegmentsToBeMovedPerInstance(Map<String, Map<String, String>> oldAssignment,
      Map<String, Map<String, String>> newAssignment); }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceTableWithHelixAutoRebalanceStrategy( Map<String, Map<String, String>> currentAssignment, List<String> instances, int replication) { LinkedHashMap<String, Integer> states = new LinkedHashMap<>(); states.put(SegmentStateModel.ONLINE, replication); AutoRebalanceStrategy autoRebalanceStrategy = new AutoRebalanceStrategy(null, new ArrayList<>(currentAssignment.keySet()), states); Map<String, Map<String, String>> currentAssignmentCopy = new TreeMap<>(); for (Map.Entry<String, Map<String, String>> entry : currentAssignment.entrySet()) { String segmentName = entry.getKey(); Map<String, String> instanceStateMap = entry.getValue(); currentAssignmentCopy.put(segmentName, new TreeMap<>(instanceStateMap)); } return autoRebalanceStrategy.computePartitionAssignment(instances, instances, currentAssignmentCopy, null) .getMapFields(); } private  SegmentAssignmentUtils(); static Map<String, String> getInstanceStateMap(Collection<String> instances, String state); static Map<String, Integer> getNumSegmentsToBeMovedPerInstance(Map<String, Map<String, String>> oldAssignment,
      Map<String, Map<String, String>> newAssignment);  }