Fix failing test in ShardMovementStrategyTests (#9420)

Signed-off-by: Poojita Raj <[email protected]>

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

@@ -169,12 +169,25 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting
         clusterSettings.addSettingsUpdateConsumer(THRESHOLD_SETTING, this::setThreshold);
     }
 
+    /**
+     * Changes in deprecated setting SHARD_MOVE_PRIMARY_FIRST_SETTING affect value of its replacement setting SHARD_MOVEMENT_STRATEGY_SETTING.
+     */
     private void setMovePrimaryFirst(boolean movePrimaryFirst) {
         this.movePrimaryFirst = movePrimaryFirst;
+        setShardMovementStrategy(this.shardMovementStrategy);
     }
 
+    /**
+     * Sets the correct Shard movement strategy to use.
+     * If users are still using deprecated setting `move_primary_first`, we want behavior to remain unchanged.
+     * In the event of changing ShardMovementStrategy setting from default setting NO_PREFERENCE to either PRIMARY_FIRST or REPLICA_FIRST, we want that
+     * to have priority over values set in move_primary_first setting.
+     */
     private void setShardMovementStrategy(ShardMovementStrategy shardMovementStrategy) {
         this.shardMovementStrategy = shardMovementStrategy;
+        if (shardMovementStrategy == ShardMovementStrategy.NO_PREFERENCE && this.movePrimaryFirst) {
+            this.shardMovementStrategy = ShardMovementStrategy.PRIMARY_FIRST;
+        }
     }
 
     private void setWeightFunction(float indexBalance, float shardBalanceFactor) {
@@ -205,7 +218,6 @@ public void allocate(RoutingAllocation allocation) {
         final ShardsBalancer localShardsBalancer = new LocalShardsBalancer(
             logger,
             allocation,
-            movePrimaryFirst,
             shardMovementStrategy,
             weightFunction,
             threshold,
@@ -227,7 +239,6 @@ public ShardAllocationDecision decideShardAllocation(final ShardRouting shard, f
         ShardsBalancer localShardsBalancer = new LocalShardsBalancer(
             logger,
             allocation,
-            movePrimaryFirst,
             shardMovementStrategy,
             weightFunction,
             threshold,
@@ -479,13 +490,12 @@ public static class Balancer extends LocalShardsBalancer {
         public Balancer(
             Logger logger,
             RoutingAllocation allocation,
-            boolean movePrimaryFirst,
             ShardMovementStrategy shardMovementStrategy,
             BalancedShardsAllocator.WeightFunction weight,
             float threshold,
             boolean preferPrimaryBalance
         ) {
-            super(logger, allocation, movePrimaryFirst, shardMovementStrategy, weight, threshold, preferPrimaryBalance);
+            super(logger, allocation, shardMovementStrategy, weight, threshold, preferPrimaryBalance);
         }
     }
 

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

@@ -58,7 +58,6 @@ public class LocalShardsBalancer extends ShardsBalancer {
     private final Map<String, BalancedShardsAllocator.ModelNode> nodes;
     private final RoutingAllocation allocation;
     private final RoutingNodes routingNodes;
-    private final boolean movePrimaryFirst;
     private final ShardMovementStrategy shardMovementStrategy;
 
     private final boolean preferPrimaryBalance;
@@ -75,15 +74,13 @@ public class LocalShardsBalancer extends ShardsBalancer {
     public LocalShardsBalancer(
         Logger logger,
         RoutingAllocation allocation,
-        boolean movePrimaryFirst,
         ShardMovementStrategy shardMovementStrategy,
         BalancedShardsAllocator.WeightFunction weight,
         float threshold,
         boolean preferPrimaryBalance
     ) {
         this.logger = logger;
         this.allocation = allocation;
-        this.movePrimaryFirst = movePrimaryFirst;
         this.weight = weight;
         this.threshold = threshold;
         this.routingNodes = allocation.routingNodes();
@@ -531,22 +528,6 @@ private void checkAndAddInEligibleTargetNode(RoutingNode targetNode) {
         }
     }
 
-    /**
-     * Returns the correct Shard movement strategy to use.
-     * If users are still using deprecated setting "move_primary_first", we want behavior to remain unchanged.
-     * In the event of changing ShardMovementStrategy setting from default setting NO_PREFERENCE to either PRIMARY_FIRST or REPLICA_FIRST, we want that
-     * to have priority over values set in move_primary_first setting.
-     */
-    private ShardMovementStrategy getShardMovementStrategy() {
-        if (shardMovementStrategy != ShardMovementStrategy.NO_PREFERENCE) {
-            return shardMovementStrategy;
-        }
-        if (movePrimaryFirst) {
-            return ShardMovementStrategy.PRIMARY_FIRST;
-        }
-        return ShardMovementStrategy.NO_PREFERENCE;
-    }
-
     /**
      * Move started shards that can not be allocated to a node anymore
      *
@@ -569,8 +550,7 @@ void moveShards() {
             checkAndAddInEligibleTargetNode(currentNode.getRoutingNode());
         }
         boolean primariesThrottled = false;
-        for (Iterator<ShardRouting> it = allocation.routingNodes().nodeInterleavedShardIterator(getShardMovementStrategy()); it
-            .hasNext();) {
+        for (Iterator<ShardRouting> it = allocation.routingNodes().nodeInterleavedShardIterator(shardMovementStrategy); it.hasNext();) {
             // Verify if the cluster concurrent recoveries have been reached.
             if (allocation.deciders().canMoveAnyShard(allocation).type() != Decision.Type.YES) {
                 logger.info(

server/src/test/java/org/opensearch/cluster/routing/ShardMovementStrategyTests.java

@@ -55,24 +55,23 @@ private static Settings.Builder getSettings(ShardMovementStrategy shardMovementS
             .put("cluster.routing.allocation.move.primary_first", movePrimaryFirst);
     }
 
-    public void testClusterGreenAfterPartialRelocationPrimaryFirstShardMovementMovePrimarySettingEnabled() throws InterruptedException {
+    public void testClusterRelocationPrimaryFirstShardMovementMovePrimarySettingEnabled() throws InterruptedException {
         testClusterGreenAfterPartialRelocation(ShardMovementStrategy.PRIMARY_FIRST, true);
     }
 
-    public void testClusterGreenAfterPartialRelocationPrimaryFirstShardMovementMovePrimarySettingDisabled() throws InterruptedException {
+    public void testClusterRelocationPrimaryFirstShardMovementMovePrimarySettingDisabled() throws InterruptedException {
         testClusterGreenAfterPartialRelocation(ShardMovementStrategy.PRIMARY_FIRST, false);
     }
 
-    public void testClusterGreenAfterPartialRelocationReplicaFirstShardMovementPrimaryFirstEnabled() throws InterruptedException {
+    public void testClusterRelocationReplicaFirstShardMovementPrimaryFirstEnabled() throws InterruptedException {
         testClusterGreenAfterPartialRelocation(ShardMovementStrategy.REPLICA_FIRST, true);
     }
 
-    public void testClusterGreenAfterPartialRelocationReplicaFirstShardMovementPrimaryFirstDisabled() throws InterruptedException {
+    public void testClusterRelocationReplicaFirstShardMovementPrimaryFirstDisabled() throws InterruptedException {
         testClusterGreenAfterPartialRelocation(ShardMovementStrategy.REPLICA_FIRST, false);
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/9178")
-    public void testClusterGreenAfterPartialRelocationNoPreferenceShardMovementPrimaryFirstEnabled() throws InterruptedException {
+    public void testClusterRelocationNoPreferenceShardMovementPrimaryFirstEnabled() throws InterruptedException {
         testClusterGreenAfterPartialRelocation(ShardMovementStrategy.NO_PREFERENCE, true);
     }