Add ability to restore partial snapshots

Closes #5742

docs/reference/modules/snapshots.asciidoc

@@ -207,6 +207,16 @@ didn't exist in the cluster. If cluster state is restored, the restored template
 cluster are added and existing templates with the same name are replaced by the restored templates. The restored
 persistent settings are added to the existing persistent settings.
 
+[float]
+=== Partial restore
+
+added[1.3.0]
+
+By default, entire restore operation will fail if one or more indices participating in the operation don't have
+snapshots of all shards available. It can occur if some shards failed to snapshot for example. It is still possible to
+restore such indices by setting `partial` to `true`. Please note, that only successfully snapshotted shards will be
+restored in this case and all missing shards will be recreated empty.
+
 
 [float]
 === Snapshot status

src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/RestoreSnapshotRequest.java

@@ -21,6 +21,7 @@
 
 import org.elasticsearch.ElasticsearchGenerationException;
 import org.elasticsearch.ElasticsearchIllegalArgumentException;
+import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.action.support.master.MasterNodeOperationRequest;
@@ -67,6 +68,8 @@ public class RestoreSnapshotRequest extends MasterNodeOperationRequest<RestoreSn
 
     private boolean includeGlobalState = true;
 
+    private boolean partial = false;
+
     private Settings settings = EMPTY_SETTINGS;
 
     RestoreSnapshotRequest() {
@@ -272,6 +275,26 @@ public boolean waitForCompletion() {
         return waitForCompletion;
     }
 
+    /**
+     * Returns true if indices with failed to snapshot shards should be partially restored.
+     *
+     * @return true if indices with failed to snapshot shards should be partially restored
+     */
+    public boolean partial() {
+        return partial;
+    }
+
+    /**
+     * Set to true to allow indices with failed to snapshot shards should be partially restored.
+     *
+     * @param partial true if indices with failed to snapshot shards should be partially restored.
+     * @return this request
+     */
+    public RestoreSnapshotRequest partial(boolean partial) {
+        this.partial = partial;
+        return this;
+    }
+
     /**
      * Sets repository-specific restore settings.
      * <p/>
@@ -405,6 +428,8 @@ public RestoreSnapshotRequest source(Map source) {
                 expandWildcardsOpen = nodeBooleanValue(entry.getValue());
             } else if (name.equals("expand_wildcards_closed") || name.equals("expandWildcardsClosed")) {
                 expandWildcardsClosed = nodeBooleanValue(entry.getValue());
+            } else if (name.equals("partial")) {
+                partial(nodeBooleanValue(entry.getValue()));
             } else if (name.equals("settings")) {
                 if (!(entry.getValue() instanceof Map)) {
                     throw new ElasticsearchIllegalArgumentException("malformed settings section, should indices an inner object");
@@ -511,6 +536,9 @@ public void readFrom(StreamInput in) throws IOException {
         renameReplacement = in.readOptionalString();
         waitForCompletion = in.readBoolean();
         includeGlobalState = in.readBoolean();
+        if (in.getVersion().onOrAfter(Version.V_1_3_0)) {
+            partial = in.readBoolean();
+        }
         settings = readSettingsFromStream(in);
     }
 
@@ -525,6 +553,9 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(renameReplacement);
         out.writeBoolean(waitForCompletion);
         out.writeBoolean(includeGlobalState);
+        if (out.getVersion().onOrAfter(Version.V_1_3_0)) {
+            out.writeBoolean(partial);
+        }
         writeSettingsToStream(settings, out);
     }
 }

src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/RestoreSnapshotRequestBuilder.java

@@ -208,6 +208,17 @@ public RestoreSnapshotRequestBuilder setRestoreGlobalState(boolean restoreGlobal
         return this;
     }
 
+    /**
+     * If set to true the restore procedure will restore partially snapshotted indices
+     *
+     * @param partial true if partially snapshotted indices should be restored
+     * @return this request
+     */
+    public RestoreSnapshotRequestBuilder setPartial(boolean partial) {
+        request.partial(partial);
+        return this;
+    }
+
     @Override
     protected void doExecute(ActionListener<RestoreSnapshotResponse> listener) {
         client.restoreSnapshot(request, listener);

src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/TransportRestoreSnapshotAction.java

@@ -75,15 +75,10 @@ protected ClusterBlockException checkBlock(RestoreSnapshotRequest request, Clust
 
     @Override
     protected void masterOperation(final RestoreSnapshotRequest request, ClusterState state, final ActionListener<RestoreSnapshotResponse> listener) throws ElasticsearchException {
-        RestoreService.RestoreRequest restoreRequest =
-                new RestoreService.RestoreRequest("restore_snapshot[" + request.snapshot() + "]", request.repository(), request.snapshot())
-                        .indices(request.indices())
-                        .indicesOptions(request.indicesOptions())
-                        .renamePattern(request.renamePattern())
-                        .renameReplacement(request.renameReplacement())
-                        .includeGlobalState(request.includeGlobalState())
-                        .settings(request.settings())
-                        .masterNodeTimeout(request.masterNodeTimeout());
+        RestoreService.RestoreRequest restoreRequest = new RestoreService.RestoreRequest(
+                "restore_snapshot[" + request.snapshot() + "]", request.repository(), request.snapshot(),
+                request.indices(), request.indicesOptions(), request.renamePattern(), request.renameReplacement(),
+                request.settings(), request.masterNodeTimeout(), request.includeGlobalState(), request.partial());
         restoreService.restoreSnapshot(restoreRequest, new RestoreSnapshotListener() {
             @Override
             public void onResponse(RestoreInfo restoreInfo) {

src/main/java/org/elasticsearch/cluster/routing/IndexRoutingTable.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.cluster.routing;
 
+import com.carrotsearch.hppc.IntSet;
 import com.carrotsearch.hppc.cursors.IntCursor;
 import com.carrotsearch.hppc.cursors.IntObjectCursor;
 import com.google.common.collect.ImmutableList;
@@ -381,28 +382,33 @@ public Builder initializeAsRecovery(IndexMetaData indexMetaData) {
         /**
          * Initializes a new empty index, to be restored from a snapshot
          */
-        public Builder initializeAsNewRestore(IndexMetaData indexMetaData, RestoreSource restoreSource) {
-            return initializeAsRestore(indexMetaData, restoreSource, true);
+        public Builder initializeAsNewRestore(IndexMetaData indexMetaData, RestoreSource restoreSource, IntSet ignoreShards) {
+            return initializeAsRestore(indexMetaData, restoreSource, ignoreShards, true);
         }
 
         /**
          * Initializes an existing index, to be restored from a snapshot
          */
         public Builder initializeAsRestore(IndexMetaData indexMetaData, RestoreSource restoreSource) {
-            return initializeAsRestore(indexMetaData, restoreSource, false);
+            return initializeAsRestore(indexMetaData, restoreSource, null, false);
         }
 
         /**
          * Initializes an index, to be restored from snapshot
          */
-        private Builder initializeAsRestore(IndexMetaData indexMetaData, RestoreSource restoreSource, boolean asNew) {
+        private Builder initializeAsRestore(IndexMetaData indexMetaData, RestoreSource restoreSource, IntSet ignoreShards, boolean asNew) {
             if (!shards.isEmpty()) {
                 throw new ElasticsearchIllegalStateException("trying to initialize an index with fresh shards, but already has shards created");
             }
             for (int shardId = 0; shardId < indexMetaData.numberOfShards(); shardId++) {
                 IndexShardRoutingTable.Builder indexShardRoutingBuilder = new IndexShardRoutingTable.Builder(new ShardId(indexMetaData.index(), shardId), asNew ? false : true);
                 for (int i = 0; i <= indexMetaData.numberOfReplicas(); i++) {
-                    indexShardRoutingBuilder.addShard(new ImmutableShardRouting(index, shardId, null, null, i == 0 ? restoreSource : null, i == 0, ShardRoutingState.UNASSIGNED, 0));
+                    if (asNew && ignoreShards.contains(shardId)) {
+                        // This shards wasn't completely snapshotted - restore it as new shard
+                        indexShardRoutingBuilder.addShard(new ImmutableShardRouting(index, shardId, null, i == 0, ShardRoutingState.UNASSIGNED, 0));
+                    } else {
+                        indexShardRoutingBuilder.addShard(new ImmutableShardRouting(index, shardId, null, null, i == 0 ? restoreSource : null, i == 0, ShardRoutingState.UNASSIGNED, 0));
+                    }
                 }
                 shards.put(shardId, indexShardRoutingBuilder.build());
             }

src/main/java/org/elasticsearch/cluster/routing/RoutingTable.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.cluster.routing;
 
+import com.carrotsearch.hppc.IntSet;
 import com.google.common.collect.*;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.IndexMetaData;
@@ -417,9 +418,9 @@ public Builder addAsRestore(IndexMetaData indexMetaData, RestoreSource restoreSo
             return this;
         }
 
-        public Builder addAsNewRestore(IndexMetaData indexMetaData, RestoreSource restoreSource) {
+        public Builder addAsNewRestore(IndexMetaData indexMetaData, RestoreSource restoreSource, IntSet ignoreShards) {
             IndexRoutingTable.Builder indexRoutingBuilder = new IndexRoutingTable.Builder(indexMetaData.index())
-                    .initializeAsNewRestore(indexMetaData, restoreSource);
+                    .initializeAsNewRestore(indexMetaData, restoreSource, ignoreShards);
             add(indexRoutingBuilder);
             return this;
         }