Switch to ShardsBatchGatewayAllocator in batch enabled mode for shard…

… explain allocation API (opensearch-project#14107) Signed-off-by: Swetha Guptha <[email protected]> Co-authored-by: Swetha Guptha <[email protected]>
kkewwei · Jun 12, 2024 · 2e13e9c · 2e13e9c
1 parent 5093901
commit 2e13e9c
Showing 1 changed file with 15 additions and 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
@@ -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;