@VisibleForTesting
  static SingleSegmentAssignment getNextSingleSegmentAssignment(Map<String, String> currentInstanceStateMap,
      Map<String, String> targetInstanceStateMap, int minAvailableReplicas, boolean lowDiskMode,
      Map<String, Integer> numSegmentsToOffloadMap, Map<Pair<Set<String>, Set<String>>, Set<String>> assignmentMap) {
    Map<String, String> nextInstanceStateMap = new TreeMap<>();

    // Assign the segment the same way as other segments if the current and target instances are the same. We need this
    // to guarantee the mirror servers for replica-group based routing strategies.
    Set<String> currentInstances = currentInstanceStateMap.keySet();
    Set<String> targetInstances = targetInstanceStateMap.keySet();
    Pair<Set<String>, Set<String>> assignmentKey = Pair.of(currentInstances, targetInstances);
    Set<String> instancesToAssign = assignmentMap.get(assignmentKey);
    if (instancesToAssign != null) {
      Set<String> availableInstances = new TreeSet<>();
      for (String instanceName : instancesToAssign) {
        String currentInstanceState = currentInstanceStateMap.get(instanceName);
        String targetInstanceState = targetInstanceStateMap.get(instanceName);
        if (currentInstanceState != null) {
          availableInstances.add(instanceName);
          // Use target instance state if available in case the state changes
          nextInstanceStateMap.put(instanceName,
              targetInstanceState != null ? targetInstanceState : currentInstanceState);
        } else {
          nextInstanceStateMap.put(instanceName, targetInstanceState);
        }
      }
      return new SingleSegmentAssignment(nextInstanceStateMap, availableInstances);
    }

    // Add all the common instances
    // Use target instance state in case the state changes
    for (Map.Entry<String, String> entry : targetInstanceStateMap.entrySet()) {
      String instanceName = entry.getKey();
      if (currentInstanceStateMap.containsKey(instanceName)) {
        nextInstanceStateMap.put(instanceName, entry.getValue());
      }
    }

    // Add current instances until the min available replicas achieved
    int numInstancesToKeep = minAvailableReplicas - nextInstanceStateMap.size();
    if (numInstancesToKeep > 0) {
      // Sort instances by number of segments to offload, and keep the ones with the least segments to offload
      List<Triple<String, String, Integer>> instancesInfo =
          getSortedInstancesOnNumSegmentsToOffload(currentInstanceStateMap, nextInstanceStateMap,
              numSegmentsToOffloadMap);
      numInstancesToKeep = Integer.min(numInstancesToKeep, instancesInfo.size());
      for (int i = 0; i < numInstancesToKeep; i++) {
        Triple<String, String, Integer> instanceInfo = instancesInfo.get(i);
        nextInstanceStateMap.put(instanceInfo.getLeft(), instanceInfo.getMiddle());
      }
    }
    Set<String> availableInstances = new TreeSet<>(nextInstanceStateMap.keySet());

    // After achieving the min available replicas, when low disk mode is enabled, only add new instances when all
    // current instances exist in the next assignment.
    // We want to first drop the extra instances as one step, then add the target instances as another step to avoid the
    // case where segments are first added to the instance before other segments are dropped from the instance, which
    // might cause server running out of disk. Note that even if segment addition and drop happen in the same step,
    // there is no guarantee that server process the segment drop before the segment addition.
    if (!lowDiskMode || currentInstanceStateMap.size() == nextInstanceStateMap.size()) {
      int numInstancesToAdd = targetInstanceStateMap.size() - nextInstanceStateMap.size();
      if (numInstancesToAdd > 0) {
        // Sort instances by number of segments to offload, and add the ones with the least segments to offload
        List<Triple<String, String, Integer>> instancesInfo =
            getSortedInstancesOnNumSegmentsToOffload(targetInstanceStateMap, nextInstanceStateMap,
                numSegmentsToOffloadMap);
        for (int i = 0; i < numInstancesToAdd; i++) {
          Triple<String, String, Integer> instanceInfo = instancesInfo.get(i);
          nextInstanceStateMap.put(instanceInfo.getLeft(), instanceInfo.getMiddle());
        }
      }
    }

    assignmentMap.put(assignmentKey, nextInstanceStateMap.keySet());
    return new SingleSegmentAssignment(nextInstanceStateMap, availableInstances);
  }
@Test
  public void testTwoMinAvailableReplicasWithLowDiskMode() {
    // With 3 common instances, first assignment should keep the common instances and remove the not common instance
    Map<String, String> currentInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host3", "host4"), ONLINE);
    Map<String, String> targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host3", "host5"), ONLINE);
    TableRebalancer.SingleSegmentAssignment assignment =
        getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host3"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2", "host3")));
    // Second assignment should be the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2", "host3")));

    // With 2 common instances, first assignment should keep the common instances and remove the not common instances
    targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host5", "host6"), ONLINE);
    assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Second assignment should be the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));

    // With 1 common instance, fist assignment should keep the common instance, and 1 more instance from current
    // assignment
    targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host5", "host6", "host7"), ONLINE);
    assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Second assignment should add 2 instances from target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host5", "host6"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Third assignment should remove the old instance from current assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host5", "host6"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host5", "host6")));
    // Fourth assignment should make the assignment the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host5", "host6")));

    // Without common instance, fist assignment should keep 2 instances from current assignment
    targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host5", "host6", "host7", "host8"), ONLINE);
    assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Second assignment should add 2 instances from target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host5", "host6"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Third assignment should remove the old instances from current assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host5", "host6"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host5", "host6")));
    // Fourth assignment should make the assignment the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host5", "host6")));

    // With increasing number of replicas, fist assignment should keep 2 instances from current assignment
    targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host5", "host6", "host7", "host8", "host9"), ONLINE);
    assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Second assignment should add 3 instances from target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host5", "host6", "host7"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Third assignment should remove the old instances from current assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host5", "host6", "host7"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host5", "host6", "host7")));
    // Fourth assignment should make the assignment the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host5", "host6", "host7")));

    // With decreasing number of replicas, fist assignment should keep 2 instances from current assignment
    targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host5", "host6", "host7"), ONLINE);
    assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Second assignment should add 1 instance from target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host5"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host2")));
    // Third assignment should remove 1 old instance from current assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host5"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host5")));
    // Forth assignment should add 1 more instance from target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host5", "host6"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host1", "host5")));
    // Fifth assignment should remove the other old instance from current assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host5", "host6"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host5", "host6")));
    // Sixth assignment should make the assignment the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host5", "host6")));

    // With increasing from 1 replica, fist assignment should keep the instance from current assignment, and add 2
    // instances from target assignment
    // NOTE: This is the best we can do because we don't have 2 replicas available
    currentInstanceStateMap = SegmentAssignmentUtils.getInstanceStateMap(Collections.singletonList("host1"), ONLINE);
    targetInstanceStateMap =
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host2", "host3", "host4"), ONLINE);
    assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host3"), ONLINE));
    assertEquals(assignment._availableInstances, Collections.singleton("host1"));
    // Second assignment should remove the old instance from current assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap,
        SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host2", "host3"), ONLINE));
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host2", "host3")));
    // Third assignment should make the assignment the same as target assignment
    assignment = getNextSingleSegmentAssignment(assignment._instanceStateMap, targetInstanceStateMap, 2, true);
    assertEquals(assignment._instanceStateMap, targetInstanceStateMap);
    assertEquals(assignment._availableInstances, new TreeSet<>(Arrays.asList("host2", "host3")));
  }