Make writable partition selection account for partitions with more than 3 replicas

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,6 @@ 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

@@ -341,6 +341,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;
@@ -349,10 +350,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);
     }
 
@@ -364,6 +368,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) {
@@ -432,7 +437,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);
@@ -487,20 +491,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

@@ -232,7 +232,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));
   }
 
   /**