Refactor and reuse EntryContext for propagation

Author: ywangd

server/src/main/java/org/elasticsearch/cluster/SnapshotsInProgress.java

@@ -59,7 +59,6 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -1272,17 +1271,18 @@ public Entry withClones(Map<RepositoryShardId, ShardSnapshotStatus> updatedClone
          * In the special case where this instance has not yet made any progress on any shard this method just returns
          * {@code null} since no abort is needed and the snapshot can simply be removed from the cluster state outright.
          *
+         * @param completedAssignedQueuedShards Map to accumulate assigned-queued shards that get aborted in this entry
          * @return aborted snapshot entry or {@code null} if entry can be removed from the cluster state directly
          */
         @Nullable
-        public Entry abort(String localNodeId, BiConsumer<ShardId, ShardSnapshotStatus> abortedAssignedQueuedShardConsumer) {
+        public Entry abort(Map<ShardId, Tuple<Snapshot, ShardSnapshotStatus>> completedAssignedQueuedShards) {
             final Map<ShardId, ShardSnapshotStatus> shardsBuilder = new HashMap<>();
             boolean completed = true;
             boolean allQueued = true;
             for (Map.Entry<ShardId, ShardSnapshotStatus> shardEntry : shards.entrySet()) {
                 ShardSnapshotStatus status = shardEntry.getValue();
                 final var isAssignedQueued = status.isAssignedQueued();
-                allQueued &= (status.state() == ShardState.QUEUED && isAssignedQueued == false);
+                allQueued &= status.state() == ShardState.QUEUED;
                 if (status.state().completed() == false) {
                     final String nodeId = status.nodeId();
                     if (isAssignedQueued == false) {
@@ -1293,21 +1293,14 @@ public Entry abort(String localNodeId, BiConsumer<ShardId, ShardSnapshotStatus>
                             "aborted by snapshot deletion"
                         );
                     } else {
+                        // Record the deleted assigned-queued shard snapshot so that we can reassign to the first QUEUED one
+                        // in later snapshots.
+                        final var old = completedAssignedQueuedShards.put(shardEntry.getKey(), new Tuple<>(snapshot(), status));
+                        assert old == null : old;
                         assert isClone() == false
                             : "The state queued with generation should not be possible for a clone entry [" + this + "]";
                         final String reason = "assigned-queued aborted by snapshot deletion";
-                        // Assigned QUEUED transitions to ABORTED (incomplete) and is completed by a separate cluster state update
-                        status = new ShardSnapshotStatus(
-                            null, // use null nodeId to differentiate it from aborted INIT shard snapshot
-                            ShardState.ABORTED,
-                            status.generation(),
-                            reason
-                        );
-                        // Accumulate the updates needed to complete the aborted assigned-queued shard snapshots
-                        abortedAssignedQueuedShardConsumer.accept(
-                            shardEntry.getKey(),
-                            new ShardSnapshotStatus(localNodeId, ShardState.FAILED, status.generation, reason)
-                        );
+                        status = new ShardSnapshotStatus(nodeId, ShardState.FAILED, status.generation(), reason);
                     }
                 }
                 completed &= status.state().completed();

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

@@ -104,8 +104,7 @@ private boolean enabled() {
 
     private static boolean isRunningOnDataNode(SnapshotsInProgress.ShardSnapshotStatus shardSnapshotStatus) {
         return shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.INIT
-            // Aborted shard snapshot may still be running on the data node unless it was assigned-queued, i.e. never actually started
-            || (shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.ABORTED
-                && shardSnapshotStatus.isAbortedAssignedQueued() == false);
+            // Aborted shard snapshot is still running on the data node until it is FAILED
+            || shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.ABORTED;
     }
 }