Exempt server-injected preferences from strict weighted routing check…

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

@@ -265,30 +265,40 @@ public GroupShardsIterator<ShardIterator> searchShards(
         for (IndexShardRoutingTable shard : shards) {
 
             IndexMetadata indexMetadataForShard = indexMetadata(clusterState, shard.shardId.getIndex().getName());
-            if (indexMetadataForShard.isRemoteSnapshot() && (preference == null || preference.isEmpty())) {
-                preference = Preference.PRIMARY.type();
+            // Server-injected preferences are scoped per index: compute the effective preference for each shard
+            // instead of mutating the caller-supplied preference, so that a preference injected for one index
+            // (e.g. a remote snapshot or writable warm index) does not leak into other indices of the same
+            // multi-index search request.
+            String effectivePreference = preference;
+            boolean serverInjectedPreference = false;
+            if (indexMetadataForShard.isRemoteSnapshot() && (effectivePreference == null || effectivePreference.isEmpty())) {
+                effectivePreference = Preference.PRIMARY.type();
+                serverInjectedPreference = true;
             }
 
             if (FeatureFlags.isEnabled(FeatureFlags.WRITABLE_WARM_INDEX_EXPERIMENTAL_FLAG)
                 && indexMetadataForShard.getSettings().getAsBoolean(IndexModule.IS_WARM_INDEX_SETTING.getKey(), false)
-                && (preference == null || preference.isEmpty())) {
-                preference = Preference.PRIMARY_FIRST.type();
+                && (effectivePreference == null || effectivePreference.isEmpty())) {
+                effectivePreference = Preference.PRIMARY_FIRST.type();
+                serverInjectedPreference = true;
             }
 
-            if (preference == null || preference.isEmpty()) {
+            if (effectivePreference == null || effectivePreference.isEmpty()) {
                 if (indexMetadataForShard.getNumberOfSearchOnlyReplicas() > 0 && isStrictSearchOnlyShardRouting) {
-                    preference = Preference.SEARCH_REPLICA.type();
+                    effectivePreference = Preference.SEARCH_REPLICA.type();
+                    serverInjectedPreference = true;
                 }
             }
 
             ShardIterator iterator = preferenceActiveShardIterator(
                 shard,
                 clusterState.nodes().getLocalNodeId(),
                 clusterState.nodes(),
-                preference,
+                effectivePreference,
                 collectorService,
                 nodeCounts,
-                clusterState.metadata().weightedRoutingMetadata()
+                clusterState.metadata().weightedRoutingMetadata(),
+                serverInjectedPreference
             );
             if (iterator != null) {
                 shardIterators.computeIfAbsent(iterator.shardId().getIndex(), k -> new ArrayList<>()).add(iterator);
@@ -362,6 +372,28 @@ private ShardIterator preferenceActiveShardIterator(
         @Nullable ResponseCollectorService collectorService,
         @Nullable Map<String, Long> nodeCounts,
         @Nullable WeightedRoutingMetadata weightedRoutingMetadata
+    ) {
+        return preferenceActiveShardIterator(
+            indexShard,
+            localNodeId,
+            nodes,
+            preference,
+            collectorService,
+            nodeCounts,
+            weightedRoutingMetadata,
+            false
+        );
+    }
+
+    private ShardIterator preferenceActiveShardIterator(
+        IndexShardRoutingTable indexShard,
+        String localNodeId,
+        DiscoveryNodes nodes,
+        @Nullable String preference,
+        @Nullable ResponseCollectorService collectorService,
+        @Nullable Map<String, Long> nodeCounts,
+        @Nullable WeightedRoutingMetadata weightedRoutingMetadata,
+        boolean serverInjectedPreference
     ) {
         if (preference == null || preference.isEmpty()) {
             return shardRoutings(indexShard, nodes, collectorService, nodeCounts, weightedRoutingMetadata);
@@ -399,7 +431,7 @@ private ShardIterator preferenceActiveShardIterator(
                 }
             }
             preferenceType = Preference.parse(preference);
-            checkPreferenceBasedRoutingAllowed(preferenceType, weightedRoutingMetadata);
+            checkPreferenceBasedRoutingAllowed(preferenceType, weightedRoutingMetadata, serverInjectedPreference);
             switch (preferenceType) {
                 case PREFER_NODES:
                     final Set<String> nodesIds = Arrays.stream(preference.substring(Preference.PREFER_NODES.type().length() + 1).split(","))
@@ -565,12 +597,22 @@ private static int calculateShardIdOfChild(
         return indexMetadata.getSplitShardsMetadata().getShardIdOfHash(rootShardId, hash, includeInProgressChild);
     }
 
-    private void checkPreferenceBasedRoutingAllowed(Preference preference, @Nullable WeightedRoutingMetadata weightedRoutingMetadata) {
-        if (WeightedRoutingUtils.shouldPerformStrictWeightedRouting(
-            isStrictWeightedShardRouting,
-            ignoreWeightedRouting,
-            weightedRoutingMetadata
-        ) && WEIGHTED_ROUTING_RESTRICTED_PREFERENCES.contains(preference)) {
+    // package-private for testing
+    void checkPreferenceBasedRoutingAllowed(
+        Preference preference,
+        @Nullable WeightedRoutingMetadata weightedRoutingMetadata,
+        boolean serverInjectedPreference
+    ) {
+        // Preferences injected by the server for specific index types (e.g. remote snapshot or writable warm
+        // indices) are the engine's own routing decision and are exempt from this check, which exists to stop
+        // caller-supplied preferences from bypassing weighted shard routing.
+        if (serverInjectedPreference == false
+            && WeightedRoutingUtils.shouldPerformStrictWeightedRouting(
+                isStrictWeightedShardRouting,
+                ignoreWeightedRouting,
+                weightedRoutingMetadata
+            )
+            && WEIGHTED_ROUTING_RESTRICTED_PREFERENCES.contains(preference)) {
             throw new PreferenceBasedSearchNotAllowedException(
                 "Preference type based routing not allowed with strict weighted shard routing enabled"
             );

server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java

@@ -1119,6 +1119,141 @@ public void testPartialIndexPrimaryDefault() throws Exception {
         }
     }
 
+    public void testServerInjectedPreferenceDoesNotLeakAcrossIndices() throws Exception {
+        final int numIndices = 4;
+        final int numShards = 3;
+        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 remote snapshot index, for which the server injects the _primary preference
+        final String remoteSnapshotIndex = 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("testServerInjectedPreferenceDoesNotLeakAcrossIndices");
+            clusterService = ClusterServiceUtils.createClusterService(threadPool);
+
+            IndexMetadata remoteSnapshotIndexMetadata = IndexMetadata.builder(remoteSnapshotIndex)
+                .settings(
+                    Settings.builder()
+                        .put(state.metadata().index(remoteSnapshotIndex).getSettings())
+                        .put(IndexModule.INDEX_STORE_TYPE_SETTING.getKey(), IndexModule.Type.REMOTE_SNAPSHOT.getSettingsKey())
+                        .build()
+                )
+                .build();
+            Metadata.Builder metadataBuilder = Metadata.builder(state.metadata())
+                .put(remoteSnapshotIndexMetadata, false)
+                .generateClusterUuidIfNeeded();
+            state = ClusterState.builder(state).metadata(metadataBuilder.build()).build();
+
+            GroupShardsIterator<ShardIterator> groupIterator = opRouting.searchShards(state, indexNames, null, null);
+            assertThat("One group per index shard", groupIterator.size(), equalTo(numIndices * numShards));
+
+            for (ShardIterator shardIterator : groupIterator) {
+                if (remoteSnapshotIndex.equals(shardIterator.shardId().getIndexName())) {
+                    // Server-injected _primary preference applies to the remote snapshot index
+                    assertThat("Only the primary should be returned for the remote snapshot index", shardIterator.size(), equalTo(1));
+                    assertTrue("Returned shard should be the primary", shardIterator.nextOrNull().primary());
+                } else {
+                    // The injected preference must not leak into the other indices of the same request:
+                    // they should route across all active shard copies
+                    assertThat("All shard copies should be returned for a regular index", shardIterator.size(), equalTo(1 + numReplicas));
+                }
+            }
+        } finally {
+            IOUtils.close(clusterService);
+            terminate(threadPool);
+        }
+    }
+
+    public void testStrictWeightedRoutingExemptsOnlyServerInjectedPreferences() {
+        OperationRouting opRouting = new OperationRouting(
+            Settings.builder().put(OperationRouting.STRICT_WEIGHTED_SHARD_ROUTING_ENABLED.getKey(), true).build(),
+            new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)
+        );
+        WeightedRouting weightedRouting = new WeightedRouting("zone", Map.of("a", 1.0, "b", 1.0, "c", 0.0));
+        WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting, 0);
+
+        // A caller-supplied restricted preference must be rejected under strict weighted routing
+        expectThrows(
+            PreferenceBasedSearchNotAllowedException.class,
+            () -> opRouting.checkPreferenceBasedRoutingAllowed(Preference.PREFER_NODES, weightedRoutingMetadata, false)
+        );
+
+        // The same preference must be allowed when it was injected by the server itself
+        opRouting.checkPreferenceBasedRoutingAllowed(Preference.PREFER_NODES, weightedRoutingMetadata, true);
+    }
+
+    @LockFeatureFlag(WRITABLE_WARM_INDEX_EXPERIMENTAL_FLAG)
+    @SuppressForbidden(reason = "feature flag overrides")
+    public void testWarmIndexSearchNotBlockedByStrictWeightedRouting() throws Exception {
+        final int numShards = 2;
+        final int numReplicas = 2;
+        final String[] indexNames = { "test0", "test1" };
+        // The first index is a writable warm index, for which the server injects the _primary_first preference
+        final String warmIndex = indexNames[0];
+        ClusterService clusterService = null;
+        ThreadPool threadPool = null;
+
+        try {
+            ClusterState state = clusterStateForWeightedRouting(indexNames, numShards, numReplicas);
+
+            Settings setting = Settings.builder()
+                .put("cluster.routing.allocation.awareness.attributes", "zone")
+                .put(OperationRouting.STRICT_WEIGHTED_SHARD_ROUTING_ENABLED.getKey(), true)
+                .build();
+
+            threadPool = new TestThreadPool("testWarmIndexSearchNotBlockedByStrictWeightedRouting");
+            clusterService = ClusterServiceUtils.createClusterService(threadPool);
+
+            OperationRouting opRouting = new OperationRouting(
+                setting,
+                new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)
+            );
+
+            IndexMetadata warmIndexMetadata = IndexMetadata.builder(warmIndex)
+                .settings(
+                    Settings.builder()
+                        .put(state.metadata().index(warmIndex).getSettings())
+                        .put(IndexModule.IS_WARM_INDEX_SETTING.getKey(), true)
+                        .build()
+                )
+                .build();
+            Metadata.Builder metadataBuilder = Metadata.builder(state.metadata())
+                .put(warmIndexMetadata, false)
+                .generateClusterUuidIfNeeded();
+            state = ClusterState.builder(state).metadata(metadataBuilder.build()).build();
+
+            // Weighted shard routing is active with a zero-weight zone (e.g. Multi-AZ with Standby)
+            Map<String, Double> weights = Map.of("a", 1.0, "b", 1.0, "c", 0.0);
+            state = setWeightedRoutingWeights(state, weights);
+            ClusterServiceUtils.setState(clusterService, ClusterState.builder(state));
+
+            // A search with no caller-supplied preference must not be rejected by the strict weighted
+            // routing preference check, even though the server injects _primary_first for the warm index
+            GroupShardsIterator<ShardIterator> groupIterator = opRouting.searchShards(state, indexNames, null, null);
+            assertThat("One group per index shard", groupIterator.size(), equalTo(indexNames.length * numShards));
+
+            for (ShardIterator shardIterator : groupIterator) {
+                if (warmIndex.equals(shardIterator.shardId().getIndexName())) {
+                    assertTrue("Primary should be returned first for the warm index", shardIterator.nextOrNull().primary());
+                }
+            }
+        } finally {
+            IOUtils.close(clusterService);
+            terminate(threadPool);
+        }
+    }
+
     public void testSearchReplicaDefaultRouting() throws Exception {
         final int numShards = 1;
         final int numReplicas = 2;