Cache the shard routings with no weight for faster access

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

@@ -34,7 +34,6 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.opensearch.cluster.metadata.WeightedRoutingMetadata;
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.common.Nullable;
@@ -63,7 +62,6 @@
 import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 import static java.util.Collections.emptyMap;
 
@@ -96,8 +94,8 @@ public class IndexShardRoutingTable implements Iterable<ShardRouting> {
     private volatile Map<AttributesKey, AttributesRoutings> initializingShardsByAttributes = emptyMap();
     private final Object shardsByAttributeMutex = new Object();
     private final Object shardsByWeightMutex = new Object();
-    private volatile Map<WeightedRoutingKey, List<ShardRouting>> activeShardsByWeight = emptyMap();
-    private volatile Map<WeightedRoutingKey, List<ShardRouting>> initializingShardsByWeight = emptyMap();
+    private volatile Map<WeightedRoutingKey, WeightedShardRoutings> activeShardsByWeight = emptyMap();
+    private volatile Map<WeightedRoutingKey, WeightedShardRoutings> initializingShardsByWeight = emptyMap();
 
     private static final Logger logger = LogManager.getLogger(IndexShardRoutingTable.class);
 
@@ -249,7 +247,7 @@ public List<ShardRouting> assignedShards() {
         return this.assignedShards;
     }
 
-    public Map<WeightedRoutingKey, List<ShardRouting>> getActiveShardsByWeight() {
+    public Map<WeightedRoutingKey, WeightedShardRoutings> getActiveShardsByWeight() {
         return activeShardsByWeight;
     }
 
@@ -338,23 +336,7 @@ public ShardIterator activeInitializingShardsWeightedIt(
         // append shards for attribute value with weight zero, so that shard search requests can be tried on
         // shard copies in case of request failure from other attribute values.
         if (isFailOpenEnabled) {
-            try {
-                Stream<String> keys = weightedRouting.weights()
-                    .entrySet()
-                    .stream()
-                    .filter(entry -> entry.getValue().intValue() == WeightedRoutingMetadata.WEIGHED_AWAY_WEIGHT)
-                    .map(Map.Entry::getKey);
-                keys.forEach(key -> {
-                    ShardIterator iterator = onlyNodeSelectorActiveInitializingShardsIt(weightedRouting.attributeName() + ":" + key, nodes);
-                    while (iterator.remaining() > 0) {
-                        ordered.add(iterator.nextOrNull());
-                    }
-                });
-            } catch (IllegalArgumentException e) {
-                // this exception is thrown by {@link onlyNodeSelectorActiveInitializingShardsIt} in case count of shard
-                // copies found is zero
-                logger.debug("no shard copies found for shard id [{}] for node attribute with weight zero", shardId);
-            }
+            ordered.addAll(activeInitializingShardsWithoutWeights(weightedRouting, nodes, defaultWeight));
         }
 
         return new PlainShardIterator(shardId, ordered);
@@ -378,6 +360,18 @@ private List<ShardRouting> activeInitializingShardsWithWeights(
         return orderedListWithDistinctShards;
     }
 
+    private List<ShardRouting> activeInitializingShardsWithoutWeights(
+        WeightedRouting weightedRouting,
+        DiscoveryNodes nodes,
+        double defaultWeight
+    ) {
+        List<ShardRouting> ordered = new ArrayList<>(getActiveShardsWithoutWeight(weightedRouting, nodes, defaultWeight));
+        if (!allInitializingShards.isEmpty()) {
+            ordered.addAll(getInitializingShardsWithoutWeight(weightedRouting, nodes, defaultWeight));
+        }
+        return ordered.stream().distinct().collect(Collectors.toList());
+    }
+
     /**
      * Returns a list containing shard routings ordered using weighted round-robin scheduling.
      */
@@ -949,20 +943,55 @@ public int hashCode() {
         }
     }
 
+    @PublicApi(since = "2.14.0")
+    public static class WeightedShardRoutings {
+        private final List<ShardRouting> shardRoutingsWithWeight;
+        private final List<ShardRouting> shardRoutingWithoutWeight;
+
+        public WeightedShardRoutings(List<ShardRouting> shardRoutingsWithWeight, List<ShardRouting> shardRoutingWithoutWeight) {
+            this.shardRoutingsWithWeight = Collections.unmodifiableList(shardRoutingsWithWeight);
+            this.shardRoutingWithoutWeight = Collections.unmodifiableList(shardRoutingWithoutWeight);
+        }
+
+        public List<ShardRouting> getShardRoutingsWithWeight() {
+            return shardRoutingsWithWeight;
+        }
+
+        public List<ShardRouting> getShardRoutingWithoutWeight() {
+            return shardRoutingWithoutWeight;
+        }
+    }
+
     /**
      * *
      * Gets active shard routing from memory if available, else calculates and put it in memory.
      */
     private List<ShardRouting> getActiveShardsByWeight(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) {
         WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
-        List<ShardRouting> shardRoutings = activeShardsByWeight.get(key);
-        if (shardRoutings == null) {
-            synchronized (shardsByWeightMutex) {
-                shardRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight);
-                activeShardsByWeight = new MapBuilder().put(key, shardRoutings).immutableMap();
-            }
+        if (activeShardsByWeight.get(key) == null) {
+            populateActiveShardWeightsMap(weightedRouting, nodes, defaultWeight);
+        }
+        return activeShardsByWeight.get(key).getShardRoutingsWithWeight();
+    }
+
+    private List<ShardRouting> getActiveShardsWithoutWeight(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) {
+        WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
+        if (activeShardsByWeight.get(key) == null) {
+            populateActiveShardWeightsMap(weightedRouting, nodes, defaultWeight);
+        }
+        return activeShardsByWeight.get(key).getShardRoutingWithoutWeight();
+    }
+
+    private void populateActiveShardWeightsMap(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) {
+        WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
+        List<ShardRouting> weightedRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight);
+        List<ShardRouting> nonWeightedRoutings = activeShards.stream()
+            .filter(shard -> !weightedRoutings.contains(shard))
+            .collect(Collectors.toUnmodifiableList());
+        synchronized (shardsByWeightMutex) {
+            activeShardsByWeight = new MapBuilder().put(key, new WeightedShardRoutings(weightedRoutings, nonWeightedRoutings))
+                .immutableMap();
         }
-        return shardRoutings;
     }
 
     /**
@@ -971,14 +1000,34 @@ private List<ShardRouting> getActiveShardsByWeight(WeightedRouting weightedRouti
      */
     private List<ShardRouting> getInitializingShardsByWeight(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) {
         WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
-        List<ShardRouting> shardRoutings = initializingShardsByWeight.get(key);
-        if (shardRoutings == null) {
-            synchronized (shardsByWeightMutex) {
-                shardRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight);
-                initializingShardsByWeight = new MapBuilder().put(key, shardRoutings).immutableMap();
-            }
+        if (initializingShardsByWeight.get(key) == null) {
+            populateInitializingShardWeightsMap(weightedRouting, nodes, defaultWeight);
+        }
+        return initializingShardsByWeight.get(key).getShardRoutingsWithWeight();
+    }
+
+    private List<ShardRouting> getInitializingShardsWithoutWeight(
+        WeightedRouting weightedRouting,
+        DiscoveryNodes nodes,
+        double defaultWeight
+    ) {
+        WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
+        if (initializingShardsByWeight.get(key) == null) {
+            populateInitializingShardWeightsMap(weightedRouting, nodes, defaultWeight);
+        }
+        return initializingShardsByWeight.get(key).getShardRoutingWithoutWeight();
+    }
+
+    private void populateInitializingShardWeightsMap(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) {
+        WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
+        List<ShardRouting> weightedRoutings = shardsOrderedByWeight(allInitializingShards, weightedRouting, nodes, defaultWeight);
+        List<ShardRouting> nonWeightedRoutings = allInitializingShards.stream()
+            .filter(shard -> !weightedRoutings.contains(shard))
+            .collect(Collectors.toUnmodifiableList());
+        synchronized (shardsByWeightMutex) {
+            initializingShardsByWeight = new MapBuilder().put(key, new WeightedShardRoutings(weightedRoutings, nonWeightedRoutings))
+                .immutableMap();
         }
-        return shardRoutings;
     }
 
     /**

server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java

@@ -700,9 +700,18 @@ public void testWeightedRoutingWithDifferentWeights() {
                 .shard(0)
                 .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1, false, null);
             assertEquals(1, shardIterator.size());
-            shardRouting = shardIterator.nextOrNull();
-            assertNotNull(shardRouting);
-            assertFalse(Arrays.asList("node2", "node1").contains(shardRouting.currentNodeId()));
+            assertEquals("node3", shardIterator.nextOrNull().currentNodeId());
+
+            weights = Map.of("zone1", -1.0, "zone2", 0.0, "zone3", 1.0);
+            weightedRouting = new WeightedRouting("zone", weights);
+            shardIterator = clusterState.routingTable()
+                .index("test")
+                .shard(0)
+                .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1, true, null);
+            assertEquals(3, shardIterator.size());
+            assertEquals("node3", shardIterator.nextOrNull().currentNodeId());
+            assertNotEquals("node3", shardIterator.nextOrNull().currentNodeId());
+            assertNotEquals("node3", shardIterator.nextOrNull().currentNodeId());
 
             weights = Map.of("zone1", 3.0, "zone2", 2.0, "zone3", 0.0);
             weightedRouting = new WeightedRouting("zone", weights);
@@ -711,8 +720,9 @@ public void testWeightedRoutingWithDifferentWeights() {
                 .shard(0)
                 .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1, true, null);
             assertEquals(3, shardIterator.size());
-            shardRouting = shardIterator.nextOrNull();
-            assertNotNull(shardRouting);
+            assertNotEquals("node3", shardIterator.nextOrNull().currentNodeId());
+            assertNotEquals("node3", shardIterator.nextOrNull().currentNodeId());
+            assertEquals("node3", shardIterator.nextOrNull().currentNodeId());
         } finally {
             terminate(threadPool);
         }
@@ -887,6 +897,7 @@ public void testWeightedRoutingShardState() {
                 shardRouting = shardIterator.nextOrNull();
                 assertNotNull(shardRouting);
                 requestCount.put(shardRouting.currentNodeId(), requestCount.getOrDefault(shardRouting.currentNodeId(), 0) + 1);
+
             }
             assertEquals(3, requestCount.get("node1").intValue());
             assertEquals(2, requestCount.get("node2").intValue());