Add support for handling failed shards in the cache

Signed-off-by: Aman Khare <[email protected]>
amkhar · Feb 28, 2024 · 0daa63e · 0daa63e
1 parent a6921f0
commit 0daa63e
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/gateway/AsyncShardFetch.java b/server/src/main/java/org/opensearch/gateway/AsyncShardFetch.java
@@ -170,6 +170,7 @@ public synchronized FetchResult<T> fetchData(DiscoveryNodes nodes, Map<ShardId,
         }
 
         cache.fillShardCacheWithDataNodes(nodes);
+        List<ShardId> failedShards = clearFailedShards();
         List<String> nodeIds = cache.findNodesToFetch();
         if (nodeIds.isEmpty() == false) {
             // mark all node as fetching and go ahead and async fetch them
@@ -205,12 +206,25 @@ public synchronized FetchResult<T> fetchData(DiscoveryNodes nodes, Map<ShardId,
                         + allIgnoreNodesMap.values().stream().mapToInt(Set::size).sum()
                         + "]"
                 );
+            } else if (failedShards.isEmpty() == false) {
+                // trigger a reroute if there are any shards failed, to make sure they're picked up in next run
+                logger.trace("triggering another reroute for failed shards in {}", reroutingKey);
+                reroute("shards-failed", "shards failed in "+ reroutingKey);
             }
 
             return new FetchResult<>(fetchData, allIgnoreNodesMap);
         }
     }
 
+    private List<ShardId> clearFailedShards() {
+        // get failed shards from previous fetch and remove them
+        List<ShardId> failedShards = cache.getFailedShards();
+        if (failedShards !=null && failedShards.isEmpty() == false ) {
+            shardAttributesMap.keySet().removeIf(failedShards::contains);
+        }
+        return failedShards;
+    }
+
     /**
      * Called by the response handler of the async action to fetch data. Verifies that its still working
      * on the same cache generation, otherwise the results are discarded. It then goes and fills the relevant data for

diff --git a/server/src/main/java/org/opensearch/gateway/BaseShardCache.java b/server/src/main/java/org/opensearch/gateway/BaseShardCache.java
@@ -73,6 +73,13 @@ protected BaseShardCache(Logger logger, String logKey, String type) {
      */
     public abstract K getData(DiscoveryNode node);
 
+    /**
+     * Provide the list of shards which got failures, these shards should be removed
+     * @return list of failed shards
+     */
+    public abstract List<ShardId> getFailedShards();
+
+
     @NonNull
     public abstract Map<String, ? extends BaseNodeEntry> getCache();
 

diff --git a/server/src/main/java/org/opensearch/gateway/ShardCache.java b/server/src/main/java/org/opensearch/gateway/ShardCache.java
@@ -14,7 +14,9 @@
 import org.opensearch.common.Nullable;
 import org.opensearch.core.index.shard.ShardId;
 
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -55,6 +57,12 @@ public void clearShardCache(ShardId shardId) {
         cache.clear();
     }
 
+    @Override
+    public List<ShardId> getFailedShards() {
+        // Single shard cache does not need to return that shard itself because handleFailure will take care of retries
+        return Collections.emptyList();
+    }
+
     /**
      * A node entry, holding the state of the fetched data for a specific shard
      * for a giving node.