Added tests

Signed-off-by: Smit Patel <patelsmit32123@gmail.com>
opensearch-project · Jan 22, 2025 · feacabc · feacabc
1 parent 95a4058
commit feacabc
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 1 deletion.
diff --git a/...ain/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java b/...ain/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
@@ -513,7 +513,13 @@ static class WeightFunction {
         private AllocationConstraints constraints;
         private RebalanceConstraints rebalanceConstraints;
 
-        WeightFunction(float indexBalance, float shardBalance, float preferPrimaryBalanceBuffer, long primaryConstraintThreshold, boolean preferPrimaryShardBalance) {
+        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);

diff --git a/...er/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java b/...er/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.
      */

diff --git a/test/framework/src/main/java/org/opensearch/cluster/OpenSearchAllocationTestCase.java b/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())