Add null check for discovery nodes and single mutex for weighted shar…

…d iterator
opensearch-project · Sep 19, 2022 · 38a9be5 · 38a9be5
1 parent 01239bf
commit 38a9be5
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java b/server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java
@@ -85,8 +85,7 @@ public class IndexShardRoutingTable implements Iterable<ShardRouting> {
     private volatile Map<AttributesKey, AttributesRoutings> activeShardsByAttributes = emptyMap();
     private volatile Map<AttributesKey, AttributesRoutings> initializingShardsByAttributes = emptyMap();
     private final Object shardsByAttributeMutex = new Object();
-    private final Object activeShardsByWeightMutex = new Object();
-    private final Object initializingShardsByWeightMutex = new Object();
+    private final Object shardsByWeightMutex = new Object();
     private volatile Map<WeightedRoutingKey, List<ShardRouting>> activeShardsByWeight = emptyMap();
     private volatile Map<WeightedRoutingKey, List<ShardRouting>> initializingShardsByWeight = emptyMap();
 
@@ -357,10 +356,12 @@ private List<WeightedRoundRobin.Entity<ShardRouting>> calculateShardWeight(
         List<WeightedRoundRobin.Entity<ShardRouting>> shardsWithWeights = new ArrayList<>();
         for (ShardRouting shard : shards) {
             DiscoveryNode node = nodes.get(shard.currentNodeId());
-            String attVal = node.getAttributes().get(weightedRouting.attributeName());
-            // If weight for a zone is not defined, considering it as 1 by default
-            Double weight = weightedRouting.weights().getOrDefault(attVal, defaultWeight);
-            shardsWithWeights.add(new WeightedRoundRobin.Entity<>(weight, shard));
+            if (node != null) {
+                String attVal = node.getAttributes().get(weightedRouting.attributeName());
+                // If weight for a zone is not defined, considering it as 1 by default
+                Double weight = weightedRouting.weights().getOrDefault(attVal, defaultWeight);
+                shardsWithWeights.add(new WeightedRoundRobin.Entity<>(weight, shard));
+            }
         }
         return shardsWithWeights;
     }
@@ -807,7 +808,7 @@ private List<ShardRouting> getActiveShardsByWeight(WeightedRouting weightedRouti
         WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
         List<ShardRouting> shardRoutings = activeShardsByWeight.get(key);
         if (shardRoutings == null) {
-            synchronized (activeShardsByWeightMutex) {
+            synchronized (shardsByWeightMutex) {
                 shardRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight);
                 activeShardsByWeight = new MapBuilder().put(key, shardRoutings).immutableMap();
             }
@@ -823,7 +824,7 @@ private List<ShardRouting> getInitializingShardsByWeight(WeightedRouting weighte
         WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting);
         List<ShardRouting> shardRoutings = initializingShardsByWeight.get(key);
         if (shardRoutings == null) {
-            synchronized (initializingShardsByWeightMutex) {
+            synchronized (shardsByWeightMutex) {
                 shardRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight);
                 initializingShardsByWeight = new MapBuilder().put(key, shardRoutings).immutableMap();
             }