Cache the shard routings with no weight for faster access

The list of shards to run a query is determined for every request and
the weight of the nodes guides the shard selection. Currently, IndexRoutingTable
caches the shard routings with weight for faster access. But, during cases
where the fail open option is enabled, shards with no weight is also returned
lower in the order along with shards with weights. They will be used as fall
back if the shards with weights can't be used due to some error.

The shard routing with no weight is not cached, hence it does a full loop for
every request, this impacts the search latency when the number of shards to
query or the number of nodes in the cluster is high. The latency impact is
very high when both the number of shards and the number of nodes are high.

This change introduces a caching mechanism for shard routing with no weights
similar to the existing cache for shard routing with weights.

Signed-off-by: Prabhakar Sithanandam <prabhakar.s87@gmail.com>

Author: Prabhakar Sithanandam

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());