SNAPSHOT: Keep SnapshotsInProgress State in Sync with Routing Table

server/src/main/java/org/elasticsearch/cluster/routing/RoutingChangesObserver.java

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.cluster.routing;
 
+import org.elasticsearch.index.shard.ShardId;
+
 /**
  * Records changes made to {@link RoutingNodes} during an allocation round.
  */
@@ -211,4 +213,51 @@ public void initializedReplicaReinitialized(ShardRouting oldReplica, ShardRoutin
             }
         }
     }
+
+    abstract class AbstractChangedShardObserver extends AbstractRoutingChangesObserver {
+
+        @Override
+        public void shardInitialized(ShardRouting unassignedShard, ShardRouting initializedShard) {
+            onChanged(unassignedShard.shardId());
+        }
+
+        @Override
+        public void shardStarted(ShardRouting initializingShard, ShardRouting startedShard) {
+            onChanged(initializingShard.shardId());
+        }
+        @Override
+        public void relocationStarted(ShardRouting startedShard, ShardRouting targetRelocatingShard) {
+            onChanged(startedShard.shardId());
+        }
+        @Override
+        public void unassignedInfoUpdated(ShardRouting unassignedShard, UnassignedInfo newUnassignedInfo) {
+            onChanged(unassignedShard.shardId());
+        }
+        @Override
+        public void shardFailed(ShardRouting failedShard, UnassignedInfo unassignedInfo) {
+            onChanged(failedShard.shardId());
+        }
+        @Override
+        public void relocationCompleted(ShardRouting removedRelocationSource) {
+            onChanged(removedRelocationSource.shardId());
+        }
+        @Override
+        public void relocationSourceRemoved(ShardRouting removedReplicaRelocationSource) {
+            onChanged(removedReplicaRelocationSource.shardId());
+        }
+        @Override
+        public void startedPrimaryReinitialized(ShardRouting startedPrimaryShard, ShardRouting initializedShard) {
+            onChanged(startedPrimaryShard.shardId());
+        }
+        @Override
+        public void replicaPromoted(ShardRouting replicaShard) {
+            onChanged(replicaShard.shardId());
+        }
+        @Override
+        public void initializedReplicaReinitialized(ShardRouting oldReplica, ShardRouting reinitializedReplica) {
+            onChanged(oldReplica.shardId());
+        }
+
+        protected abstract void onChanged(ShardId shardId);
+    }
 }

server/src/main/java/org/elasticsearch/cluster/routing/allocation/AllocationService.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.cluster.ClusterInfoService;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.RestoreInProgress;
+import org.elasticsearch.cluster.SnapshotsInProgress;
 import org.elasticsearch.cluster.health.ClusterHealthStatus;
 import org.elasticsearch.cluster.health.ClusterStateHealth;
 import org.elasticsearch.cluster.metadata.AutoExpandReplicas;
@@ -136,15 +137,29 @@ private ClusterState buildResult(ClusterState oldState, RoutingAllocation alloca
         final ClusterState.Builder newStateBuilder = ClusterState.builder(oldState)
             .routingTable(newRoutingTable)
             .metaData(newMetaData);
+        ImmutableOpenMap.Builder<String, ClusterState.Custom> customsBuilder = null;
         final RestoreInProgress restoreInProgress = allocation.custom(RestoreInProgress.TYPE);
         if (restoreInProgress != null) {
             RestoreInProgress updatedRestoreInProgress = allocation.updateRestoreInfoWithRoutingChanges(restoreInProgress);
             if (updatedRestoreInProgress != restoreInProgress) {
-                ImmutableOpenMap.Builder<String, ClusterState.Custom> customsBuilder = ImmutableOpenMap.builder(allocation.getCustoms());
+                customsBuilder = ImmutableOpenMap.builder(allocation.getCustoms());
                 customsBuilder.put(RestoreInProgress.TYPE, updatedRestoreInProgress);
-                newStateBuilder.customs(customsBuilder.build());
             }
         }
+        final SnapshotsInProgress snapshotsInProgress = allocation.custom(SnapshotsInProgress.TYPE);
+        if (snapshotsInProgress != null) {
+            SnapshotsInProgress updatedSnapshotsInProgress =
+                allocation.updateSnapshotsWithRoutingChanges(snapshotsInProgress, newRoutingTable);
+            if (updatedSnapshotsInProgress != snapshotsInProgress) {
+                if (customsBuilder == null) {
+                    customsBuilder = ImmutableOpenMap.builder(allocation.getCustoms());
+                }
+                customsBuilder.put(SnapshotsInProgress.TYPE, updatedSnapshotsInProgress);
+            }
+        }
+        if (customsBuilder != null) {
+            newStateBuilder.customs(customsBuilder.build());
+        }
         return newStateBuilder.build();
     }
 

server/src/main/java/org/elasticsearch/cluster/routing/allocation/RoutingAllocation.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.cluster.ClusterInfo;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.RestoreInProgress;
+import org.elasticsearch.cluster.SnapshotsInProgress;
 import org.elasticsearch.cluster.metadata.MetaData;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.routing.RoutingChangesObserver;
@@ -38,6 +39,7 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import org.elasticsearch.snapshots.SnapshotsService.SnapshotsInProgressUpdater;
 
 import static java.util.Collections.emptySet;
 import static java.util.Collections.unmodifiableSet;
@@ -76,11 +78,11 @@ public class RoutingAllocation {
     private final IndexMetaDataUpdater indexMetaDataUpdater = new IndexMetaDataUpdater();
     private final RoutingNodesChangedObserver nodesChangedObserver = new RoutingNodesChangedObserver();
     private final RestoreInProgressUpdater restoreInProgressUpdater = new RestoreInProgressUpdater();
+    private final SnapshotsInProgressUpdater snapshotsInProgressUpdater = new SnapshotsInProgressUpdater();
     private final RoutingChangesObserver routingChangesObserver = new RoutingChangesObserver.DelegatingRoutingChangesObserver(
-        nodesChangedObserver, indexMetaDataUpdater, restoreInProgressUpdater
+        nodesChangedObserver, indexMetaDataUpdater, restoreInProgressUpdater, snapshotsInProgressUpdater
     );
 
-
     /**
      * Creates a new {@link RoutingAllocation}
      *  @param deciders {@link AllocationDeciders} to used to make decisions for routing allocations
@@ -251,6 +253,14 @@ public RestoreInProgress updateRestoreInfoWithRoutingChanges(RestoreInProgress r
         return restoreInProgressUpdater.applyChanges(restoreInProgress);
     }
 
+    /**
+     * Returns updated {@link SnapshotsInProgress} based on the changes that were made to the routing nodes
+     */
+    public SnapshotsInProgress updateSnapshotsWithRoutingChanges(SnapshotsInProgress snapshotsInProgress,
+        RoutingTable newRoutingTable) {
+        return snapshotsInProgressUpdater.applyChanges(snapshotsInProgress, newRoutingTable);
+    }
+
     /**
      * Returns true iff changes were made to the routing nodes
      */

server/src/main/java/org/elasticsearch/snapshots/SnapshotShardsService.java

@@ -313,7 +313,7 @@ private void processIndexShardSnapshots(ClusterChangedEvent event) {
             shutdownLock.unlock();
         }
 
-        // We have new shards to starts
+        // We have new shards to start
         if (newSnapshots.isEmpty() == false) {
             Executor executor = threadPool.executor(ThreadPool.Names.SNAPSHOT);
             for (final Map.Entry<Snapshot, Map<ShardId, IndexShardSnapshotStatus>> entry : newSnapshots.entrySet()) {

server/src/main/java/org/elasticsearch/snapshots/SnapshotsService.java

@@ -21,6 +21,7 @@
 
 import com.carrotsearch.hppc.cursors.ObjectCursor;
 import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
+import java.util.stream.StreamSupport;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
@@ -45,6 +46,7 @@
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.routing.IndexRoutingTable;
 import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.RoutingChangesObserver;
 import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.service.ClusterService;
@@ -657,12 +659,28 @@ public void applyClusterState(ClusterChangedEvent event) {
                 }
                 removeFinishedSnapshotFromClusterState(event);
                 finalizeSnapshotDeletionFromPreviousMaster(event);
+                // TODO org.elasticsearch.snapshots.SharedClusterSnapshotRestoreIT.testDeleteOrphanSnapshot fails right after election here
+                assert event.previousState().nodes().isLocalNodeElectedMaster() || assertConsistency(event.state());
             }
         } catch (Exception e) {
             logger.warn("Failed to update snapshot state ", e);
         }
     }
 
+    private boolean assertConsistency(ClusterState state) {
+        SnapshotsInProgress snapshotsInProgress = state.custom(SnapshotsInProgress.TYPE);
+        if (snapshotsInProgress != null) {
+            assert snapshotsInProgress == updateWithRoutingTable(
+                snapshotsInProgress.entries().stream().flatMap(entry -> {
+                    Iterable<ShardId> iterable = () -> entry.shards().keysIt();
+                    return StreamSupport.stream(iterable.spliterator(), false);
+                }).collect(Collectors.toSet()),
+                snapshotsInProgress, state.routingTable()
+            ) : "SnapshotsInProgress state [" + snapshotsInProgress + "] not in sync with routing table [" + state.routingTable() + "].";
+        }
+        return true;
+    }
+
     /**
      * Finalizes a snapshot deletion in progress if the current node is the master but it
      * was not master in the previous cluster state and there is still a lingering snapshot
@@ -1539,6 +1557,128 @@ public interface SnapshotCompletionListener {
         void onSnapshotFailure(Snapshot snapshot, Exception e);
     }
 
+    public static final class SnapshotsInProgressUpdater extends RoutingChangesObserver.AbstractChangedShardObserver {
+
+        private final Set<ShardId> shardChanges = new HashSet<>();
+
+        public SnapshotsInProgress applyChanges(SnapshotsInProgress oldSnapshot, RoutingTable newRoutingTable) {
+            return updateWithRoutingTable(shardChanges, oldSnapshot, newRoutingTable);
+        }
+
+        @Override
+        protected void onChanged(ShardId shardId) {
+            shardChanges.add(shardId);
+        }
+    }
+
+    private static SnapshotsInProgress updateWithRoutingTable(Set<ShardId> shardIds, SnapshotsInProgress oldSnapshot,
+        RoutingTable newRoutingTable) {
+        if (oldSnapshot == null || shardIds.isEmpty()) {
+            return oldSnapshot;
+        }
+        List<SnapshotsInProgress.Entry> entries = new ArrayList<>();
+        boolean snapshotsInProgressChanged = false;
+        for (SnapshotsInProgress.Entry entry : oldSnapshot.entries()) {
+            ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> shardsBuilder = null;
+            for (ShardId shardId : shardIds) {
+                final ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards = entry.shards();
+                final ShardSnapshotStatus currentStatus = shards.get(shardId);
+                if (currentStatus != null && currentStatus.state().completed() == false) {
+                    IndexShardRoutingTable routingTable = newRoutingTable.shardRoutingTableOrNull(shardId);
+                    assert routingTable != null;
+                    final ShardSnapshotStatus newStatus = Optional.ofNullable(routingTable)
+                        .map(IndexShardRoutingTable::primaryShard)
+                        .map(
+                            primaryShardRouting -> determineShardSnapshotStatus(currentStatus, primaryShardRouting)
+                        )
+                        .orElse(failedStatus(null, "missing shard"));
+                    if (newStatus != currentStatus) {
+                        if (shardsBuilder == null) {
+                            shardsBuilder = ImmutableOpenMap.builder(shards);
+                        }
+                        shardsBuilder.put(shardId, newStatus);
+                    }
+                }
+            }
+            if (shardsBuilder == null) {
+                entries.add(entry);
+            } else {
+                snapshotsInProgressChanged = true;
+                ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards = shardsBuilder.build();
+                entries.add(
+                    new SnapshotsInProgress.Entry(
+                        entry,
+                        completed(shards.values()) ? State.SUCCESS : entry.state(),
+                        shards
+                    )
+                );
+            }
+        }
+        if (snapshotsInProgressChanged) {
+            return new SnapshotsInProgress(entries);
+        }
+        return oldSnapshot;
+    }
+
+    private static ShardSnapshotStatus determineShardSnapshotStatus(final ShardSnapshotStatus currentStatus,
+        final ShardRouting primaryShardRouting) {
+        final State currentState = currentStatus.state();
+        final ShardSnapshotStatus newStatus;
+        if (primaryShardRouting.active() == false) {
+            if (primaryShardRouting.initializing() && currentState == State.WAITING) {
+                newStatus = currentStatus;
+            } else {
+                newStatus = failedStatus(
+                    primaryShardRouting.currentNodeId(),
+                    primaryShardRouting.unassignedInfo().getReason().toString()
+                );
+            }
+        } else if (primaryShardRouting.started()) {
+            switch (currentState) {
+                case WAITING:
+                    newStatus = new ShardSnapshotStatus(primaryShardRouting.currentNodeId());
+                    break;
+                case INIT: {
+                    String currentNodeId = currentStatus.nodeId();
+                    assert currentNodeId != null;
+                    if (primaryShardRouting.currentNodeId().equals(currentNodeId)) {
+                        newStatus = currentStatus;
+                    } else {
+                        newStatus = failedStatus(currentNodeId);
+                    }
+                    break;
+                }
+                case ABORTED:
+                    String currentNodeId = currentStatus.nodeId();
+                    if (currentNodeId.equals(primaryShardRouting.currentNodeId())) {
+                        newStatus = currentStatus;
+                    } else {
+                        newStatus = failedStatus(currentNodeId);
+                    }
+                    break;
+                default:
+                    newStatus = currentStatus;
+                    break;
+            }
+        } else {
+            assert primaryShardRouting.relocating();
+            if (currentState == State.INIT || currentStatus.state() == State.ABORTED) {
+                newStatus = failedStatus(currentStatus.nodeId());
+            } else {
+                newStatus = currentStatus;
+            }
+        }
+        return newStatus;
+    }
+
+    private static ShardSnapshotStatus failedStatus(String nodeId) {
+        return failedStatus(nodeId, "shard failed");
+    }
+
+    private static ShardSnapshotStatus failedStatus(String nodeId, String reason) {
+        return new ShardSnapshotStatus(nodeId, State.FAILED, reason);
+    }
+
     /**
      * Snapshot creation request
      */

server/src/test/java/org/elasticsearch/snapshots/DedicatedClusterSnapshotRestoreIT.java

@@ -78,6 +78,8 @@
 import org.elasticsearch.test.ESIntegTestCase.Scope;
 import org.elasticsearch.test.InternalTestCluster;
 import org.elasticsearch.test.TestCustomMetaData;
+import org.elasticsearch.test.disruption.BusyMasterServiceDisruption;
+import org.elasticsearch.test.disruption.ServiceDisruptionScheme;
 import org.elasticsearch.test.rest.FakeRestRequest;
 
 import java.io.IOException;
@@ -1151,6 +1153,50 @@ public void testSnapshotTotalAndIncrementalSizes() throws IOException {
         assertThat(anotherStats.getTotalSize(), is(snapshot1FileSize));
     }
 
+    public void testDataNodeRestartWithBusyMasterDuringSnapshot() throws Exception {
+        logger.info("-->  starting a master node and two data nodes");
+        internalCluster().startMasterOnlyNode();
+        internalCluster().startDataOnlyNodes(2);
+        logger.info("-->  creating repository");
+        assertAcked(client().admin().cluster().preparePutRepository("test-repo")
+            .setType("mock").setSettings(Settings.builder()
+                .put("location", randomRepoPath())
+                .put("compress", randomBoolean())
+                .put("max_snapshot_bytes_per_sec", "1000b")
+                .put("chunk_size", randomIntBetween(100, 1000), ByteSizeUnit.BYTES)));
+        assertAcked(prepareCreate("test-idx", 0, Settings.builder()
+            .put("number_of_shards", 5).put("number_of_replicas", 0)));
+        ensureGreen();
+        logger.info("-->  indexing some data");
+        final int numdocs = randomIntBetween(50, 100);
+        IndexRequestBuilder[] builders = new IndexRequestBuilder[numdocs];
+        for (int i = 0; i < builders.length; i++) {
+            builders[i] = client().prepareIndex("test-idx", "type1",
+                Integer.toString(i)).setSource("field1", "bar " + i);
+        }
+        indexRandom(true, builders);
+        flushAndRefresh();
+        final String dataNode = blockNodeWithIndex("test-repo", "test-idx");
+        logger.info("-->  snapshot");
+        client(internalCluster().getMasterName()).admin().cluster()
+            .prepareCreateSnapshot("test-repo", "test-snap").setWaitForCompletion(false).setIndices("test-idx").get();
+        ServiceDisruptionScheme disruption = new BusyMasterServiceDisruption(random(), Priority.HIGH);
+        setDisruptionScheme(disruption);
+        disruption.startDisrupting();
+        logger.info("-->  restarting data node, which should cause primary shards to be failed");
+        internalCluster().restartNode(dataNode, InternalTestCluster.EMPTY_CALLBACK);
+        unblockNode("test-repo", dataNode);
+        disruption.stopDisrupting();
+        // check that snapshot completes
+        assertBusy(() -> {
+            GetSnapshotsResponse snapshotsStatusResponse = client().admin().cluster()
+                .prepareGetSnapshots("test-repo").setSnapshots("test-snap").setIgnoreUnavailable(true).get();
+            assertEquals(1, snapshotsStatusResponse.getSnapshots().size());
+            SnapshotInfo snapshotInfo = snapshotsStatusResponse.getSnapshots().get(0);
+            assertTrue(snapshotInfo.state().toString(), snapshotInfo.state().completed());
+        }, 30, TimeUnit.SECONDS);
+    }
+
     private long calculateTotalFilesSize(List<Path> files) {
         return files.stream().mapToLong(f -> {
             try {