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 shards, RoutingAllocation allocation) { + public void allocateUnassignedBatch(Set shardRoutings, RoutingAllocation allocation) { // make Allocation Decisions for all shards - HashMap decisionMap = makeAllocationDecision(shards, allocation, logger); - assert shards.size() == decisionMap.size() : "make allocation decision didn't return allocation decision for " + "some shards"; + HashMap 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 makeAllocationDecision( - Set shards, + Set unassignedShardBatch, RoutingAllocation allocation, Logger logger ) { - HashMap allocationDecisions = new HashMap<>(); - for (ShardRouting unassignedShard : shards) { - allocationDecisions.put(unassignedShard, makeAllocationDecision(unassignedShard, allocation, logger)); - } - return allocationDecisions; + + return (HashMap) 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 fetchData(ShardRouting shard, RoutingAllocation allocation); - @Override - // to be override - public HashMap makeAllocationDecision( - Set shards, - RoutingAllocation allocation, - Logger logger - ) { - return null; - } - private static class NodeShardsResult { final List 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 fetchData(ShardRouting shard, RoutingAllocation allocation); - @Override - // to be override - public HashMap makeAllocationDecision( - Set 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. */