From 38a9be5a5f4db486f595adec2fb44fb50db39e19 Mon Sep 17 00:00:00 2001
From: Anshu Agarwal <anshukag@amazon.com>
Date: Mon, 19 Sep 2022 10:04:57 +0530
Subject: [PATCH] Add null check for discovery nodes and single mutex for
 weighted shard iterator

---
 .../cluster/routing/IndexShardRoutingTable.java | 17 +++++++++--------
 1 file changed, 9 insertions(+), 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
index fbd8050fd467c..9026e7068e9fe 100644
--- 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();
             }