From 4dd5219309e9e4484d650556cb9ac87b5f0d1670 Mon Sep 17 00:00:00 2001
From: Aman Khare <amkhar@amazon.com>
Date: Tue, 5 Sep 2023 11:40:52 +0530
Subject: [PATCH] Minor refactoring

Signed-off-by: Aman Khare <amkhar@amazon.com>
---
 .../opensearch/gateway/PrimaryShardAllocator.java   |  7 ++-----
 .../gateway/PrimaryShardBatchAllocator.java         | 13 +------------
 2 files changed, 3 insertions(+), 17 deletions(-)

diff --git a/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java b/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java
index c8783e39049d0..7a8472a4fd28b 100644
--- a/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java
+++ b/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java
@@ -51,7 +51,6 @@
 import org.opensearch.env.ShardLockObtainFailedException;
 import org.opensearch.gateway.AsyncShardFetch.FetchResult;
 import org.opensearch.gateway.TransportNodesListGatewayStartedShards.NodeGatewayStartedShards;
-import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -60,9 +59,7 @@
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
-import java.util.TreeMap;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -101,7 +98,7 @@ protected static boolean isResponsibleFor(final ShardRouting shard) {
      * @param allocation      routing allocation object
      * @return allocation decision taken for this shard
      */
-    protected AllocateUnassignedDecision skipSnapshotRestore(ShardRouting unassignedShard, RoutingAllocation allocation) {
+    protected AllocateUnassignedDecision getInEligibleShardDecision(ShardRouting unassignedShard, RoutingAllocation allocation) {
         if (isResponsibleFor(unassignedShard) == false) {
             // this allocator is not responsible for allocating this shard
             return AllocateUnassignedDecision.NOT_TAKEN;
@@ -124,7 +121,7 @@ public AllocateUnassignedDecision makeAllocationDecision(
         final RoutingAllocation allocation,
         final Logger logger
     ) {
-        AllocateUnassignedDecision decision = skipSnapshotRestore(unassignedShard, allocation);
+        AllocateUnassignedDecision decision = getInEligibleShardDecision(unassignedShard, allocation);
         if (decision != null) {
             return decision;
         }
diff --git a/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java b/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java
index 803e8dae2ef1a..26cec71b208cc 100644
--- a/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java
+++ b/server/src/main/java/org/opensearch/gateway/PrimaryShardBatchAllocator.java
@@ -33,36 +33,25 @@
 package org.opensearch.gateway;
 
 import org.apache.logging.log4j.Logger;
-import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.opensearch.cluster.node.DiscoveryNode;
-import org.opensearch.cluster.routing.RecoverySource;
 import org.opensearch.cluster.routing.RoutingNode;
 import org.opensearch.cluster.routing.RoutingNodes;
 import org.opensearch.cluster.routing.ShardRouting;
-import org.opensearch.cluster.routing.UnassignedInfo;
 import org.opensearch.cluster.routing.UnassignedInfo.AllocationStatus;
 import org.opensearch.cluster.routing.allocation.AllocateUnassignedDecision;
 import org.opensearch.cluster.routing.allocation.NodeAllocationResult;
 import org.opensearch.cluster.routing.allocation.RoutingAllocation;
-import org.opensearch.env.ShardLockObtainFailedException;
 import org.opensearch.gateway.AsyncShardFetch.FetchResult;
 import org.opensearch.gateway.TransportNodesListGatewayStartedShardsBatch.NodeGatewayStartedShardsBatch;
 import org.opensearch.gateway.TransportNodesListGatewayStartedShardsBatch.NodeGatewayStartedShards;
-import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.TreeMap;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 /**
  * PrimaryShardBatchAllocator is similar to {@link org.opensearch.gateway.PrimaryShardAllocator} only difference is
@@ -130,7 +119,7 @@ public HashMap<ShardRouting, AllocateUnassignedDecision> makeAllocationDecision(
             if (matchingShard == null) {
                 matchingShard = shard;
             }
-            AllocateUnassignedDecision decision = skipSnapshotRestore(matchingShard, allocation);
+            AllocateUnassignedDecision decision = getInEligibleShardDecision(matchingShard, allocation);
             if (decision != null) {
                 shardsNotEligibleForFetch.add(shard);
                 shardAllocationDecisions.put(shard, decision);