Using gt to make constraint test during rebalance

Signed-off-by: panguixin <[email protected]>
opensearch-project · Feb 12, 2025 · d4e101a · d4e101a
1 parent d0a65d3
commit d4e101a
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 9 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationConstraints.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationConstraints.java
@@ -31,8 +31,8 @@ public class AllocationConstraints {
     public AllocationConstraints() {
         this.constraints = new HashMap<>();
         this.constraints.put(INDEX_SHARD_PER_NODE_BREACH_CONSTRAINT_ID, new Constraint(isIndexShardsPerNodeBreached()));
-        this.constraints.put(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, new Constraint(isPerIndexPrimaryShardsPerNodeBreached()));
-        this.constraints.put(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, new Constraint(isPrimaryShardsPerNodeBreached(0.0f)));
+        this.constraints.put(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, new Constraint(isPerIndexPrimaryShardsPerNodeBreached(false)));
+        this.constraints.put(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, new Constraint(isPrimaryShardsPerNodeBreached(false, 0.0f)));
     }
 
     public void updateAllocationConstraint(String constraint, boolean enable) {

diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/ConstraintTypes.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/ConstraintTypes.java
@@ -66,11 +66,11 @@ public static Predicate<Constraint.ConstraintParams> isIndexShardsPerNodeBreache
      * {@link ConstraintTypes#CONSTRAINT_WEIGHT} is assigned to node resulting in lesser chances of node being selected
      * as allocation or rebalancing target
      */
-    public static Predicate<Constraint.ConstraintParams> isPerIndexPrimaryShardsPerNodeBreached() {
+    public static Predicate<Constraint.ConstraintParams> isPerIndexPrimaryShardsPerNodeBreached(boolean rebalance) {
         return (params) -> {
             int perIndexPrimaryShardCount = params.getNode().numPrimaryShards(params.getIndex());
             int perIndexAllowedPrimaryShardCount = (int) Math.ceil(params.getBalancer().avgPrimaryShardsPerNode(params.getIndex()));
-            return perIndexPrimaryShardCount >= perIndexAllowedPrimaryShardCount;
+            return rebalance ? perIndexPrimaryShardCount > perIndexAllowedPrimaryShardCount : perIndexPrimaryShardCount >= perIndexAllowedPrimaryShardCount;
         };
     }
 
@@ -79,11 +79,11 @@ public static Predicate<Constraint.ConstraintParams> isPerIndexPrimaryShardsPerN
      * constraint is used in weight calculation during allocation/rebalance both. When breached a high weight {@link ConstraintTypes#CONSTRAINT_WEIGHT}
      * is assigned to node resulting in lesser chances of node being selected as allocation/rebalance target
      */
-    public static Predicate<Constraint.ConstraintParams> isPrimaryShardsPerNodeBreached(float buffer) {
+    public static Predicate<Constraint.ConstraintParams> isPrimaryShardsPerNodeBreached(boolean rebalance, float buffer) {
         return (params) -> {
             int primaryShardCount = params.getNode().numPrimaryShards();
             int allowedPrimaryShardCount = (int) Math.ceil(params.getBalancer().avgPrimaryShardsPerNode() * (1 + buffer));
-            return primaryShardCount >= allowedPrimaryShardCount;
+            return rebalance ? primaryShardCount > allowedPrimaryShardCount : primaryShardCount >= allowedPrimaryShardCount;
         };
     }
 

diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/RebalanceConstraints.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/RebalanceConstraints.java
@@ -31,10 +31,10 @@ public class RebalanceConstraints {
 
     public RebalanceConstraints(RebalanceParameter rebalanceParameter) {
         this.constraints = new HashMap<>();
-        this.constraints.put(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, new Constraint(isPerIndexPrimaryShardsPerNodeBreached()));
+        this.constraints.put(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, new Constraint(isPerIndexPrimaryShardsPerNodeBreached(true)));
         this.constraints.put(
             CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID,
-            new Constraint(isPrimaryShardsPerNodeBreached(rebalanceParameter.getPreferPrimaryBalanceBuffer()))
+            new Constraint(isPrimaryShardsPerNodeBreached(true, rebalanceParameter.getPreferPrimaryBalanceBuffer()))
         );
     }
 

diff --git a/...ain/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java b/...ain/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
@@ -148,7 +148,7 @@ public class BalancedShardsAllocator implements ShardsAllocator {
     );
 
     /**
-     * This setting governs whether primary shards balance is desired during allocation. This is used by {@link ConstraintTypes#isPerIndexPrimaryShardsPerNodeBreached()}
+     * This setting governs whether primary shards balance is desired during allocation. This is used by {@link ConstraintTypes#isPerIndexPrimaryShardsPerNodeBreached(boolean)}
      * and {@link ConstraintTypes#isPrimaryShardsPerNodeBreached} which is used during unassigned shard allocation
      * {@link LocalShardsBalancer#allocateUnassigned()} and shard re-balance/relocation to a different node via {@link LocalShardsBalancer#balance()} .
      */