Making shard copy count a multiple of attribute count

server/src/internalClusterTest/java/org/opensearch/cluster/allocation/AwarenessAllocationIT.java

@@ -36,12 +36,14 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.admin.cluster.health.ClusterHealthResponse;
+import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.metadata.IndexMetadata.State;
 import org.opensearch.cluster.routing.IndexRoutingTable;
 import org.opensearch.cluster.routing.IndexShardRoutingTable;
 import org.opensearch.cluster.routing.ShardRouting;
+import org.opensearch.cluster.routing.allocation.AwarenessReplicaBalance;
 import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider;
 import org.opensearch.common.Priority;
 import org.opensearch.common.settings.Settings;
@@ -77,6 +79,12 @@ public void testSimpleAwareness() throws Exception {
         logger.info("--> starting 2 nodes on the same rack");
         internalCluster().startNodes(2, Settings.builder().put(commonSettings).put("node.attr.rack_id", "rack_1").build());
 
+        Settings settings = Settings.builder()
+            .put(AwarenessReplicaBalance.CLUSTER_ROUTING_ALLOCATION_AWARENESS_BALANCE_SETTING.getKey(), false)
+            .build();
+        ClusterUpdateSettingsRequest updateSettingsRequest = new ClusterUpdateSettingsRequest();
+        updateSettingsRequest.persistentSettings(settings);
+
         createIndex("test1");
         createIndex("test2");
 

server/src/internalClusterTest/java/org/opensearch/indices/settings/UpdateNumberOfReplicasIT.java

@@ -45,11 +45,11 @@
 import java.io.IOException;
 import java.util.EnumSet;
 
+import static org.hamcrest.Matchers.equalTo;
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
-import static org.hamcrest.Matchers.equalTo;
 
 @OpenSearchIntegTestCase.ClusterScope(minNumDataNodes = 2)
 public class UpdateNumberOfReplicasIT extends OpenSearchIntegTestCase {
@@ -606,4 +606,48 @@ public void testUpdateNumberOfReplicasAllowNoIndices() {
         assertThat(numberOfReplicas, equalTo(0));
     }
 
+    public void testAwarenessReplicaBalance() {
+        createIndex("aware-replica", Settings.builder().put("index.number_of_replicas", 0).build());
+        createIndex(".system-index", Settings.builder().put("index.number_of_replicas", 0).build());
+        manageReplicaBalanceSetting(true);
+        int updated = 0;
+
+        try {
+            // replica count of 1 is ideal
+            client().admin()
+                .indices()
+                .prepareUpdateSettings("aware-replica")
+                .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1))
+                .execute()
+                .actionGet();
+            updated++;
+
+            // system index - should be able to update
+            client().admin()
+                .indices()
+                .prepareUpdateSettings(".system-index")
+                .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 2))
+                .execute()
+                .actionGet();
+            updated++;
+
+            client().admin()
+                .indices()
+                .prepareUpdateSettings("aware-replica")
+                .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 2))
+                .execute()
+                .actionGet();
+            fail("should have thrown an exception about the replica  count");
+
+        } catch (IllegalArgumentException e) {
+            assertEquals(
+                "Validation Failed: 1: expected total copies needs to be a multiple of total awareness attributes [2];",
+                e.getMessage()
+            );
+            assertEquals(2, updated);
+        } finally {
+            manageReplicaBalanceSetting(false);
+        }
+    }
+
 }

server/src/internalClusterTest/java/org/opensearch/indices/template/SimpleIndexTemplateIT.java

@@ -32,18 +32,19 @@
 
 package org.opensearch.indices.template;
 
+import org.junit.After;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.action.admin.indices.alias.Alias;
 import org.opensearch.action.admin.indices.alias.get.GetAliasesResponse;
 import org.opensearch.action.admin.indices.settings.get.GetSettingsResponse;
 import org.opensearch.action.admin.indices.template.get.GetIndexTemplatesResponse;
 import org.opensearch.action.admin.indices.template.put.PutIndexTemplateRequestBuilder;
-
 import org.opensearch.action.bulk.BulkResponse;
 import org.opensearch.action.index.IndexRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.AliasMetadata;
+import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.ParsingException;
 import org.opensearch.common.bytes.BytesArray;
 import org.opensearch.common.settings.Settings;
@@ -52,12 +53,11 @@
 import org.opensearch.index.mapper.MapperParsingException;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.indices.InvalidAliasNameException;
+import org.opensearch.indices.InvalidIndexTemplateException;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.search.SearchHit;
-import org.opensearch.test.OpenSearchIntegTestCase;
 import org.opensearch.test.InternalSettingsPlugin;
-
-import org.junit.After;
+import org.opensearch.test.OpenSearchIntegTestCase;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -68,11 +68,6 @@
 import java.util.List;
 import java.util.Set;
 
-import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
-import static org.opensearch.index.query.QueryBuilders.termQuery;
-import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
-import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
-import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertRequestBuilderThrows;
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.containsString;
@@ -83,6 +78,11 @@
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.nullValue;
+import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
+import static org.opensearch.index.query.QueryBuilders.termQuery;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertRequestBuilderThrows;
 
 public class SimpleIndexTemplateIT extends OpenSearchIntegTestCase {
 
@@ -1029,4 +1029,33 @@ public void testPartitionedTemplate() throws Exception {
         GetSettingsResponse getSettingsResponse = client().admin().indices().prepareGetSettings("test_good").get();
         assertEquals("6", getSettingsResponse.getIndexToSettings().get("test_good").get("index.routing_partition_size"));
     }
+
+    public void testAwarenessReplicaBalance() throws IOException {
+        manageReplicaBalanceSetting(true);
+        try {
+            client().admin()
+                .indices()
+                .preparePutTemplate("template_1")
+                .setPatterns(Arrays.asList("a*", "b*"))
+                .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1))
+                .get();
+
+            client().admin()
+                .indices()
+                .preparePutTemplate("template_1")
+                .setPatterns(Arrays.asList("a*", "b*"))
+                .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0))
+                .get();
+
+            fail("should have thrown an exception about the replica  count");
+        } catch (InvalidIndexTemplateException e) {
+            assertEquals(
+                "index_template [template_1] invalid, cause [Validation Failed: 1: expected total copies needs to be a multiple of total awareness attributes [2];]",
+                e.getMessage()
+            );
+        } finally {
+            manageReplicaBalanceSetting(false);
+        }
+    }
+
 }

server/src/internalClusterTest/java/org/opensearch/snapshots/RestoreSnapshotIT.java

@@ -59,6 +59,14 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import static org.hamcrest.Matchers.allOf;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_REPLICAS;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS;
 import static org.opensearch.index.IndexSettings.INDEX_REFRESH_INTERVAL_SETTING;
@@ -70,14 +78,6 @@
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertIndexTemplateExists;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertIndexTemplateMissing;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertRequestBuilderThrows;
-import static org.hamcrest.Matchers.allOf;
-import static org.hamcrest.Matchers.containsString;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.not;
-import static org.hamcrest.Matchers.notNullValue;
-import static org.hamcrest.Matchers.nullValue;
 
 public class RestoreSnapshotIT extends AbstractSnapshotIntegTestCase {
 
@@ -973,4 +973,54 @@ public void testForbidDisableSoftDeletesDuringRestore() throws Exception {
         );
         assertThat(restoreError.getMessage(), containsString("cannot disable setting [index.soft_deletes.enabled] on restore"));
     }
+
+    public void testRestoreBalancedReplica() {
+        try {
+            createRepository("test-repo", "fs");
+            createIndex("test-index", Settings.builder().put("index.number_of_replicas", 0).build());
+            createIndex(".system-index", Settings.builder().put("index.number_of_replicas", 0).build());
+            ensureGreen();
+            clusterAdmin().prepareCreateSnapshot("test-repo", "snapshot-0")
+                .setIndices("test-index", ".system-index")
+                .setWaitForCompletion(true)
+                .get();
+            manageReplicaBalanceSetting(true);
+
+            final IllegalArgumentException restoreError = expectThrows(
+                IllegalArgumentException.class,
+                () -> clusterAdmin().prepareRestoreSnapshot("test-repo", "snapshot-0")
+                    .setRenamePattern("test-index")
+                    .setRenameReplacement("new-index")
+                    .setIndices("test-index")
+                    .get()
+            );
+            assertThat(
+                restoreError.getMessage(),
+                containsString("expected total copies needs to be a multiple of total awareness attributes [2]")
+            );
+
+            RestoreSnapshotResponse restoreSnapshotResponse = clusterAdmin().prepareRestoreSnapshot("test-repo", "snapshot-0")
+                .setRenamePattern(".system-index")
+                .setRenameReplacement(".system-index-restore-1")
+                .setWaitForCompletion(true)
+                .setIndices(".system-index")
+                .execute()
+                .actionGet();
+
+            assertThat(restoreSnapshotResponse.getRestoreInfo().totalShards(), greaterThan(0));
+
+            restoreSnapshotResponse = clusterAdmin().prepareRestoreSnapshot("test-repo", "snapshot-0")
+                .setRenamePattern("test-index")
+                .setRenameReplacement("new-index")
+                .setIndexSettings(Settings.builder().put("index.number_of_replicas", 1).build())
+                .setWaitForCompletion(true)
+                .setIndices("test-index")
+                .execute()
+                .actionGet();
+
+            assertThat(restoreSnapshotResponse.getRestoreInfo().totalShards(), greaterThan(0));
+        } finally {
+            manageReplicaBalanceSetting(false);
+        }
+    }
 }

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

@@ -59,6 +59,7 @@
 import org.opensearch.cluster.routing.ShardRouting;
 import org.opensearch.cluster.routing.ShardRoutingState;
 import org.opensearch.cluster.routing.allocation.AllocationService;
+import org.opensearch.cluster.routing.allocation.AwarenessReplicaBalance;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.Priority;
@@ -102,6 +103,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
@@ -145,6 +147,7 @@ public class MetadataCreateIndexService {
     private final ShardLimitValidator shardLimitValidator;
     private final boolean forbidPrivateIndexSettings;
     private final Set<IndexSettingProvider> indexSettingProviders = new HashSet<>();
+    private AwarenessReplicaBalance awarenessReplicaBalance;
 
     public MetadataCreateIndexService(
         final Settings settings,
@@ -158,7 +161,8 @@ public MetadataCreateIndexService(
         final ThreadPool threadPool,
         final NamedXContentRegistry xContentRegistry,
         final SystemIndices systemIndices,
-        final boolean forbidPrivateIndexSettings
+        final boolean forbidPrivateIndexSettings,
+        final AwarenessReplicaBalance awarenessReplicaBalance
     ) {
         this.settings = settings;
         this.clusterService = clusterService;
@@ -172,6 +176,7 @@ public MetadataCreateIndexService(
         this.systemIndices = systemIndices;
         this.forbidPrivateIndexSettings = forbidPrivateIndexSettings;
         this.shardLimitValidator = shardLimitValidator;
+        this.awarenessReplicaBalance = awarenessReplicaBalance;
     }
 
     /**
@@ -1150,7 +1155,7 @@ private void validate(CreateIndexClusterStateUpdateRequest request, ClusterState
 
     public void validateIndexSettings(String indexName, final Settings settings, final boolean forbidPrivateIndexSettings)
         throws IndexCreationException {
-        List<String> validationErrors = getIndexSettingsValidationErrors(settings, forbidPrivateIndexSettings);
+        List<String> validationErrors = getIndexSettingsValidationErrors(settings, forbidPrivateIndexSettings, indexName);
 
         if (validationErrors.isEmpty() == false) {
             ValidationException validationException = new ValidationException();
@@ -1159,11 +1164,31 @@ public void validateIndexSettings(String indexName, final Settings settings, fin
         }
     }
 
-    List<String> getIndexSettingsValidationErrors(final Settings settings, final boolean forbidPrivateIndexSettings) {
+    List<String> getIndexSettingsValidationErrors(final Settings settings, final boolean forbidPrivateIndexSettings, String indexName) {
+        List<String> validationErrors = getIndexSettingsValidationErrors(settings, forbidPrivateIndexSettings, Optional.of(indexName));
+        return validationErrors;
+    }
+
+    List<String> getIndexSettingsValidationErrors(
+        final Settings settings,
+        final boolean forbidPrivateIndexSettings,
+        Optional<String> indexName
+    ) {
         List<String> validationErrors = validateIndexCustomPath(settings, env.sharedDataDir());
         if (forbidPrivateIndexSettings) {
             validationErrors.addAll(validatePrivateSettingsNotExplicitlySet(settings, indexScopedSettings));
         }
+        if (indexName.isEmpty() || indexName.get().charAt(0) != '.') {
+            // Apply aware replica balance only to non system indices
+            int replicaCount = settings.getAsInt(
+                IndexMetadata.SETTING_NUMBER_OF_REPLICAS,
+                INDEX_NUMBER_OF_REPLICAS_SETTING.getDefault(Settings.EMPTY)
+            );
+            Optional<String> error = awarenessReplicaBalance.validate(replicaCount);
+            if (error.isPresent()) {
+                validationErrors.add(error.get());
+            }
+        }
         return validationErrors;
     }
 

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

@@ -1471,7 +1471,11 @@ private void validate(String name, @Nullable Settings settings, List<String> ind
                     validationErrors.add(t.getMessage());
                 }
             }
-            List<String> indexSettingsValidation = metadataCreateIndexService.getIndexSettingsValidationErrors(settings, true);
+            List<String> indexSettingsValidation = metadataCreateIndexService.getIndexSettingsValidationErrors(
+                settings,
+                true,
+                Optional.empty()
+            );
             validationErrors.addAll(indexSettingsValidation);
         }