improve node counting a bit

Author: ywangd

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

@@ -1231,7 +1231,7 @@ public Entry withClones(Map<RepositoryShardId, ShardSnapshotStatus> updatedClone
          * @return aborted snapshot entry or {@code null} if entry can be removed from the cluster state directly
          */
         @Nullable
-        public Entry abort(BiConsumer<ShardId, ShardSnapshotStatus> abortedAssignedQueuedShardConsumer) {
+        public Entry abort(String localNodeId, BiConsumer<ShardId, ShardSnapshotStatus> abortedAssignedQueuedShardConsumer) {
             final Map<ShardId, ShardSnapshotStatus> shardsBuilder = new HashMap<>();
             boolean completed = true;
             boolean allQueued = true;
@@ -1262,7 +1262,7 @@ assert isClone() == false
                         // Accumulate the updates needed to complete the aborted QUEUED with generation shard snapshots
                         abortedAssignedQueuedShardConsumer.accept(
                             shardEntry.getKey(),
-                            new ShardSnapshotStatus(null, ShardState.FAILED, status.generation, reason)
+                            new ShardSnapshotStatus(localNodeId, ShardState.FAILED, status.generation, reason)
                         );
                     }
                 }

server/src/main/java/org/elasticsearch/common/settings/ClusterSettings.java

@@ -578,6 +578,7 @@ public void apply(Settings value, Settings current, Settings previous) {
         HandshakingTransportAddressConnector.PROBE_CONNECT_TIMEOUT_SETTING,
         HandshakingTransportAddressConnector.PROBE_HANDSHAKE_TIMEOUT_SETTING,
         SnapshotsService.MAX_CONCURRENT_SNAPSHOT_OPERATIONS_SETTING,
+        SnapshotsService.SHARD_SNAPSHOT_PER_NODE_LIMIT_SETTING,
         RestoreService.REFRESH_REPO_UUID_ON_RESTORE_SETTING,
         FsHealthService.ENABLED_SETTING,
         FsHealthService.REFRESH_INTERVAL_SETTING,

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

@@ -44,8 +44,7 @@ public PerNodeShardSnapshotCounter(
                     return;
                 }
                 for (var shardSnapshotStatus : entry.shards().values()) {
-                    // TODO: consider more states as active, e.g. Abort on data node?
-                    if (shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.INIT) {
+                    if (isRunningOnDataNode(shardSnapshotStatus)) {
                         perNodeCounts.computeIfPresent(shardSnapshotStatus.nodeId(), (nodeId, count) -> count + 1);
                     }
                 }
@@ -54,6 +53,13 @@ public PerNodeShardSnapshotCounter(
         }
     }
 
+    private static boolean isRunningOnDataNode(SnapshotsInProgress.ShardSnapshotStatus shardSnapshotStatus) {
+        return shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.INIT
+            || (shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.ABORTED && shardSnapshotStatus.reason() != null
+            // TODO: find a better way to check for abort on data nodes
+                && shardSnapshotStatus.reason().startsWith("assigned-queued aborted") == false);
+    }
+
     public boolean tryStartShardSnapshotOnNode(String nodeId) {
         if (enabled() == false) {
             return true;

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

@@ -2305,6 +2305,7 @@ public ClusterState execute(ClusterState currentState) {
                             && snapshotIdsRequiringCleanup.contains(existing.snapshot().getSnapshotId())) {
                             // snapshot is started - mark every non completed shard as aborted
                             final SnapshotsInProgress.Entry abortedEntry = existing.abort(
+                                currentState.nodes().getLocalNodeId(),
                                 ((shardId, shardSnapshotStatus) -> completeAbortedAssignedQueuedRunnables.add(
                                     () -> innerUpdateSnapshotState(
                                         existing.snapshot(),
@@ -3820,9 +3821,7 @@ private <T> void applyShardSnapshotUpdate(
                     assert updatedShardSnapshotStatus.isActive() == false : updatedShardSnapshotStatus;
                 }
 
-                if (shardSnapshotStatusUpdate.isClone() == false
-                    && existing.state() == ShardState.INIT // TODO: this is not right, should consider aborted ones on data nodes as well
-                    && updatedShardSnapshotStatus.state() != ShardState.INIT) {
+                if (shardSnapshotStatusUpdate.isClone() == false && changeReleasesDataNode(existing, updatedShardSnapshotStatus)) {
                     perNodeShardSnapshotCounter.completeShardSnapshotOnNode(updatedShardSnapshotStatus.nodeId());
                 }
                 logger.trace(
@@ -3836,6 +3835,18 @@ private <T> void applyShardSnapshotUpdate(
                 executedUpdates.add(shardSnapshotStatusUpdate);
             }
 
+            private boolean changeReleasesDataNode(ShardSnapshotStatus previous, ShardSnapshotStatus current) {
+                if (previous.state() == ShardState.INIT) {
+                    return current.state().completed() || current.state() == ShardState.PAUSED_FOR_NODE_REMOVAL;
+                }
+                // TODO find a better to check for abort on data nodes
+                if (previous.state() == ShardState.ABORTED
+                    && (previous.reason() == null || previous.reason().startsWith("assigned-queued aborted") == false)) {
+                    return current.state().completed();
+                }
+                return false;
+            }
+
             private void tryStartNextTaskAfterCloneUpdated(RepositoryShardId repoShardId, ShardSnapshotStatus updatedState) {
                 // the update was already executed on the clone operation it applied to, now we check if it may be possible to
                 // start a shard snapshot or clone operation on the current entry