Make writable partition selection account for partitions with more th…

…an 3 replicas (#1349)

Currently, selecting writable partition for PUT operation only picks
partitions with highest local replica count. This will become a problem
when "move replica" is rolled out. In the intermediate state of "move
replica", particular partition may have 6 replicas and if we stick with
current logic. All the write traffic will be routed to this particular
partition. Hence, this PR we make changes to allow partition selection
to pick partitions with replicas count >= min replica count(specified in
ClusterMapConfig). In most cases, min replica count = 3.

ambry-api/src/main/java/com.github.ambry/config/ClusterMapConfig.java

@@ -216,6 +216,14 @@ public class ClusterMapConfig {
   @Config("clustermap.replica.catchup.target")
   public final int clustermapReplicaCatchupTarget;
 
+  /**
+   * The minimum number of replicas in local datacenter required for a partition to serve PUT request. This is used to
+   * get writable partitions for PUT operation. Any partition with replica count larger than or equal to this number is
+   * acceptable to be considered as a candidate.
+   */
+  @Config("clustermap.writable.partition.min.replica.count")
+  public final int clustermapWritablePartitionMinReplicaCount;
+
   public ClusterMapConfig(VerifiableProperties verifiableProperties) {
     clusterMapFixedTimeoutDatanodeErrorThreshold =
         verifiableProperties.getIntInRange("clustermap.fixedtimeout.datanode.error.threshold", 3, 1, 100);
@@ -261,5 +269,7 @@ public ClusterMapConfig(VerifiableProperties verifiableProperties) {
         verifiableProperties.getLongInRange("clustermap.replica.catchup.acceptable.lag.bytes", 0L, 0L, Long.MAX_VALUE);
     clustermapReplicaCatchupTarget =
         verifiableProperties.getIntInRange("clustermap.replica.catchup.target", 0, 0, Integer.MAX_VALUE);
+    clustermapWritablePartitionMinReplicaCount =
+        verifiableProperties.getIntInRange("clustermap.writable.partition.min.replica.count", 3, 0, Integer.MAX_VALUE);
   }
 }

ambry-clustermap/src/main/java/com.github.ambry.clustermap/ClusterMapUtils.java

@@ -364,6 +364,7 @@ static boolean areAllReplicasForPartitionUp(PartitionId partition) {
    * Not thread safe.
    */
   static class PartitionSelectionHelper {
+    private final int minimumLocalReplicaCount;
     private Collection<? extends PartitionId> allPartitions;
     private Map<String, SortedMap<Integer, List<PartitionId>>> partitionIdsByClassAndLocalReplicaCount;
     private Map<PartitionId, List<ReplicaId>> partitionIdToLocalReplicas;
@@ -372,10 +373,13 @@ static class PartitionSelectionHelper {
     /**
      * @param allPartitions the list of all {@link PartitionId}s
      * @param localDatacenterName the name of the local datacenter. Can be null if datacenter specific replica counts
-     *                            are not required.
+     * @param minimumLocalReplicaCount the minimum number of replicas in local datacenter. This is used when selecting
+     *                                 writable partitions.
      */
-    PartitionSelectionHelper(Collection<? extends PartitionId> allPartitions, String localDatacenterName) {
+    PartitionSelectionHelper(Collection<? extends PartitionId> allPartitions, String localDatacenterName,
+        int minimumLocalReplicaCount) {
       this.localDatacenterName = localDatacenterName;
+      this.minimumLocalReplicaCount = minimumLocalReplicaCount;
       updatePartitions(allPartitions, localDatacenterName);
     }
 
@@ -387,6 +391,7 @@ static class PartitionSelectionHelper {
      */
     void updatePartitions(Collection<? extends PartitionId> allPartitions, String localDatacenterName) {
       this.allPartitions = allPartitions;
+      // todo when new partitions added into clustermap, dynamically update these two maps.
       partitionIdsByClassAndLocalReplicaCount = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
       partitionIdToLocalReplicas = new HashMap<>();
       for (PartitionId partition : allPartitions) {
@@ -455,7 +460,6 @@ List<PartitionId> getWritablePartitions(String partitionClass) {
     PartitionId getRandomWritablePartition(String partitionClass, List<PartitionId> partitionsToExclude) {
       PartitionId anyWritablePartition = null;
       List<PartitionId> partitionsInClass = getPartitionsInClass(partitionClass, true);
-
       int workingSize = partitionsInClass.size();
       while (workingSize > 0) {
         int randomIndex = ThreadLocalRandom.current().nextInt(workingSize);
@@ -510,20 +514,23 @@ private boolean areAllLocalReplicasForPartitionUp(PartitionId partitionId) {
      * Returns the partitions belonging to the {@code partitionClass}. Returns all partitions if {@code partitionClass}
      * is {@code null}.
      * @param partitionClass the class of the partitions desired.
-     * @param highestReplicaCountOnly if {@code true}, returns only the partitions with the highest number of replicas
-     *                                in the local datacenter.
+     * @param minimumReplicaCountRequired if {@code true}, returns only the partitions with the number of replicas in
+     *                                    local datacenter that is larger than or equal to minimum required count.
      * @return the partitions belonging to the {@code partitionClass}. Returns all partitions if {@code partitionClass}
      * is {@code null}.
      */
-    private List<PartitionId> getPartitionsInClass(String partitionClass, boolean highestReplicaCountOnly) {
+    private List<PartitionId> getPartitionsInClass(String partitionClass, boolean minimumReplicaCountRequired) {
       List<PartitionId> toReturn = new ArrayList<>();
       if (partitionClass == null) {
         toReturn.addAll(allPartitions);
       } else if (partitionIdsByClassAndLocalReplicaCount.containsKey(partitionClass)) {
         SortedMap<Integer, List<PartitionId>> partitionsByReplicaCount =
             partitionIdsByClassAndLocalReplicaCount.get(partitionClass);
-        if (highestReplicaCountOnly) {
-          toReturn.addAll(partitionsByReplicaCount.get(partitionsByReplicaCount.lastKey()));
+        if (minimumReplicaCountRequired) {
+          // get partitions with replica count >= min replica count specified in ClusterMapConfig
+          for (List<PartitionId> partitionIds : partitionsByReplicaCount.tailMap(minimumLocalReplicaCount).values()) {
+            toReturn.addAll(partitionIds);
+          }
         } else {
           for (List<PartitionId> partitionIds : partitionIdsByClassAndLocalReplicaCount.get(partitionClass).values()) {
             toReturn.addAll(partitionIds);

ambry-clustermap/src/main/java/com.github.ambry.clustermap/HelixBootstrapUpgradeUtil.java

@@ -329,7 +329,7 @@ private HelixBootstrapUpgradeUtil(String hardwareLayoutPath, String partitionLay
     } else {
       staticClusterMap = (new StaticClusterAgentsFactory(clusterMapConfig, new PartitionLayout(
           new HardwareLayout(new JSONObject(Utils.readStringFromFile(hardwareLayoutPath)), clusterMapConfig),
-          null))).getClusterMap();
+          clusterMapConfig))).getClusterMap();
     }
     String clusterNameInStaticClusterMap = staticClusterMap.partitionLayout.getClusterName();
     clusterName = clusterNamePrefix + clusterNameInStaticClusterMap;

ambry-clustermap/src/main/java/com.github.ambry.clustermap/HelixClusterManager.java

@@ -218,7 +218,8 @@ public HelixClusterManager(ClusterMapConfig clusterMapConfig, String instanceNam
     }
     localDatacenterId = dcToDcZkInfo.get(clusterMapConfig.clusterMapDatacenterName).dcZkInfo.getDcId();
     partitionSelectionHelper =
-        new PartitionSelectionHelper(partitionMap.values(), clusterMapConfig.clusterMapDatacenterName);
+        new PartitionSelectionHelper(partitionMap.values(), clusterMapConfig.clusterMapDatacenterName,
+            clusterMapConfig.clustermapWritablePartitionMinReplicaCount);
   }
 
   /**

ambry-clustermap/src/main/java/com.github.ambry.clustermap/PartitionLayout.java

@@ -13,6 +13,7 @@
  */
 package com.github.ambry.clustermap;
 
+import com.github.ambry.config.ClusterMapConfig;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
@@ -52,16 +53,16 @@ class PartitionLayout {
    * Create a PartitionLayout
    * @param hardwareLayout the {@link HardwareLayout} to use.
    * @param jsonObject the {@link JSONObject} that represents the partition layout
-   * @param localDatacenterName the name of the local datacenter. Can be {@code null}.
+   * @param clusterMapConfig the {@link ClusterMapConfig} to use.
    * @throws JSONException
    */
-  public PartitionLayout(HardwareLayout hardwareLayout, JSONObject jsonObject, String localDatacenterName)
+  public PartitionLayout(HardwareLayout hardwareLayout, JSONObject jsonObject, ClusterMapConfig clusterMapConfig)
       throws JSONException {
     if (logger.isTraceEnabled()) {
       logger.trace("PartitionLayout " + hardwareLayout + ", " + jsonObject.toString());
     }
     this.hardwareLayout = hardwareLayout;
-    this.localDatacenterName = localDatacenterName;
+    this.localDatacenterName = clusterMapConfig.clusterMapDatacenterName;
     this.clusterName = jsonObject.getString("clusterName");
     this.version = jsonObject.getLong("version");
     this.partitionMap = new HashMap<>();
@@ -70,26 +71,28 @@ public PartitionLayout(HardwareLayout hardwareLayout, JSONObject jsonObject, Str
     }
 
     validate();
-    partitionSelectionHelper = new ClusterMapUtils.PartitionSelectionHelper(partitionMap.values(), localDatacenterName);
+    partitionSelectionHelper = new ClusterMapUtils.PartitionSelectionHelper(partitionMap.values(), localDatacenterName,
+        clusterMapConfig.clustermapWritablePartitionMinReplicaCount);
   }
 
   /**
    * Constructor for initial PartitionLayout
    * @param hardwareLayout the {@link JSONObject} that represents the partition layout
-   * @param localDatacenterName the name of the local datacenter. Can be {@code null}.
+   * @param clusterMapConfig the {@link ClusterMapConfig} to use.
    */
-  public PartitionLayout(HardwareLayout hardwareLayout, String localDatacenterName) {
+  public PartitionLayout(HardwareLayout hardwareLayout, ClusterMapConfig clusterMapConfig) {
     if (logger.isTraceEnabled()) {
       logger.trace("PartitionLayout " + hardwareLayout);
     }
     this.hardwareLayout = hardwareLayout;
-    this.localDatacenterName = localDatacenterName;
+    this.localDatacenterName = clusterMapConfig.clusterMapDatacenterName;
     this.clusterName = hardwareLayout.getClusterName();
     this.version = 1;
     this.maxPartitionId = MinPartitionId;
     this.partitionMap = new HashMap<>();
     validate();
-    partitionSelectionHelper = new ClusterMapUtils.PartitionSelectionHelper(partitionMap.values(), localDatacenterName);
+    partitionSelectionHelper = new ClusterMapUtils.PartitionSelectionHelper(partitionMap.values(), localDatacenterName,
+        clusterMapConfig.clustermapWritablePartitionMinReplicaCount);
   }
 
   public HardwareLayout getHardwareLayout() {

ambry-clustermap/src/main/java/com.github.ambry.clustermap/StaticClusterAgentsFactory.java

@@ -52,7 +52,7 @@ public StaticClusterAgentsFactory(ClusterMapConfig clusterMapConfig, String hard
       String partitionLayoutFilePath) throws JSONException, IOException {
     this(clusterMapConfig, new PartitionLayout(
         new HardwareLayout(new JSONObject(readStringFromFile(hardwareLayoutFilePath)), clusterMapConfig),
-        new JSONObject(readStringFromFile(partitionLayoutFilePath)), clusterMapConfig.clusterMapDatacenterName));
+        new JSONObject(readStringFromFile(partitionLayoutFilePath)), clusterMapConfig));
   }
 
   /**