Fixed primary shard balance constraints getting unset

server/src/main/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,13 @@ 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 +534,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 +565,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);
+        }
     }
 
     /**

server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java

@@ -250,6 +250,68 @@ public void testPrimaryBalanceWithPreferPrimaryBalanceSetting() {
         assertTrue(balanceFailed <= 1);
     }
 
+    /**
+     * This test verifies primary shard balance is not reset when primary shard balance factors are changed.
+     */
+    public void testPreferPrimaryBalanceSettingRemainsUnchangedWithUpdatingPrimaryBalanceFactors() {
+        final int numberOfNodes = 5;
+        final int numberOfIndices = 5;
+        final int numberOfShards = 25;
+        final int numberOfReplicas = 1;
+        final int numberOfRuns = 5;
+        int balanceFailed = 0;
+
+        Settings settings = getSettingsBuilderForPrimaryBalance().build();
+        ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+        AllocationService strategy = createAllocationService(
+            settings,
+            clusterSettings,
+            new TestGatewayAllocator(),
+            SNAPSHOT_INFO_SERVICE_WITH_NO_SHARD_SIZES
+        );
+        for (int i = 0; i < numberOfRuns; i++) {
+            ClusterState clusterState = initCluster(strategy, numberOfIndices, numberOfNodes, numberOfShards, numberOfReplicas);
+            clusterState = removeOneNode(clusterState, strategy);
+            logger.info(ShardAllocations.printShardDistribution(clusterState));
+            try {
+                verifyPerIndexPrimaryBalance(clusterState);
+            } catch (AssertionError e) {
+                balanceFailed++;
+                logger.info("Unexpected assertion failure");
+            }
+        }
+        assertTrue(balanceFailed <= 1);
+
+        // Update settings & apply
+        Settings updatedSettings = getSettingsBuilderForPrimaryBalance().put(
+            BalancedShardsAllocator.INDEX_BALANCE_FACTOR_SETTING.getKey(),
+            BalancedShardsAllocator.INDEX_BALANCE_FACTOR_SETTING.get(settings) + 0.01f
+        )
+            .put(
+                BalancedShardsAllocator.SHARD_BALANCE_FACTOR_SETTING.getKey(),
+                BalancedShardsAllocator.SHARD_BALANCE_FACTOR_SETTING.get(settings) + 0.01f
+            )
+            .put(
+                BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER.getKey(),
+                BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER.get(settings) + 0.01f
+            )
+            .build();
+        clusterSettings.applySettings(updatedSettings);
+        // Test if primary shard balancing still works
+        for (int i = 0; i < numberOfRuns; i++) {
+            ClusterState clusterState = initCluster(strategy, numberOfIndices, numberOfNodes, numberOfShards, numberOfReplicas);
+            clusterState = removeOneNode(clusterState, strategy);
+            logger.info(ShardAllocations.printShardDistribution(clusterState));
+            try {
+                verifyPerIndexPrimaryBalance(clusterState);
+            } catch (AssertionError e) {
+                balanceFailed++;
+                logger.info("Unexpected assertion failure");
+            }
+        }
+        assertTrue(balanceFailed <= 1);
+    }
+
     /**
      * This test verifies primary shard balance is attained  setting.
      */

test/framework/src/main/java/org/opensearch/cluster/OpenSearchAllocationTestCase.java

@@ -142,6 +142,21 @@ public static MockAllocationService createAllocationService(
         );
     }
 
+    public static MockAllocationService createAllocationService(
+        Settings settings,
+        ClusterSettings clusterSettings,
+        GatewayAllocator gatewayAllocator,
+        SnapshotsInfoService snapshotsInfoService
+    ) {
+        return new MockAllocationService(
+            randomAllocationDeciders(settings, EMPTY_CLUSTER_SETTINGS, random()),
+            gatewayAllocator,
+            new BalancedShardsAllocator(settings, clusterSettings),
+            EmptyClusterInfoService.INSTANCE,
+            snapshotsInfoService
+        );
+    }
+
     public static AllocationDeciders randomAllocationDeciders(Settings settings, ClusterSettings clusterSettings, Random random) {
         List<AllocationDecider> deciders = new ArrayList<>(
             ClusterModule.createAllocationDeciders(settings, clusterSettings, Collections.emptyList())