Route search traffic to _primary_first for warm index (#14934)

Signed-off-by: Neetika Singhal <neetiks@amazon.com>
opensearch-project · Jul 31, 2024 · 0324edd · 0324edd
1 parent e7ee950
commit 0324edd
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java b/server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java
@@ -42,8 +42,10 @@
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.Strings;
 import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.index.IndexModule;
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.node.ResponseCollectorService;
 
@@ -245,6 +247,13 @@ public GroupShardsIterator<ShardIterator> searchShards(
                 preference = Preference.PRIMARY.type();
             }
 
+            if (FeatureFlags.isEnabled(FeatureFlags.TIERED_REMOTE_INDEX)
+                && IndexModule.DataLocalityType.PARTIAL.name()
+                    .equals(indexMetadataForShard.getSettings().get(IndexModule.INDEX_STORE_LOCALITY_SETTING.getKey()))
+                && (preference == null || preference.isEmpty())) {
+                preference = Preference.PRIMARY_FIRST.type();
+            }
+
             ShardIterator iterator = preferenceActiveShardIterator(
                 shard,
                 clusterState.nodes().getLocalNodeId(),

diff --git a/server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java b/server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java
@@ -41,9 +41,11 @@
 import org.opensearch.cluster.node.DiscoveryNodeRole;
 import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.SuppressForbidden;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.core.index.Index;
 import org.opensearch.core.index.shard.ShardId;
@@ -1054,6 +1056,68 @@ public void testSearchableSnapshotPrimaryDefault() throws Exception {
         }
     }
 
+    @SuppressForbidden(reason = "feature flag overrides")
+    public void testPartialIndexPrimaryDefault() throws Exception {
+        System.setProperty(FeatureFlags.TIERED_REMOTE_INDEX, "true");
+        final int numIndices = 1;
+        final int numShards = 2;
+        final int numReplicas = 2;
+        final String[] indexNames = new String[numIndices];
+        for (int i = 0; i < numIndices; i++) {
+            indexNames[i] = "test" + i;
+        }
+        // The first index is a partial index
+        final String indexName = indexNames[0];
+        ClusterService clusterService = null;
+        ThreadPool threadPool = null;
+
+        try {
+            OperationRouting opRouting = new OperationRouting(
+                Settings.EMPTY,
+                new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)
+            );
+
+            ClusterState state = ClusterStateCreationUtils.stateWithAssignedPrimariesAndReplicas(indexNames, numShards, numReplicas);
+            threadPool = new TestThreadPool("testPartialIndexPrimaryDefault");
+            clusterService = ClusterServiceUtils.createClusterService(threadPool);
+
+            // Update the index config within the cluster state to modify the index to a partial index
+            IndexMetadata partialIndexMetadata = IndexMetadata.builder(indexName)
+                .settings(
+                    Settings.builder()
+                        .put(state.metadata().index(indexName).getSettings())
+                        .put(IndexModule.INDEX_STORE_LOCALITY_SETTING.getKey(), IndexModule.DataLocalityType.PARTIAL)
+                        .build()
+                )
+                .build();
+            Metadata.Builder metadataBuilder = Metadata.builder(state.metadata())
+                .put(partialIndexMetadata, false)
+                .generateClusterUuidIfNeeded();
+            state = ClusterState.builder(state).metadata(metadataBuilder.build()).build();
+
+            // Verify default preference is primary only
+            GroupShardsIterator<ShardIterator> groupIterator = opRouting.searchShards(state, indexNames, null, null);
+            assertThat("One group per index shard", groupIterator.size(), equalTo(numIndices * numShards));
+
+            for (ShardIterator shardIterator : groupIterator) {
+                assertTrue("Only primary should exist with no preference", shardIterator.nextOrNull().primary());
+            }
+
+            // Verify alternative preference can be applied to a partial index
+            groupIterator = opRouting.searchShards(state, indexNames, null, "_replica");
+            assertThat("One group per index shard", groupIterator.size(), equalTo(numIndices * numShards));
+
+            for (ShardIterator shardIterator : groupIterator) {
+                assertThat("Replica shards will be returned", shardIterator.size(), equalTo(numReplicas));
+                assertFalse("Returned shard should be a replica", shardIterator.nextOrNull().primary());
+            }
+        } finally {
+            IOUtils.close(clusterService);
+            terminate(threadPool);
+            System.setProperty(FeatureFlags.TIERED_REMOTE_INDEX, "false");
+        }
+    }
+
     private DiscoveryNode[] setupNodes() {
         // Sets up two data nodes in zone-a and one data node in zone-b
         List<String> zones = Arrays.asList("a", "a", "b");