From e1a632fd8a88b0fad3d11708dd389c88eb0eeaa3 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Fri, 14 Jun 2024 12:49:16 +0530
Subject: [PATCH] Switch to ShardsBatchGatewayAllocator in batch enabled mode
 for shard explain allocation API (#14107) (#14340)

(cherry picked from commit 2e13e9cb5b3507e9e7e85b73012c7ccd84b6844f)

Signed-off-by: Swetha Guptha <gupthasg@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Swetha Guptha <gupthasg@amazon.com>
---
 .../routing/allocation/AllocationService.java | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationService.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationService.java
index 58b9753a11d73..ee3f4701c9f4d 100644
--- a/server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationService.java
+++ b/server/src/main/java/org/opensearch/cluster/routing/allocation/AllocationService.java
@@ -584,10 +584,7 @@ private void allocateExistingUnassignedShards(RoutingAllocation allocation) {
         /*
          Use batch mode if enabled and there is no custom allocator set for Allocation service
          */
-        Boolean batchModeEnabled = EXISTING_SHARDS_ALLOCATOR_BATCH_MODE.get(settings);
-        if (batchModeEnabled
-            && allocation.nodes().getMinNodeVersion().onOrAfter(Version.V_2_14_0)
-            && existingShardsAllocators.size() == 2) {
+        if (isBatchModeEnabled(allocation)) {
             /*
              If we do not have any custom allocator set then we will be using ShardsBatchGatewayAllocator
              Currently AllocationService will not run any custom Allocator that implements allocateAllUnassignedShards
@@ -724,13 +721,24 @@ private AllocateUnassignedDecision explainUnassignedShardAllocation(ShardRouting
 
     private ExistingShardsAllocator getAllocatorForShard(ShardRouting shardRouting, RoutingAllocation routingAllocation) {
         assert assertInitialized();
-        final String allocatorName = ExistingShardsAllocator.EXISTING_SHARDS_ALLOCATOR_SETTING.get(
-            routingAllocation.metadata().getIndexSafe(shardRouting.index()).getSettings()
-        );
+        String allocatorName;
+        if (isBatchModeEnabled(routingAllocation)) {
+            allocatorName = ShardsBatchGatewayAllocator.ALLOCATOR_NAME;
+        } else {
+            allocatorName = ExistingShardsAllocator.EXISTING_SHARDS_ALLOCATOR_SETTING.get(
+                routingAllocation.metadata().getIndexSafe(shardRouting.index()).getSettings()
+            );
+        }
         final ExistingShardsAllocator existingShardsAllocator = existingShardsAllocators.get(allocatorName);
         return existingShardsAllocator != null ? existingShardsAllocator : new NotFoundAllocator(allocatorName);
     }
 
+    private boolean isBatchModeEnabled(RoutingAllocation routingAllocation) {
+        return EXISTING_SHARDS_ALLOCATOR_BATCH_MODE.get(settings)
+            && routingAllocation.nodes().getMinNodeVersion().onOrAfter(Version.V_2_14_0)
+            && existingShardsAllocators.size() == 2;
+    }
+
     private boolean assertInitialized() {
         assert existingShardsAllocators != null : "must have set allocators first";
         return true;