@Test public void testRebalanceReplicaGroupBasedTable() { int numSegments = 90; List<String> segments = SegmentAssignmentTestUtils.getNameList(SEGMENT_NAME_PREFIX, numSegments); Map<Integer, List<String>> partitionIdToSegmentsMap = Collections.singletonMap(0, segments); int numInstances = 9; List<String> instances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, numInstances); InstancePartitions instancePartitions = new InstancePartitions(null); int numInstancesPerReplicaGroup = numInstances / NUM_REPLICAS; int instanceIdToAdd = 0; for (int replicaGroupId = 0; replicaGroupId < NUM_REPLICAS; replicaGroupId++) { List<String> instancesForReplicaGroup = new ArrayList<>(numInstancesPerReplicaGroup); for (int i = 0; i < numInstancesPerReplicaGroup; i++) { instancesForReplicaGroup.add(instances.get(instanceIdToAdd++)); } instancePartitions.setInstances(0, replicaGroupId, instancesForReplicaGroup); } Map<String, Map<String, String>> currentAssignment = new TreeMap<>(); for (int segmentId = 0; segmentId < numSegments; segmentId++) { List<String> instancesAssigned = new ArrayList<>(NUM_REPLICAS); for (int replicaGroupId = 0; replicaGroupId < NUM_REPLICAS; replicaGroupId++) { int assignedInstanceId = segmentId % numInstancesPerReplicaGroup + replicaGroupId * numInstancesPerReplicaGroup; instancesAssigned.add(instances.get(assignedInstanceId)); } currentAssignment.put(segments.get(segmentId), 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 .rebalanceReplicaGroupBasedTable(currentAssignment, instancePartitions, partitionIdToSegmentsMap), currentAssignment); List<String> newInstances = new ArrayList<>(numInstances); List<String> newReplicaGroup0Instances = new ArrayList<>(instancePartitions.getInstances(0, 0)); String newReplicaGroup0Instance = INSTANCE_NAME_PREFIX + 9; newReplicaGroup0Instances.set(0, newReplicaGroup0Instance); newInstances.addAll(newReplicaGroup0Instances); List<String> newReplicaGroup1Instances = new ArrayList<>(instancePartitions.getInstances(0, 1)); String newReplicaGroup1Instance = INSTANCE_NAME_PREFIX + 10; newReplicaGroup1Instances.set(1, newReplicaGroup1Instance); newInstances.addAll(newReplicaGroup1Instances); List<String> newReplicaGroup2Instances = instancePartitions.getInstances(0, 2); newInstances.addAll(newReplicaGroup2Instances); InstancePartitions newInstancePartitions = new InstancePartitions(null); newInstancePartitions.setInstances(0, 0, newReplicaGroup0Instances); newInstancePartitions.setInstances(0, 1, newReplicaGroup1Instances); newInstancePartitions.setInstances(0, 2, newReplicaGroup2Instances); Map<String, Map<String, String>> newAssignment = SegmentAssignmentUtils .rebalanceReplicaGroupBasedTable(currentAssignment, newInstancePartitions, partitionIdToSegmentsMap); assertEquals(newAssignment.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(), 2); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newReplicaGroup0Instance), numSegmentsPerInstance); assertEquals((int) numSegmentsToBeMovedPerInstance.get(newReplicaGroup1Instance), numSegmentsPerInstance); String oldReplicaGroup0Instance = INSTANCE_NAME_PREFIX + 0; String oldReplicaGroup1Instance = INSTANCE_NAME_PREFIX + 4; for (String segmentName : segments) { Map<String, String> oldInstanceStateMap = currentAssignment.get(segmentName); if (oldInstanceStateMap.containsKey(oldReplicaGroup0Instance)) { assertTrue(newAssignment.get(segmentName).containsKey(newReplicaGroup0Instance)); } if (oldInstanceStateMap.containsKey(oldReplicaGroup1Instance)) { assertTrue(newAssignment.get(segmentName).containsKey(newReplicaGroup1Instance)); } } int newNumInstances = numInstances - 3; int newNumInstancesPerReplicaGroup = newNumInstances / NUM_REPLICAS; newInstances = new ArrayList<>(newNumInstances); for (int replicaGroupId = 0; replicaGroupId < NUM_REPLICAS; replicaGroupId++) { List<String> newInstancesForReplicaGroup = instancePartitions.getInstances(0, replicaGroupId).subList(0, newNumInstancesPerReplicaGroup); newInstancePartitions.setInstances(0, replicaGroupId, newInstancesForReplicaGroup); newInstances.addAll(newInstancesForReplicaGroup); } newAssignment = SegmentAssignmentUtils .rebalanceReplicaGroupBasedTable(currentAssignment, newInstancePartitions, partitionIdToSegmentsMap); 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(), newNumInstances); for (String instanceName : newInstances) { assertEquals((int) numSegmentsToBeMovedPerInstance.get(instanceName), newNumSegmentsPerInstance - numSegmentsPerInstance); } newNumInstances = numInstances + 6; newNumInstancesPerReplicaGroup = newNumInstances / NUM_REPLICAS; newInstances = SegmentAssignmentTestUtils.getNameList(INSTANCE_NAME_PREFIX, newNumInstances); instanceIdToAdd = numInstances; for (int replicaGroupId = 0; replicaGroupId < NUM_REPLICAS; replicaGroupId++) { List<String> newInstancesForReplicaGroup = new ArrayList<>(instancePartitions.getInstances(0, replicaGroupId)); for (int i = 0; i < newNumInstancesPerReplicaGroup - numInstancesPerReplicaGroup; i++) { newInstancesForReplicaGroup.add(newInstances.get(instanceIdToAdd++)); } newInstancePartitions.setInstances(0, replicaGroupId, newInstancesForReplicaGroup); } newAssignment = SegmentAssignmentUtils .rebalanceReplicaGroupBasedTable(currentAssignment, newInstancePartitions, partitionIdToSegmentsMap); 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]; newNumSegmentsPerInstance = numSegments * NUM_REPLICAS / newNumInstances; Arrays.fill(expectedNumSegmentsAssignedPerInstance, newNumSegmentsPerInstance); assertEquals(numSegmentsAssignedPerInstance, expectedNumSegmentsAssignedPerInstance); numSegmentsToBeMovedPerInstance = SegmentAssignmentUtils.getNumSegmentsToBeMovedPerInstance(currentAssignment, newAssignment); assertEquals(numSegmentsToBeMovedPerInstance.size(), 6); for (int instanceId = numInstances; instanceId < newNumInstances; instanceId++) { assertEquals((int) numSegmentsToBeMovedPerInstance.get(newInstances.get(instanceId)), newNumSegmentsPerInstance); } newInstances = SegmentAssignmentTestUtils.getNameList("i_", numInstances); instanceIdToAdd = 0; for (int replicaGroupId = 0; replicaGroupId < NUM_REPLICAS; replicaGroupId++) { List<String> instancesForReplicaGroup = new ArrayList<>(numInstancesPerReplicaGroup); for (int i = 0; i < numInstancesPerReplicaGroup; i++) { instancesForReplicaGroup.add(newInstances.get(instanceIdToAdd++)); } newInstancePartitions.setInstances(0, replicaGroupId, instancesForReplicaGroup); } newAssignment = SegmentAssignmentUtils .rebalanceReplicaGroupBasedTable(currentAssignment, newInstancePartitions, partitionIdToSegmentsMap); 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>> rebalanceReplicaGroupBasedTable( Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, Map<Integer, List<String>> partitionIdToSegmentsMap) { Map<String, Map<String, String>> newAssignment = new TreeMap<>(); int numPartitions = instancePartitions.getNumPartitions(); for (Map.Entry<Integer, List<String>> entry : partitionIdToSegmentsMap.entrySet()) { int partitionId = entry.getKey() % numPartitions; SegmentAssignmentUtils .rebalanceReplicaGroupBasedPartition(currentAssignment, instancePartitions, partitionId, entry.getValue(), newAssignment); } return newAssignment; }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceReplicaGroupBasedTable( Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, Map<Integer, List<String>> partitionIdToSegmentsMap) { Map<String, Map<String, String>> newAssignment = new TreeMap<>(); int numPartitions = instancePartitions.getNumPartitions(); for (Map.Entry<Integer, List<String>> entry : partitionIdToSegmentsMap.entrySet()) { int partitionId = entry.getKey() % numPartitions; SegmentAssignmentUtils .rebalanceReplicaGroupBasedPartition(currentAssignment, instancePartitions, partitionId, entry.getValue(), newAssignment); } return newAssignment; } }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceReplicaGroupBasedTable( Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, Map<Integer, List<String>> partitionIdToSegmentsMap) { Map<String, Map<String, String>> newAssignment = new TreeMap<>(); int numPartitions = instancePartitions.getNumPartitions(); for (Map.Entry<Integer, List<String>> entry : partitionIdToSegmentsMap.entrySet()) { int partitionId = entry.getKey() % numPartitions; SegmentAssignmentUtils .rebalanceReplicaGroupBasedPartition(currentAssignment, instancePartitions, partitionId, entry.getValue(), newAssignment); } return newAssignment; } private  SegmentAssignmentUtils(); }
SegmentAssignmentUtils { static Map<String, Map<String, String>> rebalanceReplicaGroupBasedTable( Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, Map<Integer, List<String>> partitionIdToSegmentsMap) { Map<String, Map<String, String>> newAssignment = new TreeMap<>(); int numPartitions = instancePartitions.getNumPartitions(); for (Map.Entry<Integer, List<String>> entry : partitionIdToSegmentsMap.entrySet()) { int partitionId = entry.getKey() % numPartitions; SegmentAssignmentUtils .rebalanceReplicaGroupBasedPartition(currentAssignment, instancePartitions, partitionId, entry.getValue(), newAssignment); } return newAssignment; } 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>> rebalanceReplicaGroupBasedTable( Map<String, Map<String, String>> currentAssignment, InstancePartitions instancePartitions, Map<Integer, List<String>> partitionIdToSegmentsMap) { Map<String, Map<String, String>> newAssignment = new TreeMap<>(); int numPartitions = instancePartitions.getNumPartitions(); for (Map.Entry<Integer, List<String>> entry : partitionIdToSegmentsMap.entrySet()) { int partitionId = entry.getKey() % numPartitions; SegmentAssignmentUtils .rebalanceReplicaGroupBasedPartition(currentAssignment, instancePartitions, partitionId, entry.getValue(), newAssignment); } return newAssignment; } 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);  }