Add support to restore only unassigned shards of an index

Signed-off-by: Sachin Kale <[email protected]>
opensearch-project · Jul 20, 2023 · 446187d · 446187d
1 parent 2ba1157
commit 446187d
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 12 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/routing/IndexRoutingTable.java b/server/src/main/java/org/opensearch/cluster/routing/IndexRoutingTable.java
@@ -450,7 +450,7 @@ public Builder initializeAsRestore(IndexMetadata indexMetadata, SnapshotRecovery
         /**
          * Initializes an existing index, to be restored from remote store
          */
-        public Builder initializeAsRemoteStoreRestore(IndexMetadata indexMetadata, RemoteStoreRecoverySource recoverySource) {
+        public Builder initializeAsRemoteStoreRestore(IndexMetadata indexMetadata, RemoteStoreRecoverySource recoverySource, Map<ShardId, ShardRouting> activeShards) {
             final UnassignedInfo unassignedInfo = new UnassignedInfo(
                 UnassignedInfo.Reason.EXISTING_INDEX_RESTORED,
                 "restore_source[remote_store]"
@@ -462,7 +462,11 @@ public Builder initializeAsRemoteStoreRestore(IndexMetadata indexMetadata, Remot
             for (int shardNumber = 0; shardNumber < indexMetadata.getNumberOfShards(); shardNumber++) {
                 ShardId shardId = new ShardId(index, shardNumber);
                 IndexShardRoutingTable.Builder indexShardRoutingBuilder = new IndexShardRoutingTable.Builder(shardId);
-                indexShardRoutingBuilder.addShard(ShardRouting.newUnassigned(shardId, true, recoverySource, unassignedInfo));
+                if (activeShards.containsKey(shardId)) {
+                    indexShardRoutingBuilder.addShard(activeShards.get(shardId));
+                } else {
+                    indexShardRoutingBuilder.addShard(ShardRouting.newUnassigned(shardId, true, recoverySource, unassignedInfo));
+                }
                 shards.put(shardNumber, indexShardRoutingBuilder.build());
             }
             return this;

diff --git a/server/src/main/java/org/opensearch/cluster/routing/RoutingTable.java b/server/src/main/java/org/opensearch/cluster/routing/RoutingTable.java
@@ -562,9 +562,9 @@ public Builder addAsFromOpenToClose(IndexMetadata indexMetadata) {
             return add(indexRoutingBuilder);
         }
 
-        public Builder addAsRemoteStoreRestore(IndexMetadata indexMetadata, RemoteStoreRecoverySource recoverySource) {
+        public Builder addAsRemoteStoreRestore(IndexMetadata indexMetadata, RemoteStoreRecoverySource recoverySource, Map<ShardId, ShardRouting> activeShards) {
             IndexRoutingTable.Builder indexRoutingBuilder = new IndexRoutingTable.Builder(indexMetadata.getIndex())
-                .initializeAsRemoteStoreRestore(indexMetadata, recoverySource);
+                .initializeAsRemoteStoreRestore(indexMetadata, recoverySource, activeShards);
             add(indexRoutingBuilder);
             return this;
         }

diff --git a/server/src/main/java/org/opensearch/snapshots/RestoreService.java b/server/src/main/java/org/opensearch/snapshots/RestoreService.java
@@ -234,34 +234,37 @@ public ClusterState execute(ClusterState currentState) {
                         continue;
                     }
                     if (currentIndexMetadata.getSettings().getAsBoolean(SETTING_REMOTE_STORE_ENABLED, false)) {
-                        if (currentIndexMetadata.getState() != IndexMetadata.State.CLOSE) {
+                        /*if (currentIndexMetadata.getState() != IndexMetadata.State.CLOSE) {
                             throw new IllegalStateException(
                                 "cannot restore index ["
                                     + index
                                     + "] because an open index "
                                     + "with same name already exists in the cluster. Close the existing index"
                             );
-                        }
+                        }*/
+                        /*
                         IndexMetadata updatedIndexMetadata = IndexMetadata.builder(currentIndexMetadata)
                             .state(IndexMetadata.State.OPEN)
                             .version(1 + currentIndexMetadata.getVersion())
                             .mappingVersion(1 + currentIndexMetadata.getMappingVersion())
                             .settingsVersion(1 + currentIndexMetadata.getSettingsVersion())
                             .aliasesVersion(1 + currentIndexMetadata.getAliasesVersion())
                             .build();
+                         */
 
-                        IndexId indexId = new IndexId(index, updatedIndexMetadata.getIndexUUID());
+                        IndexId indexId = new IndexId(index, currentIndexMetadata.getIndexUUID());
 
                         RemoteStoreRecoverySource recoverySource = new RemoteStoreRecoverySource(
                             restoreUUID,
-                            updatedIndexMetadata.getCreationVersion(),
+                            currentIndexMetadata.getCreationVersion(),
                             indexId
                         );
-                        rtBuilder.addAsRemoteStoreRestore(updatedIndexMetadata, recoverySource);
-                        blocks.updateBlocks(updatedIndexMetadata);
-                        mdBuilder.put(updatedIndexMetadata, true);
+                        Map<ShardId, ShardRouting> activeShards = currentState.routingTable().index(index).randomAllActiveShardsIt().getShardRoutings().stream().collect(Collectors.toMap(ShardRouting::shardId, Function.identity()));
+                        rtBuilder.addAsRemoteStoreRestore(currentIndexMetadata, recoverySource, activeShards);
+                        blocks.updateBlocks(currentIndexMetadata);
+                        mdBuilder.put(currentIndexMetadata, true);
                         indicesToBeRestored.add(index);
-                        totalShards += updatedIndexMetadata.getNumberOfShards();
+                        totalShards += currentIndexMetadata.getNumberOfShards();
                     } else {
                         logger.warn("Remote store is not enabled for index: {}", index);
                     }