Add batch async shard fetch transport action for replica #8218

server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java

@@ -879,27 +879,36 @@ public void testShardStoreFetchNodeNotConnectedUsingBatchAction() {
 
     public void testShardStoreFetchCorruptedIndexUsingBatchAction() throws Exception {
         internalCluster().startNodes(2);
-        String indexName = "test";
-        prepareIndices(new String[] { indexName }, 1, 1);
-        Map<ShardId, ShardAttributes> shardAttributesMap = prepareRequestMap(new String[] { indexName }, 1);
-        Index index = resolveIndex(indexName);
-        ShardId shardId = new ShardId(index, 0);
-        ClusterSearchShardsResponse searchShardsResponse = client().admin().cluster().prepareSearchShards(indexName).get();
+        String index1Name = "test1";
+        String index2Name = "test2";
+        prepareIndices(new String[] { index1Name, index2Name }, 1, 1);
+        Map<ShardId, ShardAttributes> shardAttributesMap = prepareRequestMap(new String[] { index1Name, index2Name }, 1);
+        Index index1 = resolveIndex(index1Name);
+        ShardId shardId1 = new ShardId(index1, 0);
+        ClusterSearchShardsResponse searchShardsResponse = client().admin().cluster().prepareSearchShards(index1Name).get();
         assertEquals(2, searchShardsResponse.getNodes().length);
-        corruptShard(searchShardsResponse.getNodes()[0].getName(), shardId);
-        corruptShard(searchShardsResponse.getNodes()[1].getName(), shardId);
+
+        // corrupt test1 index shards
+        corruptShard(searchShardsResponse.getNodes()[0].getName(), shardId1);
+        corruptShard(searchShardsResponse.getNodes()[1].getName(), shardId1);
         ClusterRerouteResponse clusterRerouteResponse = client().admin().cluster().prepareReroute().setRetryFailed(false).get();
         DiscoveryNode[] discoveryNodes = getDiscoveryNodes();
         TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch response;
         response = ActionTestUtils.executeBlocking(
             internalCluster().getInstance(TransportNodesListShardStoreMetadataBatch.class),
             new TransportNodesListShardStoreMetadataBatch.Request(shardAttributesMap, discoveryNodes)
         );
-        TransportNodesListShardStoreMetadataBatch.NodeStoreFilesMetadata nodeStoreFilesMetadata = response.getNodesMap()
+        Map<ShardId, TransportNodesListShardStoreMetadataBatch.NodeStoreFilesMetadata> nodeStoreFilesMetadata = response.getNodesMap()
             .get(discoveryNodes[0].getId())
-            .getNodeStoreFilesMetadataBatch()
-            .get(shardId);
-        assertNodeStoreFilesMetadataFailureCase(nodeStoreFilesMetadata, shardId);
+            .getNodeStoreFilesMetadataBatch();
+        // We don't store exception in case of corrupt index, rather just return an empty response
+        assertNull(nodeStoreFilesMetadata.get(shardId1).getException());
+        assertEquals(shardId1, nodeStoreFilesMetadata.get(shardId1).storeFilesMetadata().shardId());
+        assertTrue(nodeStoreFilesMetadata.get(shardId1).storeFilesMetadata().isEmpty());
+
+        Index index2 = resolveIndex(index2Name);
+        ShardId shardId2 = new ShardId(index2, 0);
+        assertNodeStoreFilesMetadataSuccessCase(nodeStoreFilesMetadata.get(shardId2), shardId2);
     }
 
     private void prepareIndices(String[] indices, int numberOfPrimaryShards, int numberOfReplicaShards) {
@@ -930,21 +939,11 @@ private TransportNodesListShardStoreMetadataBatch.NodesStoreFilesMetadataBatch p
         );
     }
 
-    private void assertNodeStoreFilesMetadataFailureCase(
-        TransportNodesListShardStoreMetadataBatch.NodeStoreFilesMetadata nodeStoreFilesMetadata,
-        ShardId shardId
-    ) {
-        assertNotNull(nodeStoreFilesMetadata.getStoreFileFetchException());
-        TransportNodesListShardStoreMetadataHelper.StoreFilesMetadata storeFileMetadata = nodeStoreFilesMetadata.storeFilesMetadata();
-        assertEquals(shardId, storeFileMetadata.shardId());
-        assertTrue(storeFileMetadata.peerRecoveryRetentionLeases().isEmpty());
-    }
-
     private void assertNodeStoreFilesMetadataSuccessCase(
         TransportNodesListShardStoreMetadataBatch.NodeStoreFilesMetadata nodeStoreFilesMetadata,
         ShardId shardId
     ) {
-        assertNull(nodeStoreFilesMetadata.getStoreFileFetchException());
+        assertNull(nodeStoreFilesMetadata.getException());
         TransportNodesListShardStoreMetadataHelper.StoreFilesMetadata storeFileMetadata = nodeStoreFilesMetadata.storeFilesMetadata();
         assertFalse(storeFileMetadata.isEmpty());
         assertEquals(shardId, storeFileMetadata.shardId());

server/src/main/java/org/opensearch/gateway/BaseShardResponse.java

@@ -0,0 +1,51 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.gateway;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+/**
+ * Base response class for shard response.
+ *
+ * @opensearch.internal
+ */
+public abstract class BaseShardResponse {
+
+    private final Exception storeException;
+
+    public BaseShardResponse(Exception storeException) {
+        this.storeException = storeException;
+    }
+
+    public abstract boolean isEmpty();
+
+    public Exception getException() {
+        return storeException;
+    }
+
+    public BaseShardResponse(StreamInput in) throws IOException {
+        if (in.readBoolean()) {
+            storeException = in.readException();
+        } else {
+            storeException = null;
+        }
+    }
+
+    public void writeTo(StreamOutput out) throws IOException {
+        if (storeException != null) {
+            out.writeBoolean(true);
+            out.writeException(storeException);
+        } else {
+            out.writeBoolean(false);
+        }
+    }
+}

server/src/main/java/org/opensearch/indices/store/TransportNodesListShardStoreMetadata.java

@@ -52,20 +52,18 @@
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.gateway.AsyncShardFetch;
-import org.opensearch.index.store.Store;
 import org.opensearch.indices.IndicesService;
 import org.opensearch.indices.store.TransportNodesListShardStoreMetadataHelper.StoreFilesMetadata;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportRequest;
 import org.opensearch.transport.TransportService;
 
 import java.io.IOException;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
-import static org.opensearch.indices.store.TransportNodesListShardStoreMetadataHelper.getListShardMetadataOnLocalNode;
+import static org.opensearch.indices.store.TransportNodesListShardStoreMetadataHelper.listShardMetadataInternal;
 
 /**
  * Metadata for shard stores from a list of transport nodes
@@ -157,19 +155,7 @@
 
     private StoreFilesMetadata listStoreMetadata(NodeRequest request) throws IOException {
         final ShardId shardId = request.getShardId();
-        try {
-            return getListShardMetadataOnLocalNode(
-                logger,
-                shardId,
-                nodeEnv,
-                indicesService,
-                request.getCustomDataPath(),
-                settings,
-                clusterService
-            );
-        } catch (IOException e) {
-            return new StoreFilesMetadata(shardId, Store.MetadataSnapshot.EMPTY, Collections.emptyList());
-        }
+        return listShardMetadataInternal(logger, shardId, nodeEnv, indicesService, request.getCustomDataPath(), settings, clusterService);
     }
 
     /**

server/src/main/java/org/opensearch/indices/store/TransportNodesListShardStoreMetadataBatch.java

@@ -27,6 +27,7 @@
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.gateway.AsyncShardFetch;
+import org.opensearch.gateway.BaseShardResponse;
 import org.opensearch.index.store.Store;
 import org.opensearch.indices.IndicesService;
 import org.opensearch.indices.store.TransportNodesListShardStoreMetadataHelper.StoreFilesMetadata;
@@ -41,6 +42,8 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.opensearch.indices.store.TransportNodesListShardStoreMetadataHelper.INDEX_NOT_FOUND;
+
 /**
  * Transport action for fetching the batch of shard stores Metadata from a list of transport nodes
  *
@@ -98,17 +101,17 @@
         DiscoveryNode[] nodes,
         ActionListener<NodesStoreFilesMetadataBatch> listener
     ) {
         execute(new TransportNodesListShardStoreMetadataBatch.Request(shardAttributes, nodes), listener);
     }
 
     @Override
     protected NodeRequest newNodeRequest(Request request) {
         return new NodeRequest(request);
     }
 
     @Override
     protected NodeStoreFilesMetadataBatch newNodeResponse(StreamInput in) throws IOException {
         return new NodeStoreFilesMetadataBatch(in);
     }
 
     @Override
@@ -117,16 +120,16 @@
         List<NodeStoreFilesMetadataBatch> responses,
         List<FailedNodeException> failures
     ) {
         return new NodesStoreFilesMetadataBatch(clusterService.getClusterName(), responses, failures);
     }
 
     @Override
     protected NodeStoreFilesMetadataBatch nodeOperation(NodeRequest request) {
         try {
             return new NodeStoreFilesMetadataBatch(clusterService.localNode(), listStoreMetadata(request));
         } catch (IOException e) {
             throw new OpenSearchException(
                 "Failed to list store metadata for shards [" + request.getShardAttributes().keySet().stream().map(ShardId::toString) + "]",
                 e
             );
         }
@@ -137,28 +140,37 @@
      * In this case we fetch the shard store files for batch of shards instead of one shard.
      */
     private Map<ShardId, NodeStoreFilesMetadata> listStoreMetadata(NodeRequest request) throws IOException {
         Map<ShardId, NodeStoreFilesMetadata> shardStoreMetadataMap = new HashMap<ShardId, NodeStoreFilesMetadata>();
-        for (ShardAttributes shardAttributes : request.getShardAttributes().values()) {
-            final ShardId shardId = shardAttributes.getShardId();
+        for (Map.Entry<ShardId, ShardAttributes> shardAttributes : request.getShardAttributes().entrySet()) {
+            final ShardId shardId = shardAttributes.getKey();
             try {
-                StoreFilesMetadata storeFilesMetadata = TransportNodesListShardStoreMetadataHelper.getListShardMetadataOnLocalNode(
+                StoreFilesMetadata storeFilesMetadata = TransportNodesListShardStoreMetadataHelper.listShardMetadataInternal(
                     logger,
                     shardId,
                     nodeEnv,
                     indicesService,
-                    shardAttributes.getCustomDataPath(),
+                    shardAttributes.getValue().getCustomDataPath(),
                     settings,
                     clusterService
                 );
                 shardStoreMetadataMap.put(shardId, new NodeStoreFilesMetadata(storeFilesMetadata, null));
-            } catch (IOException | OpenSearchException e) {
-                shardStoreMetadataMap.put(
-                    shardId,
-                    new NodeStoreFilesMetadata(new StoreFilesMetadata(shardId, Store.MetadataSnapshot.EMPTY, Collections.emptyList()), e)
-                );
+            } catch (Exception e) {
+                // should return null in case of known exceptions being returned from listShardMetadataInternal method.
+                if (e.getMessage().contains(INDEX_NOT_FOUND)) {
+                    shardStoreMetadataMap.put(shardId, null);
+                } else {
+                    // return actual exception as it is for unknown exceptions
+                    shardStoreMetadataMap.put(
+                        shardId,
+                        new NodeStoreFilesMetadata(
+                            new StoreFilesMetadata(shardId, Store.MetadataSnapshot.EMPTY, Collections.emptyList()),
+                            e
+                        )
+                    );
+                }
             }
         }
         return shardStoreMetadataMap;
     }
 
     /**
@@ -172,24 +184,24 @@
         private final Map<ShardId, ShardAttributes> shardAttributes;
 
         public Request(StreamInput in) throws IOException {
             super(in);
             shardAttributes = in.readMap(ShardId::new, ShardAttributes::new);
         }
 
         public Request(Map<ShardId, ShardAttributes> shardAttributes, DiscoveryNode[] nodes) {
             super(nodes);
             this.shardAttributes = Objects.requireNonNull(shardAttributes);
         }
 
         public Map<ShardId, ShardAttributes> getShardAttributes() {
             return shardAttributes;
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
             out.writeMap(shardAttributes, (o, k) -> k.writeTo(o), (o, v) -> v.writeTo(o));
         }
     }
 
     /**
@@ -200,26 +212,26 @@
     public static class NodesStoreFilesMetadataBatch extends BaseNodesResponse<NodeStoreFilesMetadataBatch> {
 
         public NodesStoreFilesMetadataBatch(StreamInput in) throws IOException {
             super(in);
         }
 
         public NodesStoreFilesMetadataBatch(
             ClusterName clusterName,
             List<NodeStoreFilesMetadataBatch> nodes,
             List<FailedNodeException> failures
         ) {
             super(clusterName, nodes, failures);
         }
 
         @Override
         protected List<NodeStoreFilesMetadataBatch> readNodesFrom(StreamInput in) throws IOException {
             return in.readList(NodeStoreFilesMetadataBatch::new);
         }
 
         @Override
         protected void writeNodesTo(StreamOutput out, List<NodeStoreFilesMetadataBatch> nodes) throws IOException {
             out.writeList(nodes);
         }
     }
 
     /**
@@ -227,55 +239,46 @@
      *
      * @opensearch.internal
      */
-    public static class NodeStoreFilesMetadata {
+    public static class NodeStoreFilesMetadata extends BaseShardResponse {
 
         private StoreFilesMetadata storeFilesMetadata;
-        private Exception storeFileFetchException;
 
-        public NodeStoreFilesMetadata(StoreFilesMetadata storeFilesMetadata) {
-            this.storeFilesMetadata = storeFilesMetadata;
-            this.storeFileFetchException = null;
+        @Override
+        public boolean isEmpty() {
+            return storeFilesMetadata == null || storeFilesMetadata().isEmpty() && getException() == null;
         }
 
         public NodeStoreFilesMetadata(StreamInput in) throws IOException {
-            storeFilesMetadata = new StoreFilesMetadata(in);
+            super(in);
             if (in.readBoolean()) {
-                this.storeFileFetchException = in.readException();
+                this.storeFilesMetadata = new StoreFilesMetadata(in);
             } else {
-                this.storeFileFetchException = null;
+                this.storeFilesMetadata = null;
             }
         }
 
         public NodeStoreFilesMetadata(StoreFilesMetadata storeFilesMetadata, Exception storeFileFetchException) {
+            super(storeFileFetchException);
             this.storeFilesMetadata = storeFilesMetadata;
-            this.storeFileFetchException = storeFileFetchException;
         }
 
         public StoreFilesMetadata storeFilesMetadata() {
             return storeFilesMetadata;
         }
 
-        public static NodeStoreFilesMetadata readListShardStoreNodeOperationResponse(StreamInput in) throws IOException {
-            return new NodeStoreFilesMetadata(in);
-        }
-
         public void writeTo(StreamOutput out) throws IOException {
-            storeFilesMetadata.writeTo(out);
-            if (storeFileFetchException != null) {
+            super.writeTo(out);
+            if (storeFilesMetadata != null) {
                 out.writeBoolean(true);
-                out.writeException(storeFileFetchException);
+                storeFilesMetadata.writeTo(out);
             } else {
                 out.writeBoolean(false);
             }
         }
 
-        public Exception getStoreFileFetchException() {
-            return storeFileFetchException;
-        }
-
         @Override
         public String toString() {
             return "[[" + storeFilesMetadata + "]]";
         }
     }
 
@@ -289,22 +292,22 @@
         private final Map<ShardId, ShardAttributes> shardAttributes;
 
         public NodeRequest(StreamInput in) throws IOException {
             super(in);
             shardAttributes = in.readMap(ShardId::new, ShardAttributes::new);
         }
 
         public NodeRequest(Request request) {
             this.shardAttributes = Objects.requireNonNull(request.getShardAttributes());
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
             out.writeMap(shardAttributes, (o, k) -> k.writeTo(o), (o, v) -> v.writeTo(o));
         }
 
         public Map<ShardId, ShardAttributes> getShardAttributes() {
             return shardAttributes;
         }
     }
 
@@ -316,24 +319,24 @@
         private final Map<ShardId, NodeStoreFilesMetadata> nodeStoreFilesMetadataBatch;
 
         protected NodeStoreFilesMetadataBatch(StreamInput in) throws IOException {
             super(in);
             this.nodeStoreFilesMetadataBatch = in.readMap(ShardId::new, NodeStoreFilesMetadata::new);
         }
 
         public NodeStoreFilesMetadataBatch(DiscoveryNode node, Map<ShardId, NodeStoreFilesMetadata> nodeStoreFilesMetadataBatch) {
             super(node);
             this.nodeStoreFilesMetadataBatch = nodeStoreFilesMetadataBatch;
         }
 
         public Map<ShardId, NodeStoreFilesMetadata> getNodeStoreFilesMetadataBatch() {
             return this.nodeStoreFilesMetadataBatch;
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
             out.writeMap(nodeStoreFilesMetadataBatch, (o, k) -> k.writeTo(o), (o, v) -> v.writeTo(o));
         }
     }
 
 }

server/src/main/java/org/opensearch/indices/store/TransportNodesListShardStoreMetadataHelper.java

@@ -47,8 +47,11 @@
  *
  * @opensearch.internal
  */
 public class TransportNodesListShardStoreMetadataHelper {
-    public static StoreFilesMetadata getListShardMetadataOnLocalNode(
+
+    public static final String INDEX_NOT_FOUND = "node doesn't have meta data for index ";
+
+    public static StoreFilesMetadata listShardMetadataInternal(
         Logger logger,
         final ShardId shardId,
         NodeEnvironment nodeEnv,
@@ -57,28 +60,28 @@
         Settings settings,
         ClusterService clusterService
     ) throws IOException {
         logger.trace("listing store meta data for {}", shardId);
         long startTimeNS = System.nanoTime();
         boolean exists = false;
         try {
             IndexService indexService = indicesService.indexService(shardId.getIndex());
             if (indexService != null) {
                 IndexShard indexShard = indexService.getShardOrNull(shardId.id());
                 if (indexShard != null) {
                     try {
                         final StoreFilesMetadata storeFilesMetadata = new StoreFilesMetadata(
                             shardId,
                             indexShard.snapshotStoreMetadata(),
                             indexShard.getPeerRecoveryRetentionLeases()
                         );
                         exists = true;
                         return storeFilesMetadata;
                     } catch (org.apache.lucene.index.IndexNotFoundException e) {
                         logger.trace(new ParameterizedMessage("[{}] node is missing index, responding with empty", shardId), e);
-                        throw e;
+                        return new StoreFilesMetadata(shardId, Store.MetadataSnapshot.EMPTY, Collections.emptyList());
                     } catch (IOException e) {
                         logger.warn(new ParameterizedMessage("[{}] can't read metadata from store, responding with empty", shardId), e);
-                        throw e;
+                        return new StoreFilesMetadata(shardId, Store.MetadataSnapshot.EMPTY, Collections.emptyList());
                     }
                 }
             }
@@ -86,40 +89,40 @@
                 // TODO: Fallback for BWC with older predecessor (ES) versions.
                 // Remove this once request.getCustomDataPath() always returns non-null
                 if (indexService != null) {
                     customDataPath = indexService.getIndexSettings().customDataPath();
                 } else {
                     IndexMetadata metadata = clusterService.state().metadata().index(shardId.getIndex());
                     if (metadata != null) {
                         customDataPath = new IndexSettings(metadata, settings).customDataPath();
                     } else {
                         logger.trace("{} node doesn't have meta data for the requests index", shardId);
-                        throw new OpenSearchException("node doesn't have meta data for index " + shardId.getIndex());
+                        throw new OpenSearchException(INDEX_NOT_FOUND + shardId.getIndex());
                     }
                 }
             }
             final ShardPath shardPath = ShardPath.loadShardPath(logger, nodeEnv, shardId, customDataPath);
             if (shardPath == null) {
                 return new StoreFilesMetadata(shardId, Store.MetadataSnapshot.EMPTY, Collections.emptyList());
             }
             // note that this may fail if it can't get access to the shard lock. Since we check above there is an active shard, this means:
             // 1) a shard is being constructed, which means the cluster-manager will not use a copy of this replica
             // 2) A shard is shutting down and has not cleared it's content within lock timeout. In this case the cluster-manager may not
             // reuse local resources.
             final Store.MetadataSnapshot metadataSnapshot = Store.readMetadataSnapshot(
                 shardPath.resolveIndex(),
                 shardId,
                 nodeEnv::shardLock,
                 logger
             );
             // We use peer recovery retention leases from the primary for allocating replicas. We should always have retention leases when
             // we refresh shard info after the primary has started. Hence, we can ignore retention leases if there is no active shard.
             return new StoreFilesMetadata(shardId, metadataSnapshot, Collections.emptyList());
         } finally {
             TimeValue took = new TimeValue(System.nanoTime() - startTimeNS, TimeUnit.NANOSECONDS);
             if (exists) {
                 logger.debug("{} loaded store meta data (took [{}])", shardId, took);
             } else {
                 logger.trace("{} didn't find any store meta data to load (took [{}])", shardId, took);
             }
         }
     }
@@ -158,7 +161,7 @@
         }
 
         public ShardId shardId() {
             return this.shardId;
         }
 
         public boolean isEmpty() {
@@ -204,13 +207,13 @@
 
         @Override
         public String toString() {
             return "StoreFilesMetadata{"
                 + ", shardId="
                 + shardId
                 + ", metadataSnapshot{size="
                 + metadataSnapshot.size()
                 + ", syncId="
                 + metadataSnapshot.getSyncId()
                 + "}"
                 + '}';
         }