From b7524aed2842169daed38a7a03601436be8f7f4c Mon Sep 17 00:00:00 2001
From: Gaurav Chandani <chngau@amazon.com>
Date: Tue, 5 Dec 2023 15:15:03 +0530
Subject: [PATCH] Addressed PR comments(#8776)

1. Nit javadoc and renaming variables

Signed-off-by: Gaurav Chandani <chngau@amazon.com>
---
 .../gateway/BaseGatewayShardAllocator.java    | 27 +++++++++++--------
 .../gateway/PrimaryShardAllocator.java        | 11 --------
 .../gateway/ReplicaShardAllocator.java        | 10 -------
 3 files changed, 16 insertions(+), 32 deletions(-)

diff --git a/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java b/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java
index 865eeccb8bb82..0e4ca52d92f8e 100644
--- a/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java
+++ b/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java
@@ -49,6 +49,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * An abstract class that implements basic functionality for allocating
@@ -82,14 +83,15 @@ public void allocateUnassigned(
     }
 
     /**
-     * Allocate Set of unassigned shard  to nodes where valid copies of the shard already exists
-     * @param shards the shards to allocate
+     * Allocate Batch of unassigned shard  to nodes where valid copies of the shard already exists
+     * @param shardRoutings the shards to allocate
      * @param allocation the allocation state container object
      */
-    public void allocateUnassignedBatch(Set<ShardRouting> shards, RoutingAllocation allocation) {
+    public void allocateUnassignedBatch(Set<ShardRouting> shardRoutings, RoutingAllocation allocation) {
         // make Allocation Decisions for all shards
-        HashMap<ShardRouting, AllocateUnassignedDecision> decisionMap = makeAllocationDecision(shards, allocation, logger);
-        assert shards.size() == decisionMap.size() : "make allocation decision didn't return allocation decision for " + "some shards";
+        HashMap<ShardRouting, AllocateUnassignedDecision> decisionMap = makeAllocationDecision(shardRoutings, allocation, logger);
+        assert shardRoutings.size() == decisionMap.size() : "make allocation decision didn't return allocation decision for "
+            + "some shards";
         // get all unassigned shards iterator
         RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator();
 
@@ -159,15 +161,18 @@ public abstract AllocateUnassignedDecision makeAllocationDecision(
     );
 
     public HashMap<ShardRouting, AllocateUnassignedDecision> makeAllocationDecision(
-        Set<ShardRouting> shards,
+        Set<ShardRouting> unassignedShardBatch,
         RoutingAllocation allocation,
         Logger logger
     ) {
-        HashMap<ShardRouting, AllocateUnassignedDecision> allocationDecisions = new HashMap<>();
-        for (ShardRouting unassignedShard : shards) {
-            allocationDecisions.put(unassignedShard, makeAllocationDecision(unassignedShard, allocation, logger));
-        }
-        return allocationDecisions;
+
+        return (HashMap<ShardRouting, AllocateUnassignedDecision>) unassignedShardBatch.stream()
+            .collect(
+                Collectors.toMap(
+                    unassignedShard -> unassignedShard,
+                    unassignedShard -> makeAllocationDecision(unassignedShard, allocation, logger)
+                )
+            );
     }
 
     /**
diff --git a/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java b/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java
index 91dc6ecbda604..4dc9396751fc9 100644
--- a/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java
+++ b/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java
@@ -56,7 +56,6 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -458,16 +457,6 @@ private static NodesToAllocate buildNodesToAllocate(
 
     protected abstract FetchResult<NodeGatewayStartedShards> fetchData(ShardRouting shard, RoutingAllocation allocation);
 
-    @Override
-    // to be override
-    public HashMap<ShardRouting, AllocateUnassignedDecision> makeAllocationDecision(
-        Set<ShardRouting> shards,
-        RoutingAllocation allocation,
-        Logger logger
-    ) {
-        return null;
-    }
-
     private static class NodeShardsResult {
         final List<NodeGatewayStartedShards> orderedAllocationCandidates;
         final int allocationsFound;
diff --git a/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java b/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java
index 4dfbba8d4fe3c..c0b831b6fe4d0 100644
--- a/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java
+++ b/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java
@@ -495,16 +495,6 @@ private static boolean canPerformOperationBasedRecovery(
 
     protected abstract AsyncShardFetch.FetchResult<NodeStoreFilesMetadata> fetchData(ShardRouting shard, RoutingAllocation allocation);
 
-    @Override
-    // to be override
-    public HashMap<ShardRouting, AllocateUnassignedDecision> makeAllocationDecision(
-        Set<ShardRouting> shards,
-        RoutingAllocation allocation,
-        Logger logger
-    ) {
-        return null;
-    }
-
     /**
      * Returns a boolean indicating whether fetching shard data has been triggered at any point for the given shard.
      */