Add assigned queued and limit INIT per node

With the limit enabled, shard snapshots that would have been in the INIT
state are now put into assigned QUEUED state when the limit is reached
on the node. This new state is transitioned into INIT when a shard
snapshot completes and gives back capacity to the node. The shard
snapshot that completes can be either from the same snapshot or a
different snapshot in either the same or different repo.

Author: ywangd

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

@@ -197,6 +197,10 @@ public int count() {
         return count;
     }
 
+    public Set<ProjectRepo> repos() {
+        return entries.keySet();
+    }
+
     public Iterable<List<Entry>> entriesByRepo() {
         return () -> Iterators.map(entries.values().iterator(), byRepo -> byRepo.entries);
     }
@@ -504,7 +508,7 @@ private static boolean assertShardStateConsistent(
         int shardId,
         ShardSnapshotStatus shardSnapshotStatus
     ) {
-        if (shardSnapshotStatus.isActive()) {
+        if (shardSnapshotStatus.isActiveOrAssignedQueued()) {
             Tuple<String, Integer> plainShardId = Tuple.tuple(indexName, shardId);
             assert assignedShards.add(plainShardId) : plainShardId + " is assigned twice in " + entries;
             assert queuedShards.contains(plainShardId) == false : plainShardId + " is queued then assigned in " + entries;
@@ -752,6 +756,19 @@ public static ShardSnapshotStatus success(String nodeId, ShardSnapshotResult sha
             return new ShardSnapshotStatus(nodeId, ShardState.SUCCESS, shardSnapshotResult.getGeneration(), null, shardSnapshotResult);
         }
 
+        @SuppressForbidden(reason = "using a private constructor within the same file")
+        public static ShardSnapshotStatus assignedQueued(String nodeId, ShardGeneration generation) {
+            return new ShardSnapshotStatus(nodeId, ShardState.QUEUED, generation, null, null);
+        }
+
+        public boolean isAssignedQueued() {
+            return state == ShardState.QUEUED && nodeId != null && generation != null;
+        }
+
+        public boolean isUnassignedQueued() {
+            return this == UNASSIGNED_QUEUED || (state == ShardState.QUEUED && nodeId == null && generation == null);
+        }
+
         public ShardSnapshotStatus(
             @Nullable String nodeId,
             ShardState state,
@@ -772,8 +789,16 @@ private boolean assertConsistent() {
             assert state.failed() == false || reason != null;
             assert (state != ShardState.INIT && state != ShardState.WAITING && state != ShardState.PAUSED_FOR_NODE_REMOVAL)
                 || nodeId != null : "Null node id for state [" + state + "]";
-            assert state != ShardState.QUEUED || (nodeId == null && generation == null && reason == null)
-                : "Found unexpected non-null values for queued state shard nodeId[" + nodeId + "][" + generation + "][" + reason + "]";
+            assert state != ShardState.QUEUED || (isUnassignedQueued() || isAssignedQueued())
+                : "Found unexpected shard state=["
+                    + state
+                    + "], nodeId=["
+                    + nodeId
+                    + "], generation=["
+                    + generation
+                    + "], reason=["
+                    + reason
+                    + "]";
             assert state == ShardState.SUCCESS || shardSnapshotResult == null;
             assert shardSnapshotResult == null || shardSnapshotResult.getGeneration().equals(generation)
                 : "generation [" + generation + "] does not match result generation [" + shardSnapshotResult.getGeneration() + "]";
@@ -787,7 +812,7 @@ public static ShardSnapshotStatus readFrom(StreamInput in) throws IOException {
             final ShardGeneration generation = in.readOptionalWriteable(ShardGeneration::new);
             final String reason = in.readOptionalString();
             final ShardSnapshotResult shardSnapshotResult = in.readOptionalWriteable(ShardSnapshotResult::new);
-            if (state == ShardState.QUEUED) {
+            if (state == ShardState.QUEUED && nodeId == null && generation == null) {
                 return UNASSIGNED_QUEUED;
             }
             return new ShardSnapshotStatus(nodeId, state, generation, reason, shardSnapshotResult);
@@ -819,13 +844,18 @@ public ShardSnapshotResult shardSnapshotResult() {
          * ({@link ShardState#INIT} or {@link ShardState#ABORTED}) or about to write to it in state {@link ShardState#WAITING} or
          * {@link ShardState#PAUSED_FOR_NODE_REMOVAL}.
          */
+        // TODO: review its usage again
         public boolean isActive() {
             return switch (state) {
                 case INIT, ABORTED, WAITING, PAUSED_FOR_NODE_REMOVAL -> true;
                 case SUCCESS, FAILED, MISSING, QUEUED -> false;
             };
         }
 
+        public boolean isActiveOrAssignedQueued() {
+            return isActive() || isAssignedQueued();
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeOptionalString(nodeId);
@@ -1210,7 +1240,7 @@ public Entry abort() {
                     final String nodeId = status.nodeId();
                     status = new ShardSnapshotStatus(
                         nodeId,
-                        nodeId == null ? ShardState.FAILED : ShardState.ABORTED,
+                        nodeId == null ? ShardState.FAILED : ShardState.ABORTED, // TODO: Aborted if assigned queued?
                         status.generation(),
                         "aborted by snapshot deletion"
                     );

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

@@ -67,7 +67,8 @@ private static void addStateInformation(
         int shardId,
         String indexName
     ) {
-        if (shardState.isActive()) {
+        // Both active or assigned queued means the shard is meant to be the one with actions if node capacity allows it
+        if (shardState.isActiveOrAssignedQueued()) {
             busyIds.computeIfAbsent(indexName, k -> new HashSet<>()).add(shardId);
             assert assertGenerationConsistency(generations, indexName, shardId, shardState.generation());
         } else if (shardState.state() == SnapshotsInProgress.ShardState.SUCCESS) {

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

@@ -0,0 +1,94 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.snapshots;
+
+import org.elasticsearch.cluster.SnapshotsInProgress;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.node.DiscoveryNodeRole;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class PerNodeShardSnapshotCounter {
+
+    private final int shardSnapshotPerNodeLimit;
+    private final Map<String, Integer> perNodeCounts;
+
+    public PerNodeShardSnapshotCounter(
+        int shardSnapshotPerNodeLimit,
+        SnapshotsInProgress snapshotsInProgress,
+        DiscoveryNodes nodes,
+        boolean isStateless
+    ) {
+        this.shardSnapshotPerNodeLimit = shardSnapshotPerNodeLimit;
+        if (this.shardSnapshotPerNodeLimit <= 0) {
+            this.perNodeCounts = Map.of();
+        } else {
+            final var perNodeCounts = nodes.getDataNodes()
+                .values()
+                .stream()
+                .filter(node -> isStateless == false || node.hasRole(DiscoveryNodeRole.INDEX_ROLE.roleName()))
+                .filter(node -> snapshotsInProgress.isNodeIdForRemoval(node.getId()) == false)
+                .collect(Collectors.toMap(DiscoveryNode::getId, node -> 0));
+
+            snapshotsInProgress.asStream().forEach(entry -> {
+                if (entry.state().completed() || entry.isClone()) {
+                    return;
+                }
+                for (var shardSnapshotStatus : entry.shards().values()) {
+                    if (shardSnapshotStatus.state() == SnapshotsInProgress.ShardState.INIT) {
+                        perNodeCounts.computeIfPresent(shardSnapshotStatus.nodeId(), (nodeId, count) -> count + 1);
+                    }
+                }
+            });
+            this.perNodeCounts = perNodeCounts;
+        }
+    }
+
+    public boolean tryStartShardSnapshotOnNode(String nodeId) {
+        if (enabled() == false) {
+            return true;
+        }
+        final Integer count = perNodeCounts.get(nodeId);
+        if (count == null) {
+            return false;
+        }
+        if (count < shardSnapshotPerNodeLimit) {
+            perNodeCounts.put(nodeId, count + 1);
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    public boolean completeShardSnapshotOnNode(String nodeId) {
+        if (enabled() == false) {
+            return true;
+        }
+        final Integer count = perNodeCounts.get(nodeId);
+        if (count == null) {
+            return false;
+        }
+        if (count <= 0) {
+            return false;
+        }
+        perNodeCounts.put(nodeId, count - 1);
+        return true;
+    }
+
+    public boolean hasCapacityOnAnyNode() {
+        return enabled() == false || perNodeCounts.values().stream().anyMatch(count -> count < shardSnapshotPerNodeLimit);
+    }
+
+    private boolean enabled() {
+        return shardSnapshotPerNodeLimit > 0;
+    }
+}