[Fix Bug] Allocation and Rebalance Constraints of WeightFunction are incorrectly reset (opensearch-project#19012)

* fix primary balance strategy unset

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

* Add change log

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

* refactor

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

* refactor

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

* refactor

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

---------

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

Author: guojialiang92

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Omit maxScoreCollector in SimpleTopDocsCollectorContext when concurrent segment search enabled ([#19584](https://github.com/opensearch-project/OpenSearch/pull/19584))
 
 ### Fixed
+- Fix Allocation and Rebalance Constraints of WeightFunction are incorrectly reset ([#19012](https://github.com/opensearch-project/OpenSearch/pull/19012))
 - Fix flaky test FieldDataLoadingIT.testIndicesFieldDataCacheSizeSetting ([#19571](https://github.com/opensearch-project/OpenSearch/pull/19571))
 - Avoid primary shard failure caused by merged segment warmer exceptions ([#19436](https://github.com/opensearch-project/OpenSearch/pull/19436))
 

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationAllocationIT.java

@@ -161,6 +161,14 @@ public void testSingleIndexShardAllocation() throws Exception {
         }
         enablePreferPrimaryBalance();
 
+        // Modify other configurations, expecting that the primary balance strategy will not be affected.
+        assertAcked(
+            client().admin()
+                .cluster()
+                .prepareUpdateSettings()
+                .setPersistentSettings(Settings.builder().put(PRIMARY_SHARD_REBALANCE_BUFFER.getKey(), 0.2))
+        );
+
         ClusterState state;
         createIndex("test", maxShardCount, maxReplicaCount, true);
         ensureGreen(TimeValue.timeValueSeconds(60));

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

@@ -340,7 +340,9 @@ private void updateWeightFunction() {
             this.indexBalanceFactor,
             this.shardBalanceFactor,
             this.preferPrimaryShardRebalanceBuffer,
-            this.primaryConstraintThreshold
+            this.primaryConstraintThreshold,
+            this.preferPrimaryShardBalance,
+            this.preferPrimaryShardRebalance
         );
     }
 
@@ -552,7 +554,14 @@ 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,
+            boolean preferPrimaryShardRebalance
+        ) {
             float sum = indexBalance + shardBalance;
             if (sum <= 0.0f) {
                 throw new IllegalArgumentException("Balance factors must sum to a value > 0 but was: " + sum);
@@ -567,6 +576,10 @@ static class WeightFunction {
             this.rebalanceConstraints = new RebalanceConstraints(rebalanceParameter);
             // Enable index shard per node breach constraint
             updateAllocationConstraint(INDEX_SHARD_PER_NODE_BREACH_CONSTRAINT_ID, true);
+            updateAllocationConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateAllocationConstraint(CLUSTER_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateRebalanceConstraint(INDEX_PRIMARY_SHARD_BALANCE_CONSTRAINT_ID, preferPrimaryShardBalance);
+            updateRebalanceConstraint(CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID, preferPrimaryShardRebalance);
         }
 
         public float weightWithAllocationConstraints(ShardsBalancer balancer, ModelNode node, String index) {

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

@@ -235,7 +235,35 @@ public void testPrimaryBalanceWithPreferPrimaryBalanceSetting() {
         final int numberOfRuns = 5;
         int balanceFailed = 0;
 
-        AllocationService strategy = createAllocationService(getSettingsBuilderForPrimaryBalance().build(), new TestGatewayAllocator());
+        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.PRIMARY_SHARD_REBALANCE_BUFFER.getKey(),
+            BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER.get(settings) + 0.01f
+        ).build();
+        clusterSettings.applySettings(updatedSettings);
+
+        // Double check primary shard balance should still work
         for (int i = 0; i < numberOfRuns; i++) {
             ClusterState clusterState = initCluster(strategy, numberOfIndices, numberOfNodes, numberOfShards, numberOfReplicas);
             clusterState = removeOneNode(clusterState, strategy);

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())