Fixed primary shard balance constraints getting unset

Signed-off-by: Smit Patel <patelsmit32123@gmail.com>
opensearch-project · Jan 22, 2025 · 95a4058 · 95a4058
1 parent 2794655
commit 95a4058
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/...ain/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java b/...ain/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
@@ -309,7 +309,8 @@ private void updateWeightFunction() {
             this.indexBalanceFactor,
             this.shardBalanceFactor,
             this.preferPrimaryShardRebalanceBuffer,
-            this.primaryConstraintThreshold
+            this.primaryConstraintThreshold,
+            this.preferPrimaryShardBalance
         );
     }
 
@@ -319,9 +320,7 @@ private void updateWeightFunction() {
      */
     private void setPreferPrimaryShardBalance(boolean preferPrimaryShardBalance) {
         this.preferPrimaryShardBalance = preferPrimaryShardBalance;
-        this.weightFunction.updateAllocationConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
-        this.weightFunction.updateAllocationConstraint(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
-        this.weightFunction.updateRebalanceConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+        this.weightFunction.updatePrimaryShardBalanceConstraints(preferPrimaryShardBalance);
     }
 
     private void setPreferPrimaryShardRebalance(boolean preferPrimaryShardRebalance) {
@@ -514,7 +513,7 @@ static class WeightFunction {
         private AllocationConstraints constraints;
         private RebalanceConstraints rebalanceConstraints;
 
-        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer, long primaryConstraintThreshold) {
+        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer, long primaryConstraintThreshold, boolean preferPrimaryShardBalance) {
             float sum = indexBalance + shardBalance;
             if (sum <= 0.0f) {
                 throw new IllegalArgumentException("Balance factors must sum to a value > 0 but was: " + sum);
@@ -529,6 +528,8 @@ static class WeightFunction {
             this.rebalanceConstraints = new RebalanceConstraints(rebalanceParameter);
             // Enable index shard per node breach constraint
             updateAllocationConstraint(INDEX_SHARD_PER_NODE_BREACH_CONSTRAINT_ID, true);
+            // Set primary shard balance constraints
+            updatePrimaryShardBalanceConstraints(preferPrimaryShardBalance);
         }
 
         public float weightWithAllocationConstraints(ShardsBalancer balancer, ModelNode node, String index) {
@@ -558,6 +559,12 @@ void updateRebalanceConstraint(String constraint, boolean add) {
         void updatePrimaryConstraintThreshold(long primaryConstraintThreshold) {
             this.primaryConstraintThreshold = primaryConstraintThreshold;
         }
+
+        void updatePrimaryShardBalanceConstraints(boolean preferPrimaryShardBalance) {
+            updateAllocationConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateAllocationConstraint(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateRebalanceConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+        }
     }
 
     /**