Fixed Allocation Explain API in batch mode

Signed-off-by: Shivansh Arora <shivansh.arora@protonmail.com>
opensearch-project · Oct 3, 2023 · e35add0 · e35add0
1 parent d5a95b8
commit e35add0
Showing 1 changed file with 21 additions and 5 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/GatewayAllocator.java
@@ -188,12 +188,28 @@ protected static void innerAllocatedUnassigned(
     public AllocateUnassignedDecision explainUnassignedShardAllocation(ShardRouting unassignedShard, RoutingAllocation routingAllocation) {
         assert unassignedShard.unassigned();
         assert routingAllocation.debugDecision();
-        if (unassignedShard.primary()) {
-            assert primaryShardAllocator != null;
-            return primaryShardAllocator.makeAllocationDecision(unassignedShard, routingAllocation, logger);
+        boolean batchMode = routingAllocation.nodes().getMinNodeVersion().onOrAfter(Version.CURRENT);
+        if (batchMode) {
+            // TODO add integ test for testing this behaviour when shard is unassigned but failed many times and is ultimately removed from batch
+            if (getBatchId(unassignedShard, unassignedShard.primary()) == null) {
+                createAndUpdateBatches(routingAllocation, unassignedShard.primary());
+            }
+            assert getBatchId(unassignedShard, unassignedShard.primary()) != null;
+            if (unassignedShard.primary()) {
+                assert primaryBatchShardAllocator != null;
+                return primaryBatchShardAllocator.makeAllocationDecision(unassignedShard, routingAllocation, logger);
+            } else {
+                assert replicaBatchShardAllocator != null;
+                return replicaBatchShardAllocator.makeAllocationDecision(unassignedShard, routingAllocation, logger);
+            }
         } else {
-            assert replicaShardAllocator != null;
-            return replicaShardAllocator.makeAllocationDecision(unassignedShard, routingAllocation, logger);
+            if (unassignedShard.primary()) {
+                assert primaryShardAllocator != null;
+                return primaryShardAllocator.makeAllocationDecision(unassignedShard, routingAllocation, logger);
+            } else {
+                assert replicaShardAllocator != null;
+                return replicaShardAllocator.makeAllocationDecision(unassignedShard, routingAllocation, logger);
+            }
         }
     }