Making shard copy count a multiple of attribute count

server/src/internalClusterTest/java/org/opensearch/action/admin/indices/rollover/RolloverIT.java

@@ -61,6 +61,7 @@
 import java.util.List;
 import java.util.Set;
 
+import static org.hamcrest.Matchers.containsString;
 import static org.opensearch.index.mapper.MapperService.SINGLE_MAPPING_NAME;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -172,6 +173,7 @@ public void testRolloverWithNoWriteIndex() {
     }
 
     public void testRolloverWithIndexSettings() throws Exception {
+
         Alias testAlias = new Alias("test_alias");
         boolean explicitWriteIndex = randomBoolean();
         if (explicitWriteIndex) {
@@ -210,6 +212,47 @@ public void testRolloverWithIndexSettings() throws Exception {
         }
     }
 
+    public void testRolloverWithIndexSettingsBalancedReplica() throws Exception {
+        Alias testAlias = new Alias("test_alias");
+        boolean explicitWriteIndex = randomBoolean();
+        if (explicitWriteIndex) {
+            testAlias.writeIndex(true);
+        }
+        assertAcked(prepareCreate("test_index-2").addAlias(testAlias).get());
+        manageReplicaBalanceSetting(true);
+        index("test_index-2", "type1", "1", "field", "value");
+        flush("test_index-2");
+        final Settings settings = Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            .build();
+
+        final IllegalArgumentException restoreError = expectThrows(
+            IllegalArgumentException.class,
+            () -> client().admin().indices().prepareRolloverIndex("test_alias").settings(settings).alias(new Alias("extra_alias")).get()
+        );
+
+        assertThat(
+            restoreError.getMessage(),
+            containsString("expected total copies needs to be a multiple of total awareness attributes [2]")
+        );
+
+        final Settings balancedReplicaSettings = Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+            .build();
+
+        client().admin()
+            .indices()
+            .prepareRolloverIndex("test_alias")
+            .settings(balancedReplicaSettings)
+            .alias(new Alias("extra_alias"))
+            .waitForActiveShards(0)
+            .get();
+
+        manageReplicaBalanceSetting(false);
+    }
+
     public void testRolloverWithIndexSettingsWithoutPrefix() throws Exception {
         Alias testAlias = new Alias("test_alias");
         boolean explicitWriteIndex = randomBoolean();

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);
+        }
+    }
 }