Add Setting to adjust the primary constraint weights

Signed-off-by: Arpit Bandejiya <[email protected]>

Author: alchemist51

server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationConstraints.java

@@ -39,8 +39,8 @@ public void updateAllocationConstraint(String constraint, boolean enable) {
         this.constraints.get(constraint).setEnable(enable);
     }
 
-    public long weight(ShardsBalancer balancer, BalancedShardsAllocator.ModelNode node, String index) {
-        Constraint.ConstraintParams params = new Constraint.ConstraintParams(balancer, node, index);
+    public long weight(ShardsBalancer balancer, BalancedShardsAllocator.ModelNode node, String index, long primaryThresholdWeight) {
+        Constraint.ConstraintParams params = new Constraint.ConstraintParams(balancer, node, index, primaryThresholdWeight);
         return params.weight(constraints);
     }
 }

server/src/main/java/org/opensearch/cluster/routing/allocation/Constraint.java

@@ -14,7 +14,7 @@
 import java.util.Map;
 import java.util.function.Predicate;
 
-import static org.opensearch.cluster.routing.allocation.ConstraintTypes.CONSTRAINT_WEIGHT;
+import static org.opensearch.cluster.routing.allocation.ConstraintTypes.predicateKeyToWeightMap;
 
 /**
  * Defines a constraint useful to de-prioritize certain nodes as target of unassigned shards used in {@link AllocationConstraints} or
@@ -44,11 +44,13 @@ static class ConstraintParams {
         private ShardsBalancer balancer;
         private BalancedShardsAllocator.ModelNode node;
         private String index;
+        private long PrimaryConstraintThreshold;
 
-        ConstraintParams(ShardsBalancer balancer, BalancedShardsAllocator.ModelNode node, String index) {
+        ConstraintParams(ShardsBalancer balancer, BalancedShardsAllocator.ModelNode node, String index, long primaryConstraintThreshold) {
             this.balancer = balancer;
             this.node = node;
             this.index = index;
+            this.PrimaryConstraintThreshold = primaryConstraintThreshold;
         }
 
         public ShardsBalancer getBalancer() {
@@ -75,9 +77,12 @@ public String getIndex() {
          */
         public long weight(Map<String, Constraint> constraints) {
             long totalConstraintWeight = 0;
-            for (Constraint constraint : constraints.values()) {
+            for (Map.Entry<String, Constraint> entry : constraints.entrySet()) {
+                String key = entry.getKey();
+                Constraint constraint = entry.getValue();
                 if (constraint.test(this)) {
-                    totalConstraintWeight += CONSTRAINT_WEIGHT;
+                    double weight = predicateKeyToWeightMap(key, PrimaryConstraintThreshold);
+                    totalConstraintWeight += weight;
                 }
             }
             return totalConstraintWeight;

server/src/main/java/org/opensearch/cluster/routing/allocation/ConstraintTypes.java

@@ -86,4 +86,14 @@ public static Predicate<Constraint.ConstraintParams> isPrimaryShardsPerNodeBreac
             return primaryShardCount >= allowedPrimaryShardCount;
         };
     }
+
+    public static long predicateKeyToWeightMap(String key, long primaryConstraintWeight) {
+        switch (key) {
+            case CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID:
+            case CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID:
+                return primaryConstraintWeight;
+            default:
+                return CONSTRAINT_WEIGHT;
+        }
+    }
 }

server/src/main/java/org/opensearch/cluster/routing/allocation/RebalanceConstraints.java

@@ -42,8 +42,8 @@ public void updateRebalanceConstraint(String constraint, boolean enable) {
         this.constraints.get(constraint).setEnable(enable);
     }
 
-    public long weight(ShardsBalancer balancer, BalancedShardsAllocator.ModelNode node, String index) {
-        Constraint.ConstraintParams params = new Constraint.ConstraintParams(balancer, node, index);
+    public long weight(ShardsBalancer balancer, BalancedShardsAllocator.ModelNode node, String index, long primaryConstraintThreshold) {
+        Constraint.ConstraintParams params = new Constraint.ConstraintParams(balancer, node, index, primaryConstraintThreshold);
         return params.weight(constraints);
     }
 }

server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java

@@ -139,6 +139,14 @@ public class BalancedShardsAllocator implements ShardsAllocator {
         Property.NodeScope
     );
 
+    public static final Setting<Long> PRIMARY_CONSTRAINT_THRESHOLD_SETTING = Setting.longSetting(
+        "cluster.routing.allocation.primary_constraint.threshold",
+        10,
+        0,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     /**
      * This setting governs whether primary shards balance is desired during allocation. This is used by {@link ConstraintTypes#isPerIndexPrimaryShardsPerNodeBreached()}
      * and {@link ConstraintTypes#isPrimaryShardsPerNodeBreached} which is used during unassigned shard allocation
@@ -201,6 +209,7 @@ public class BalancedShardsAllocator implements ShardsAllocator {
     private volatile float shardBalanceFactor;
     private volatile WeightFunction weightFunction;
     private volatile float threshold;
+    private volatile long primaryConstraintThreshold;
 
     private volatile boolean ignoreThrottleInRestore;
     private volatile TimeValue allocatorTimeout;
@@ -219,6 +228,7 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting
         setIgnoreThrottleInRestore(IGNORE_THROTTLE_FOR_REMOTE_RESTORE.get(settings));
         updateWeightFunction();
         setThreshold(THRESHOLD_SETTING.get(settings));
+        setPrimaryConstraintThresholdSetting(PRIMARY_CONSTRAINT_THRESHOLD_SETTING.get(settings));
         setPreferPrimaryShardBalance(PREFER_PRIMARY_SHARD_BALANCE.get(settings));
         setPreferPrimaryShardRebalance(PREFER_PRIMARY_SHARD_REBALANCE.get(settings));
         setShardMovementStrategy(SHARD_MOVEMENT_STRATEGY_SETTING.get(settings));
@@ -231,6 +241,7 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting
         clusterSettings.addSettingsUpdateConsumer(PRIMARY_SHARD_REBALANCE_BUFFER, this::updatePreferPrimaryShardBalanceBuffer);
         clusterSettings.addSettingsUpdateConsumer(PREFER_PRIMARY_SHARD_REBALANCE, this::setPreferPrimaryShardRebalance);
         clusterSettings.addSettingsUpdateConsumer(THRESHOLD_SETTING, this::setThreshold);
+        clusterSettings.addSettingsUpdateConsumer(PRIMARY_CONSTRAINT_THRESHOLD_SETTING, this::setPrimaryConstraintThresholdSetting);
         clusterSettings.addSettingsUpdateConsumer(IGNORE_THROTTLE_FOR_REMOTE_RESTORE, this::setIgnoreThrottleInRestore);
         clusterSettings.addSettingsUpdateConsumer(ALLOCATOR_TIMEOUT_SETTING, this::setAllocatorTimeout);
     }
@@ -294,7 +305,12 @@ private void updatePreferPrimaryShardBalanceBuffer(float preferPrimaryShardBalan
     }
 
     private void updateWeightFunction() {
-        weightFunction = new WeightFunction(this.indexBalanceFactor, this.shardBalanceFactor, this.preferPrimaryShardRebalanceBuffer);
+        weightFunction = new WeightFunction(
+            this.indexBalanceFactor,
+            this.shardBalanceFactor,
+            this.preferPrimaryShardRebalanceBuffer,
+            this.primaryConstraintThreshold
+        );
     }
 
     /**
@@ -317,6 +333,11 @@ private void setThreshold(float threshold) {
         this.threshold = threshold;
     }
 
+    private void setPrimaryConstraintThresholdSetting(long threshold) {
+        this.primaryConstraintThreshold = threshold;
+        this.weightFunction.updatePrimaryConstraintThreshold(threshold);
+    }
+
     private void setAllocatorTimeout(TimeValue allocatorTimeout) {
         this.allocatorTimeout = allocatorTimeout;
     }
@@ -489,10 +510,11 @@ static class WeightFunction {
         private final float shardBalance;
         private final float theta0;
         private final float theta1;
+        private long primaryConstraintThreshold;
         private AllocationConstraints constraints;
         private RebalanceConstraints rebalanceConstraints;
 
-        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer) {
+        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer, long primaryConstraintThreshold) {
             float sum = indexBalance + shardBalance;
             if (sum <= 0.0f) {
                 throw new IllegalArgumentException("Balance factors must sum to a value > 0 but was: " + sum);
@@ -501,6 +523,7 @@ static class WeightFunction {
             theta1 = indexBalance / sum;
             this.indexBalance = indexBalance;
             this.shardBalance = shardBalance;
+            this.primaryConstraintThreshold = primaryConstraintThreshold;
             RebalanceParameter rebalanceParameter = new RebalanceParameter(preferPrimaryBalanceBuffer);
             this.constraints = new AllocationConstraints();
             this.rebalanceConstraints = new RebalanceConstraints(rebalanceParameter);
@@ -510,12 +533,12 @@ static class WeightFunction {
 
         public float weightWithAllocationConstraints(ShardsBalancer balancer, ModelNode node, String index) {
             float balancerWeight = weight(balancer, node, index);
-            return balancerWeight + constraints.weight(balancer, node, index);
+            return balancerWeight + constraints.weight(balancer, node, index, primaryConstraintThreshold);
         }
 
         public float weightWithRebalanceConstraints(ShardsBalancer balancer, ModelNode node, String index) {
             float balancerWeight = weight(balancer, node, index);
-            return balancerWeight + rebalanceConstraints.weight(balancer, node, index);
+            return balancerWeight + rebalanceConstraints.weight(balancer, node, index, primaryConstraintThreshold);
         }
 
         float weight(ShardsBalancer balancer, ModelNode node, String index) {
@@ -531,6 +554,10 @@ void updateAllocationConstraint(String constraint, boolean enable) {
         void updateRebalanceConstraint(String constraint, boolean add) {
             this.rebalanceConstraints.updateRebalanceConstraint(constraint, add);
         }
+
+        void updatePrimaryConstraintThreshold(long primaryConstraintThreshold) {
+            this.primaryConstraintThreshold = primaryConstraintThreshold;
+        }
     }
 
     /**