Add new index and cluster level settings to limit the total primary shards per node and per index

CHANGELOG-3.0.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Views, simplify data access and manipulation by providing a virtual layer over one or more indices ([#11957](https://github.com/opensearch-project/OpenSearch/pull/11957))
 - Added pull-based Ingestion (APIs, for ingestion source, a Kafka plugin, and IngestionEngine that pulls data from the ingestion source) ([#16958](https://github.com/opensearch-project/OpenSearch/pull/16958))
 - Added ConfigurationUtils to core for the ease of configuration parsing [#17223](https://github.com/opensearch-project/OpenSearch/pull/17223)
+- Added index level setting to limit the total primary shards per node per index [#17295](https://github.com/opensearch-project/OpenSearch/pull/17295)
 
 ### Dependencies
 - Update Apache Lucene to 10.1.0 ([#16366](https://github.com/opensearch-project/OpenSearch/pull/16366))

server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java

@@ -817,6 +817,7 @@ public void validate(final String value, final Map<Setting<?>, Object> settings)
     private final boolean isRemoteSnapshot;
 
     private final int indexTotalShardsPerNodeLimit;
+    private final int indexTotalPrimaryShardsPerNodeLimit;
     private final boolean isAppendOnlyIndex;
 
     private final Context context;
@@ -849,6 +850,7 @@ private IndexMetadata(
         final Map<String, RolloverInfo> rolloverInfos,
         final boolean isSystem,
         final int indexTotalShardsPerNodeLimit,
+        final int indexTotalPrimaryShardsPerNodeLimit,
         boolean isAppendOnlyIndex,
         final Context context
     ) {
@@ -887,6 +889,7 @@ private IndexMetadata(
         this.isSystem = isSystem;
         this.isRemoteSnapshot = IndexModule.Type.REMOTE_SNAPSHOT.match(this.settings);
         this.indexTotalShardsPerNodeLimit = indexTotalShardsPerNodeLimit;
+        this.indexTotalPrimaryShardsPerNodeLimit = indexTotalPrimaryShardsPerNodeLimit;
         this.isAppendOnlyIndex = isAppendOnlyIndex;
         this.context = context;
         assert numberOfShards * routingFactor == routingNumShards : routingNumShards + " must be a multiple of " + numberOfShards;
@@ -1069,6 +1072,10 @@ public int getIndexTotalShardsPerNodeLimit() {
         return this.indexTotalShardsPerNodeLimit;
     }
 
+    public int getIndexTotalPrimaryShardsPerNodeLimit() {
+        return this.indexTotalPrimaryShardsPerNodeLimit;
+    }
+
     public boolean isAppendOnlyIndex() {
         return this.isAppendOnlyIndex;
     }
@@ -1867,6 +1874,9 @@ public IndexMetadata build() {
             }
 
             final int indexTotalShardsPerNodeLimit = ShardsLimitAllocationDecider.INDEX_TOTAL_SHARDS_PER_NODE_SETTING.get(settings);
+            final int indexTotalPrimaryShardsPerNodeLimit = ShardsLimitAllocationDecider.INDEX_TOTAL_PRIMARY_SHARDS_PER_NODE_SETTING.get(
+                settings
+            );
             final boolean isAppendOnlyIndex = INDEX_APPEND_ONLY_ENABLED_SETTING.get(settings);
 
             final String uuid = settings.get(SETTING_INDEX_UUID, INDEX_UUID_NA_VALUE);
@@ -1899,6 +1909,7 @@ public IndexMetadata build() {
                 rolloverInfos,
                 isSystem,
                 indexTotalShardsPerNodeLimit,
+                indexTotalPrimaryShardsPerNodeLimit,
                 isAppendOnlyIndex,
                 context
             );

server/src/main/java/org/opensearch/cluster/routing/RoutingNode.java

@@ -393,6 +393,20 @@
         }
     }
 
+    public int numberOfOwningPrimaryShardsForIndex(final Index index) {
+        final LinkedHashSet<ShardRouting> shardRoutings = shardsByIndex.get(index);
+        if (shardRoutings == null) {
+            return 0;
+        } else {
+            return Math.toIntExact(
+                shardRoutings.stream()
+                    .filter(sr -> sr.relocating() == false)
+                    .filter(ShardRouting::primary)    // Add this filter for primary shards
+                    .count()
+            );
+        }
+    }
+
     public String prettyPrint() {
         StringBuilder sb = new StringBuilder();
         sb.append("-----node_id[").append(nodeId).append("][").append(node == null ? "X" : "V").append("]\n");

server/src/main/java/org/opensearch/cluster/routing/allocation/decider/ShardsLimitAllocationDecider.java

@@ -32,6 +32,7 @@
 
 package org.opensearch.cluster.routing.allocation.decider;
 
+import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.routing.RoutingNode;
 import org.opensearch.cluster.routing.ShardRouting;
 import org.opensearch.cluster.routing.ShardRoutingState;
@@ -40,19 +41,20 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.indices.replication.common.ReplicationType;
 
 import java.util.function.BiPredicate;
 
 /**
  * This {@link AllocationDecider} limits the number of shards per node on a per
  * index or node-wide basis. The allocator prevents a single node to hold more
- * than {@code index.routing.allocation.total_shards_per_node} per index and
+ * than {@code index.routing.allocation.total_shards_per_node} per index, {@code index.routing.allocation.total_primary_shards_per_node} per index and
  * {@code cluster.routing.allocation.total_shards_per_node} globally during the allocation
  * process. The limits of this decider can be changed in real-time via a the
  * index settings API.
  * <p>
- * If {@code index.routing.allocation.total_shards_per_node} is reset to a negative value shards
- * per index are unlimited per node. Shards currently in the
+ * If {@code index.routing.allocation.total_shards_per_node} or {@code index.routing.allocation.total_primary_shards_per_node}is reset to a negative value shards
+ * per index are unlimited per node or primary shards per index are unlimited per node respectively. Shards currently in the
  * {@link ShardRoutingState#RELOCATING relocating} state are ignored by this
  * {@link AllocationDecider} until the shard changed its state to either
  * {@link ShardRoutingState#STARTED started},
@@ -83,6 +85,18 @@ public class ShardsLimitAllocationDecider extends AllocationDecider {
         Property.IndexScope
     );
 
+    /**
+     * Controls the maximum number of primary shards per index on a single OpenSearch
+     * node for segment replication enabled indices. Negative values are interpreted as unlimited.
+     */
+    public static final Setting<Integer> INDEX_TOTAL_PRIMARY_SHARDS_PER_NODE_SETTING = Setting.intSetting(
+        "index.routing.allocation.total_primary_shards_per_node",
+        -1,
+        -1,
+        Property.Dynamic,
+        Property.IndexScope
+    );
+
     /**
      * Controls the maximum number of shards per node on a global level.
      * Negative values are interpreted as unlimited.
@@ -115,7 +129,15 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
     @Override
     public Decision canRemain(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
         return doDecide(shardRouting, node, allocation, (count, limit) -> count > limit);
+    }
 
+    /**
+     * Checks whether the given shardRouting's index uses segment replication
+     */
+    private boolean isSegRepEnabledIndex(RoutingAllocation allocation, ShardRouting shardRouting) {
+        IndexMetadata indexMetadata = allocation.metadata().getIndexSafe(shardRouting.index());
+        Settings indexSettings = indexMetadata.getSettings();
+        return IndexMetadata.INDEX_REPLICATION_TYPE_SETTING.get(indexSettings) == ReplicationType.SEGMENT;
     }
 
     private Decision doDecide(
@@ -125,11 +147,13 @@ private Decision doDecide(
         BiPredicate<Integer, Integer> decider
     ) {
         final int indexShardLimit = allocation.metadata().getIndexSafe(shardRouting.index()).getIndexTotalShardsPerNodeLimit();
+        final int indexPrimaryShardLimit = allocation.metadata()
+            .getIndexSafe(shardRouting.index())
+            .getIndexTotalPrimaryShardsPerNodeLimit();
         // Capture the limit here in case it changes during this method's
         // execution
         final int clusterShardLimit = this.clusterShardLimit;
-
-        if (indexShardLimit <= 0 && clusterShardLimit <= 0) {
+        if (indexShardLimit <= 0 && indexPrimaryShardLimit <= 0 && clusterShardLimit <= 0) {
             return allocation.decision(
                 Decision.YES,
                 NAME,
@@ -165,6 +189,22 @@ private Decision doDecide(
                 );
             }
         }
+
+        final boolean isIndexSegRepEnabled = isSegRepEnabledIndex(allocation, shardRouting);
+        if (indexPrimaryShardLimit > 0 && isIndexSegRepEnabled && shardRouting.primary()) {
+            final int indexPrimaryShardCount = node.numberOfOwningPrimaryShardsForIndex(shardRouting.index());
+            if (decider.test(indexPrimaryShardCount, indexPrimaryShardLimit)) {
+                return allocation.decision(
+                    Decision.NO,
+                    NAME,
+                    "too many primary shards [%d] allocated to this node for index [%s], index setting [%s=%d]",
+                    indexPrimaryShardCount,
+                    shardRouting.getIndexName(),
+                    INDEX_TOTAL_PRIMARY_SHARDS_PER_NODE_SETTING.getKey(),
+                    indexPrimaryShardLimit
+                );
+            }
+        }
         return allocation.decision(
             Decision.YES,
             NAME,

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

@@ -166,6 +166,7 @@ public final class IndexScopedSettings extends AbstractScopedSettings {
                 IndexSettings.MAX_SLICES_PER_PIT,
                 IndexSettings.MAX_REGEX_LENGTH_SETTING,
                 ShardsLimitAllocationDecider.INDEX_TOTAL_SHARDS_PER_NODE_SETTING,
+                ShardsLimitAllocationDecider.INDEX_TOTAL_PRIMARY_SHARDS_PER_NODE_SETTING,
                 IndexSettings.INDEX_GC_DELETES_SETTING,
                 IndexSettings.INDEX_SOFT_DELETES_SETTING,
                 IndexSettings.INDEX_SOFT_DELETES_RETENTION_OPERATIONS_SETTING,