Make shard balance aware of allocation filters

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

@@ -190,6 +190,13 @@ public class BalancedShardsAllocator implements ShardsAllocator {
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> PREFER_FILTER_AWARE_BALANCE = Setting.boolSetting(
+        "cluster.routing.allocation.balance.filter_aware",
+        false,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     public static final Setting<TimeValue> ALLOCATOR_TIMEOUT_SETTING = Setting.timeSetting(
         "cluster.routing.allocation.balanced_shards_allocator.allocator_timeout",
         TimeValue.MINUS_ONE,
@@ -238,6 +245,7 @@ private static Priority parseReroutePriority(String priorityString) {
 
     private volatile boolean preferPrimaryShardBalance;
     private volatile boolean preferPrimaryShardRebalance;
+    private volatile boolean preferFilterAwareBalance;
     private volatile float preferPrimaryShardRebalanceBuffer;
     private volatile float indexBalanceFactor;
     private volatile float shardBalanceFactor;
@@ -266,6 +274,7 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting
         setPrimaryConstraintThresholdSetting(PRIMARY_CONSTRAINT_THRESHOLD_SETTING.get(settings));
         setPreferPrimaryShardBalance(PREFER_PRIMARY_SHARD_BALANCE.get(settings));
         setPreferPrimaryShardRebalance(PREFER_PRIMARY_SHARD_REBALANCE.get(settings));
+        setPreferFilterAwareBalance(PREFER_FILTER_AWARE_BALANCE.get(settings));
         setShardMovementStrategy(SHARD_MOVEMENT_STRATEGY_SETTING.get(settings));
         setAllocatorTimeout(ALLOCATOR_TIMEOUT_SETTING.get(settings));
         setFollowUpRerouteTaskPriority(FOLLOW_UP_REROUTE_PRIORITY_SETTING.get(settings));
@@ -276,6 +285,7 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting
         clusterSettings.addSettingsUpdateConsumer(SHARD_BALANCE_FACTOR_SETTING, this::updateShardBalanceFactor);
         clusterSettings.addSettingsUpdateConsumer(PRIMARY_SHARD_REBALANCE_BUFFER, this::updatePreferPrimaryShardBalanceBuffer);
         clusterSettings.addSettingsUpdateConsumer(PREFER_PRIMARY_SHARD_REBALANCE, this::setPreferPrimaryShardRebalance);
+        clusterSettings.addSettingsUpdateConsumer(PREFER_FILTER_AWARE_BALANCE, this::setPreferFilterAwareBalance);
         clusterSettings.addSettingsUpdateConsumer(THRESHOLD_SETTING, this::setThreshold);
         clusterSettings.addSettingsUpdateConsumer(PRIMARY_CONSTRAINT_THRESHOLD_SETTING, this::setPrimaryConstraintThresholdSetting);
         clusterSettings.addSettingsUpdateConsumer(IGNORE_THROTTLE_FOR_REMOTE_RESTORE, this::setIgnoreThrottleInRestore);
@@ -368,6 +378,10 @@ private void setPreferPrimaryShardRebalance(boolean preferPrimaryShardRebalance)
         this.weightFunction.updateRebalanceConstraint(CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID, preferPrimaryShardRebalance);
     }
 
+    private void setPreferFilterAwareBalance(boolean preferFilterAwareBalance) {
+        this.preferFilterAwareBalance = preferFilterAwareBalance;
+    }
+
     private void setThreshold(float threshold) {
         this.threshold = threshold;
     }
@@ -409,6 +423,7 @@ public void allocate(RoutingAllocation allocation) {
             threshold,
             preferPrimaryShardBalance,
             preferPrimaryShardRebalance,
+            preferFilterAwareBalance,
             ignoreThrottleInRestore,
             this::allocatorTimedOut
         );
@@ -435,6 +450,7 @@ public ShardAllocationDecision decideShardAllocation(final ShardRouting shard, f
             threshold,
             preferPrimaryShardBalance,
             preferPrimaryShardRebalance,
+            preferFilterAwareBalance,
             ignoreThrottleInRestore,
             () -> false // as we don't need to check if timed out or not while just understanding ShardAllocationDecision
         );
@@ -795,7 +811,7 @@ public Balancer(
             float threshold,
             boolean preferPrimaryBalance
         ) {
-            super(logger, allocation, shardMovementStrategy, weight, threshold, preferPrimaryBalance, false, false, () -> false);
+            super(logger, allocation, shardMovementStrategy, weight, threshold, preferPrimaryBalance, false, false, false, () -> false);
         }
     }
 

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

@@ -72,6 +72,7 @@ public class LocalShardsBalancer extends ShardsBalancer {
     private final Metadata metadata;
 
     private final float avgPrimaryShardsPerNode;
+    private final int balanceNodeCount;
     private final BalancedShardsAllocator.NodeSorter sorter;
     private final Set<RoutingNode> inEligibleTargetNode;
     private final Supplier<Boolean> timedOutFunc;
@@ -85,6 +86,7 @@ public LocalShardsBalancer(
         float threshold,
         boolean preferPrimaryBalance,
         boolean preferPrimaryRebalance,
+        boolean preferFilterAwareBalance,
         boolean ignoreThrottleInRestore,
         Supplier<Boolean> timedOutFunc
     ) {
@@ -94,9 +96,21 @@ public LocalShardsBalancer(
         this.threshold = threshold;
         this.routingNodes = allocation.routingNodes();
         this.metadata = allocation.metadata();
-        avgPrimaryShardsPerNode = (float) (StreamSupport.stream(metadata.spliterator(), false)
-            .mapToInt(IndexMetadata::getNumberOfShards)
-            .sum()) / routingNodes.size();
+        int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum();
+        int balanceNodeCount = routingNodes.size();
+        if (preferFilterAwareBalance) {
+            int eligibleNodeCount = 0;
+            for (RoutingNode routingNode : routingNodes) {
+                if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) {
+                    eligibleNodeCount++;
+                }
+            }
+            if (eligibleNodeCount > 0) {
+                balanceNodeCount = eligibleNodeCount;
+            }
+        }
+        this.balanceNodeCount = balanceNodeCount;
+        avgPrimaryShardsPerNode = ((float) primarySum) / balanceNodeCount;
         nodes = Collections.unmodifiableMap(buildModelFromAssigned());
         sorter = newNodeSorter();
         inEligibleTargetNode = new HashSet<>();
@@ -119,12 +133,12 @@ private BalancedShardsAllocator.ModelNode[] nodesArray() {
      */
     @Override
     public float avgShardsPerNode(String index) {
-        return ((float) metadata.index(index).getTotalNumberOfShards()) / nodes.size();
+        return ((float) metadata.index(index).getTotalNumberOfShards()) / balanceNodeCount;
     }
 
     @Override
     public float avgPrimaryShardsPerNode(String index) {
-        return ((float) metadata.index(index).getNumberOfShards()) / nodes.size();
+        return ((float) metadata.index(index).getNumberOfShards()) / balanceNodeCount;
     }
 
     @Override
@@ -137,7 +151,7 @@ public float avgPrimaryShardsPerNode() {
      */
     @Override
     public float avgShardsPerNode() {
-        return totalShardCount / nodes.size();
+        return ((float) totalShardCount) / balanceNodeCount;
     }
 
     /**

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -286,6 +286,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER,
                 BalancedShardsAllocator.PREFER_PRIMARY_SHARD_BALANCE,
                 BalancedShardsAllocator.PREFER_PRIMARY_SHARD_REBALANCE,
+                BalancedShardsAllocator.PREFER_FILTER_AWARE_BALANCE,
                 BalancedShardsAllocator.SHARD_MOVE_PRIMARY_FIRST_SETTING,
                 BalancedShardsAllocator.SHARD_MOVEMENT_STRATEGY_SETTING,
                 BalancedShardsAllocator.THRESHOLD_SETTING,

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

@@ -1112,4 +1112,182 @@ public ShardAllocationDecision decideShardAllocation(ShardRouting shard, Routing
             }
         }
     }
+
+    /**
+     * Cluster: 8 nodes; node4..node7 excluded via cluster.routing.allocation.exclude._id
+     * Index  : 12 primaries x (1 + 1 replica) = 24 shards, all on the 4 eligible nodes
+     *          every eligible node holds exactly 6 shards.
+     *   Shard layout (each column = one node, each row = one shard on that node):
+     *
+     *   eligible (allowed)         | excluded (filter NO)
+     *   node0  node1  node2  node3 | node4  node5  node6  node7
+     *   -----  -----  -----  ----- | -----  -----  -----  -----
+     *     P      P      P      P   |  .      .      .      .
+     *     P      P      P      P   |  .      .      .      .
+     *     P      r      r      r   |  .      .      .      .
+     *     P      r      r      r   |  .      .      .      .
+     *     P      r      r      r   |  .      .      .      .
+     *     P      r      r      r   |  .      .      .      .
+     *   6P/0r  2P/4r  2P/4r  2P/4r |          (P = primary, r = replica)
+     *
+     *   Primary layout on eligible nodes: 6 / 2 / 2 / 2   (ideal = 3 / 3 / 3 / 3)
+     */
+    public void testPrimaryRebalanceIgnoresAllocationFilter() {
+        final int numberOfNodes = 8;
+        final int excludedNodeCount = 4;
+        final int numberOfShards = 12;
+        final int numberOfReplicas = 1;
+        final String indexName = "test";
+
+        // Enable prefer-primary balance & rebalance, exclude the second half of nodes.
+        Settings.Builder settingsBuilder = getSettingsBuilderForPrimaryReBalance();
+        StringBuilder excludeList = new StringBuilder();
+        for (int i = numberOfNodes - excludedNodeCount; i < numberOfNodes; i++) {
+            if (excludeList.length() > 0) {
+                excludeList.append(",");
+            }
+            excludeList.append("node").append(i);
+        }
+        settingsBuilder.put("cluster.routing.allocation.exclude._id", excludeList.toString());
+        settingsBuilder.put("cluster.routing.allocation.balance.filter_aware", true);
+
+        AllocationService strategy = createAllocationService(settingsBuilder.build(), new TestGatewayAllocator());
+
+        // Build 8 discovery nodes.
+        DiscoveryNodes.Builder discoBuilder = DiscoveryNodes.builder();
+        List<String> nodesList = new ArrayList<>(numberOfNodes);
+        for (int i = 0; i < numberOfNodes; i++) {
+            DiscoveryNode node = newNode("node" + i);
+            discoBuilder.add(node);
+            nodesList.add(node.getId());
+        }
+        discoBuilder.localNodeId(nodesList.get(0));
+        discoBuilder.clusterManagerNodeId(nodesList.get(0));
+
+        IndexMetadata indexMetadata = getIndexMetadata(indexName, numberOfShards, numberOfReplicas);
+
+        int[] primaryOwners = new int[] { 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 3, 3 };
+        int[] replicaOwners = new int[] { 1, 2, 3, 1, 2, 3, 2, 3, 1, 3, 1, 2 };
+
+        IndexRoutingTable.Builder indexRoutingTable = IndexRoutingTable.builder(indexMetadata.getIndex());
+        IndexMetadata.Builder indexMetaDataBuilder = IndexMetadata.builder(indexMetadata);
+
+        for (int shardId = 0; shardId < numberOfShards; shardId++) {
+            ShardId sid = new ShardId(indexMetadata.getIndex(), shardId);
+            IndexShardRoutingTable.Builder shardBuilder = new IndexShardRoutingTable.Builder(sid);
+            shardBuilder.addShard(
+                TestShardRouting.newShardRouting(sid, nodesList.get(primaryOwners[shardId]), true, ShardRoutingState.STARTED)
+            );
+            shardBuilder.addShard(
+                TestShardRouting.newShardRouting(sid, nodesList.get(replicaOwners[shardId]), false, ShardRoutingState.STARTED)
+            );
+            IndexShardRoutingTable shardTable = shardBuilder.build();
+            indexRoutingTable.addIndexShard(shardTable);
+            indexMetaDataBuilder.putInSyncAllocationIds(shardId, shardTable.getAllAllocationIds());
+        }
+
+        Metadata.Builder metadata = Metadata.builder();
+        metadata.persistentSettings(settingsBuilder.build());
+        metadata.put(indexMetaDataBuilder.build(), false);
+
+        RoutingTable.Builder routingTable = RoutingTable.builder();
+        routingTable.add(indexRoutingTable);
+
+        ClusterState.Builder stateBuilder = ClusterState.builder(new ClusterName("test"));
+        stateBuilder.nodes(discoBuilder);
+        stateBuilder.metadata(metadata.generateClusterUuidIfNeeded().build());
+        stateBuilder.routingTable(routingTable.build());
+        ClusterState clusterState = stateBuilder.build();
+
+        // Sanity check: initial primary distribution is 6/2/2/2 on eligible nodes and 0 on excluded.
+        int[] initialPrimaries = countPrimariesPerNode(clusterState, nodesList, indexName);
+        int[] initialTotals = countTotalShardsPerNode(clusterState, nodesList, indexName);
+        logger.info("[before reroute]\n{}", ShardAllocations.printShardDistribution(clusterState));
+        assertEquals("initial primaries on node0", 6, initialPrimaries[0]);
+        assertEquals("initial primaries on node1", 2, initialPrimaries[1]);
+        assertEquals("initial primaries on node2", 2, initialPrimaries[2]);
+        assertEquals("initial primaries on node3", 2, initialPrimaries[3]);
+        for (int i = numberOfNodes - excludedNodeCount; i < numberOfNodes; i++) {
+            assertEquals("initial primary count on excluded " + nodesList.get(i), 0, initialPrimaries[i]);
+        }
+        // Total shard count on every eligible node is identical so the generic shard-balance term is
+        // neutral -- any relocation observed later must come from the primary-only rebalance logic.
+        assertEquals("eligible node0 total shards", 6, initialTotals[0]);
+        assertEquals("eligible node1 total shards", 6, initialTotals[1]);
+        assertEquals("eligible node2 total shards", 6, initialTotals[2]);
+        assertEquals("eligible node3 total shards", 6, initialTotals[3]);
+
+        clusterState = strategy.reroute(clusterState, "reroute");
+        clusterState = applyStartedShardsUntilNoChange(clusterState, strategy);
+        for (int i = 0; i < 20; i++) {
+            ClusterState next = strategy.reroute(clusterState, "reroute-loop-" + i);
+            next = applyStartedShardsUntilNoChange(next, strategy);
+            if (next.equals(clusterState)) {
+                break;
+            }
+            clusterState = next;
+        }
+
+        logger.info("[after reroute]\n{}", ShardAllocations.printShardDistribution(clusterState));
+
+        int[] finalPrimaries = countPrimariesPerNode(clusterState, nodesList, indexName);
+
+        // Excluded nodes must still be empty.
+        for (int i = numberOfNodes - excludedNodeCount; i < numberOfNodes; i++) {
+            assertEquals("excluded node " + nodesList.get(i) + " must remain empty", 0, finalPrimaries[i]);
+        }
+
+        int totalEligiblePrimaries = 0;
+        int maxEligible = Integer.MIN_VALUE;
+        int minEligible = Integer.MAX_VALUE;
+        for (int i = 0; i < numberOfNodes - excludedNodeCount; i++) {
+            totalEligiblePrimaries += finalPrimaries[i];
+            maxEligible = Math.max(maxEligible, finalPrimaries[i]);
+            minEligible = Math.min(minEligible, finalPrimaries[i]);
+        }
+        assertEquals("all primaries should stay on eligible nodes", numberOfShards, totalEligiblePrimaries);
+        assertEquals(
+            "Primaries fully balanced across eligible nodes -- final distribution: "
+                + "node0="
+                + finalPrimaries[0]
+                + ", node1="
+                + finalPrimaries[1]
+                + ", node2="
+                + finalPrimaries[2]
+                + ", node3="
+                + finalPrimaries[3],
+            0,
+            maxEligible - minEligible
+        );
+    }
+
+    private static int[] countPrimariesPerNode(ClusterState state, List<String> nodesList, String indexName) {
+        int[] counts = new int[nodesList.size()];
+        RoutingNodes routingNodes = state.getRoutingNodes();
+        for (int i = 0; i < nodesList.size(); i++) {
+            RoutingNode node = routingNodes.node(nodesList.get(i));
+            if (node == null) {
+                continue;
+            }
+            counts[i] = (int) node.shardsWithState(indexName, STARTED).stream().filter(ShardRouting::primary).count();
+        }
+        return counts;
+    }
+
+    /**
+     * Helper: returns total shard counts (primary + replica, only STARTED) for the given index on each node.
+     */
+    private static int[] countTotalShardsPerNode(ClusterState state, List<String> nodesList, String indexName) {
+        int[] counts = new int[nodesList.size()];
+        RoutingNodes routingNodes = state.getRoutingNodes();
+        for (int i = 0; i < nodesList.size(); i++) {
+            RoutingNode node = routingNodes.node(nodesList.get(i));
+            if (node == null) {
+                continue;
+            }
+            counts[i] = node.shardsWithState(indexName, STARTED).size();
+        }
+        return counts;
+    }
+
 }