Delete batch if empty when removing shards from batch

Signed-off-by: Aman Khare <[email protected]>
Gaurav614 · Apr 12, 2024 · 9b1d553 · 9b1d553
1 parent ed9df68
commit 9b1d553
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java
@@ -251,7 +251,7 @@ else if (shardRouting.primary() == primary) {
             }
         });
 
-        refreshShardBatches(currentBatches, batchedShardsToAssign);
+        refreshShardBatches(currentBatches, batchedShardsToAssign, primary);
 
         Iterator<ShardRouting> iterator = newShardsToBatch.iterator();
         assert maxBatchSize > 0 : "Shards batch size must be greater than 0";
@@ -283,7 +283,7 @@ else if (shardRouting.primary() == primary) {
         return batchesToBeAssigned;
     }
 
-    private void refreshShardBatches(ConcurrentMap<String, ShardsBatch> currentBatches, Set<ShardId> batchedShardsToAssign) {
+    private void refreshShardBatches(ConcurrentMap<String, ShardsBatch> currentBatches, Set<ShardId> batchedShardsToAssign, boolean primary) {
         // cleanup shard from batches if they are not present in unassigned list from allocation object. This is
         // needed as AllocationService.reroute can also be called directly by API flows for example DeleteIndices.
         // So, as part of calling reroute, those shards will be removed from allocation object. It'll handle the
@@ -297,6 +297,8 @@ private void refreshShardBatches(ConcurrentMap<String, ShardsBatch> currentBatch
                     batchEntry.getValue().clearShardFromCache(shardId);
                 }
             }
+            ConcurrentMap<String, ShardsBatch> batches = primary ? batchIdToStartedShardBatch : batchIdToStoreShardBatch;
+            deleteBatchIfEmpty(batches, batchEntry.getValue().getBatchId());
         }
     }
 

diff --git a/...src/main/java/org/opensearch/indices/store/TransportNodesListShardStoreMetadataBatch.java b/...src/main/java/org/opensearch/indices/store/TransportNodesListShardStoreMetadataBatch.java
@@ -329,7 +329,13 @@ public static class NodeStoreFilesMetadataBatch extends BaseNodeResponse {
 
         protected NodeStoreFilesMetadataBatch(StreamInput in) throws IOException {
             super(in);
-            this.nodeStoreFilesMetadataBatch = in.readMap(ShardId::new, NodeStoreFilesMetadata::new);
+            this.nodeStoreFilesMetadataBatch = in.readMap(ShardId::new, i -> {
+                if (i.readBoolean()) {
+                    return new NodeStoreFilesMetadata(i);
+                } else {
+                    return null;
+                }
+            });
         }
 
         public NodeStoreFilesMetadataBatch(DiscoveryNode node, Map<ShardId, NodeStoreFilesMetadata> nodeStoreFilesMetadataBatch) {
@@ -344,7 +350,14 @@ public Map<ShardId, NodeStoreFilesMetadata> getNodeStoreFilesMetadataBatch() {
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
-            out.writeMap(nodeStoreFilesMetadataBatch, (o, k) -> k.writeTo(o), (o, v) -> v.writeTo(o));
+            out.writeMap(nodeStoreFilesMetadataBatch, (o, k) -> k.writeTo(o), (o, v) -> {
+                if (v != null) {
+                    o.writeBoolean(true);
+                    v.writeTo(o);
+                } else {
+                    o.writeBoolean(false);
+                }
+            });
         }
     }