Pre-compute hasAssignedQueuedShards

Optimize when to kickoff starting assigned queued shards.
Fix allocation filter permit release

Author: ywangd

server/src/internalClusterTest/java/org/elasticsearch/snapshots/SnapshotStressTestsIT.java

@@ -1398,6 +1398,10 @@ private void pollForAllocationFilterCompletion(
             Releasable onCompletion,
             Runnable onSuccess
         ) {
+            if (shouldStop.get()) {
+                Releasables.close(onCompletion);
+                return;
+            }
             clientExecutor.execute(mustSucceed(() -> {
                 final var clusterService = internalCluster().getCurrentMasterNodeInstance(ClusterService.class);
                 final ClusterState state = clusterService.state();

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

@@ -913,6 +913,12 @@ public static class Entry implements Writeable, ToXContentObject, RepositoryOper
          */
         private final boolean hasShardsInInitState;
 
+        /**
+         * Flag set to true in case any of the shard snapshots in {@link #shards} are {@link ShardSnapshotStatus#isAssignedQueued}.
+         * This is used to avoid having to iterate the full {@link #shards} map.
+         */
+        private final boolean hasAssignedQueuedShards;
+
         // visible for testing, use #startedEntry and copy constructors in production code
         public static Entry snapshot(
             Snapshot snapshot,
@@ -932,6 +938,7 @@ public static Entry snapshot(
             final Map<String, Index> res = Maps.newMapWithExpectedSize(indices.size());
             final Map<RepositoryShardId, ShardSnapshotStatus> byRepoShardIdBuilder = Maps.newHashMapWithExpectedSize(shards.size());
             boolean hasInitStateShards = false;
+            boolean hasAssignedQueuedShards = false;
             for (Map.Entry<ShardId, ShardSnapshotStatus> entry : shards.entrySet()) {
                 final ShardId shardId = entry.getKey();
                 final IndexId indexId = indices.get(shardId.getIndexName());
@@ -940,6 +947,7 @@ public static Entry snapshot(
                 assert existing == null || existing.equals(index) : "Conflicting indices [" + existing + "] and [" + index + "]";
                 final var shardSnapshotStatus = entry.getValue();
                 hasInitStateShards |= shardSnapshotStatus.state() == ShardState.INIT;
+                hasAssignedQueuedShards |= shardSnapshotStatus.isAssignedQueued();
                 byRepoShardIdBuilder.put(new RepositoryShardId(indexId, shardId.id()), shardSnapshotStatus);
             }
             return new Entry(
@@ -959,7 +967,8 @@ public static Entry snapshot(
                 null,
                 byRepoShardIdBuilder,
                 res,
-                hasInitStateShards
+                hasInitStateShards,
+                hasAssignedQueuedShards
             );
         }
 
@@ -991,6 +1000,7 @@ private static Entry createClone(
                 source,
                 shardStatusByRepoShardId,
                 Map.of(),
+                false,
                 false
             );
         }
@@ -1012,7 +1022,8 @@ private Entry(
             @Nullable SnapshotId source,
             Map<RepositoryShardId, ShardSnapshotStatus> shardStatusByRepoShardId,
             Map<String, Index> snapshotIndices,
-            boolean hasShardsInInitState
+            boolean hasShardsInInitState,
+            boolean hasAssignedQueuedShards
         ) {
             this.state = state;
             this.snapshot = snapshot;
@@ -1031,13 +1042,15 @@ private Entry(
             this.shardStatusByRepoShardId = Map.copyOf(shardStatusByRepoShardId);
             this.snapshotIndices = snapshotIndices;
             this.hasShardsInInitState = hasShardsInInitState;
+            this.hasAssignedQueuedShards = hasAssignedQueuedShards;
             assert assertShardsConsistent(
                 this.source,
                 this.state,
                 this.indices,
                 this.shards,
                 this.shardStatusByRepoShardId,
-                this.hasShardsInInitState
+                this.hasShardsInInitState,
+                this.hasAssignedQueuedShards
             );
         }
 
@@ -1087,14 +1100,18 @@ private static boolean assertShardsConsistent(
             Map<String, IndexId> indices,
             Map<ShardId, ShardSnapshotStatus> shards,
             Map<RepositoryShardId, ShardSnapshotStatus> statusByRepoShardId,
-            boolean hasInitStateShards
+            boolean hasInitStateShards,
+            boolean hasAssignedQueuedShards
         ) {
             if ((state == State.INIT || state == State.ABORTED) && shards.isEmpty()) {
                 return true;
             }
             if (hasInitStateShards) {
                 assert state == State.STARTED : "shouldn't have INIT-state shards in state " + state;
             }
+            if (hasAssignedQueuedShards) {
+                assert source == null : "clone entry must not have any shards in assigned-queued state";
+            }
             final Set<String> indexNames = indices.keySet();
             final Set<String> indexNamesInShards = new HashSet<>();
             shards.entrySet().forEach(s -> {
@@ -1150,7 +1167,8 @@ public Entry withRepoGen(long newRepoGen) {
                 source,
                 shardStatusByRepoShardId,
                 snapshotIndices,
-                hasShardsInInitState
+                hasShardsInInitState,
+                hasAssignedQueuedShards
             );
         }
 
@@ -1403,6 +1421,15 @@ public boolean hasShardsInInitState() {
             return hasShardsInInitState;
         }
 
+        /**
+         * See {@link #hasAssignedQueuedShards}.
+         */
+        public boolean hasAssignedQueuedShards() {
+            assert hasAssignedQueuedShards == false || isClone() == false
+                : "a clone entry must not have assigned-queued shards, but saw " + this;
+            return hasAssignedQueuedShards;
+        }
+
         public boolean partial() {
             return partial;
         }
@@ -1784,7 +1811,8 @@ public Entry apply(Entry part) {
                     null,
                     part.shardStatusByRepoShardId,
                     part.snapshotIndices,
-                    part.hasShardsInInitState
+                    part.hasShardsInInitState,
+                    part.hasAssignedQueuedShards
                 );
             }
             if (part.isClone()) {

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

@@ -95,7 +95,25 @@ public boolean hasCapacityOnAnyNode() {
         return enabled() == false || perNodeCounts.values().stream().anyMatch(count -> count < shardSnapshotPerNodeLimit);
     }
 
+    public boolean hasCapacityOnNode(String nodeId) {
+        if (enabled() == false) {
+            return true;
+        }
+        final Integer count = perNodeCounts.get(nodeId);
+        return count != null && count < shardSnapshotPerNodeLimit;
+    }
+
     private boolean enabled() {
         return shardSnapshotPerNodeLimit > 0;
     }
+
+    @Override
+    public String toString() {
+        return "PerNodeShardSnapshotCounter{"
+            + "shardSnapshotPerNodeLimit="
+            + shardSnapshotPerNodeLimit
+            + ", perNodeCounts="
+            + perNodeCounts
+            + '}';
+    }
 }

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

@@ -1342,7 +1342,6 @@ private static ImmutableOpenMap<ShardId, ShardSnapshotStatus> processWaitingShar
                 shards.put(shardId, failedState);
                 knownFailures.put(shardSnapshotEntry.getKey(), failedState);
             } else if (shardStatus.state().completed() == false && shardStatus.nodeId() != null) {
-                // TODO: This branch applies to assigned-queued shards. It seems OK since it applies to INIT as well. Double check it.
                 if (nodes.nodeExists(shardStatus.nodeId())) {
                     shards.put(shardId, shardStatus);
                 } else {
@@ -3089,16 +3088,16 @@ private static SnapshotsInProgress maybeStartAssignedQueuedShardSnapshotsForRepo
         Runnable changedCallback,
         Runnable startedCallback
     ) {
+        assert perNodeShardSnapshotCounter.hasCapacityOnAnyNode() : "no capacity left on any node " + perNodeShardSnapshotCounter;
         final List<SnapshotsInProgress.Entry> oldEntries = snapshotsInProgress.forRepo(projectRepo);
-        if (oldEntries.isEmpty()) {
+        if (oldEntries.isEmpty() || oldEntries.stream().allMatch(entry -> entry.hasAssignedQueuedShards() == false)) {
             return snapshotsInProgress;
         }
         final List<SnapshotsInProgress.Entry> newEntries = new ArrayList<>(oldEntries.size());
         for (SnapshotsInProgress.Entry entry : oldEntries) {
-            if (entry.isClone() == false && perNodeShardSnapshotCounter.hasCapacityOnAnyNode()) {
-                // TODO: Optimize this by checking whether the entry has any assigned-queued shards before building the shards map
+            if (entry.hasAssignedQueuedShards() && perNodeShardSnapshotCounter.hasCapacityOnAnyNode()) {
                 final var shardsBuilder = ImmutableOpenMap.builder(entry.shards());
-                maybeStartAssignedQueuedShardSnapshots(
+                final var changed = maybeStartAssignedQueuedShardSnapshots(
                     clusterState,
                     entry,
                     snapshotsInProgress::isNodeIdForRemoval,
@@ -3108,10 +3107,12 @@ private static SnapshotsInProgress maybeStartAssignedQueuedShardSnapshotsForRepo
                     changedCallback,
                     startedCallback
                 );
-                final var newEntry = entry.withShardStates(shardsBuilder.build());
-                newEntries.add(newEntry);
-                if (newEntry != entry) {
+                if (changed) {
+                    final var newEntry = entry.withShardStates(shardsBuilder.build());
+                    newEntries.add(newEntry);
                     newEntryConsumer.accept(newEntry);
+                } else {
+                    newEntries.add(entry);
                 }
             } else {
                 newEntries.add(entry);
@@ -3120,7 +3121,7 @@ private static SnapshotsInProgress maybeStartAssignedQueuedShardSnapshotsForRepo
         return snapshotsInProgress.createCopyWithUpdatedEntriesForRepo(projectRepo.projectId(), projectRepo.name(), newEntries);
     }
 
-    private static void maybeStartAssignedQueuedShardSnapshots(
+    private static boolean maybeStartAssignedQueuedShardSnapshots(
         ClusterState clusterState,
         SnapshotsInProgress.Entry entry,
         Predicate<String> nodeIdRemovalPredicate,
@@ -3130,41 +3131,46 @@ private static void maybeStartAssignedQueuedShardSnapshots(
         Runnable changedCallback,
         Runnable startedCallback
     ) {
-        if (perNodeShardSnapshotCounter.hasCapacityOnAnyNode()) {
-            for (var shardId : shardsBuilder.keys()) {
-                final var existingShardSnapshotStatus = shardsBuilder.get(shardId);
-                if (existingShardSnapshotStatus.isAssignedQueued() == false) {
-                    continue;
-                }
-                final IndexRoutingTable indexRouting = clusterState.routingTable(entry.projectId()).index(shardId.getIndex());
-                final ShardRouting shardRouting;
-                if (indexRouting == null) {
-                    shardRouting = null;
-                } else {
-                    shardRouting = indexRouting.shard(shardId.id()).primaryShard();
-                }
-                final var newShardSnapshotStatus = initShardSnapshotStatus(
-                    existingShardSnapshotStatus.generation(),
-                    shardRouting,
-                    nodeIdRemovalPredicate,
-                    perNodeShardSnapshotCounter
-                );
-                if (newShardSnapshotStatus.state().completed()) {
-                    // It can become complete if the shard is unassigned or deleted, i.e. state == MISSING.
-                    // We cannot directly update its status here because there maybe another snapshot for
-                    // the same shard that is QUEUED which must be updated as well, i.e. vertical update.
-                    // So we submit the status update to let it be processed in a future cluster state update.
-                    shardStatusUpdateConsumer.apply(entry.snapshot(), shardId, newShardSnapshotStatus);
-                    continue;
-                } else if (newShardSnapshotStatus.equals(existingShardSnapshotStatus) == false) {
-                    changedCallback.run();
-                    if (newShardSnapshotStatus.state() == ShardState.INIT) {
-                        startedCallback.run();
-                    }
+        assert entry.hasAssignedQueuedShards() : "entry has no assigned queued shards: " + entry;
+        assert perNodeShardSnapshotCounter.hasCapacityOnAnyNode() : "no capacity left on any node " + perNodeShardSnapshotCounter;
+        boolean changed = false;
+        for (var shardId : shardsBuilder.keys()) {
+            if (perNodeShardSnapshotCounter.hasCapacityOnAnyNode() == false) {
+                return changed;
+            }
+            final var existingShardSnapshotStatus = shardsBuilder.get(shardId);
+            if (existingShardSnapshotStatus.isAssignedQueued() == false) {
+                continue;
+            }
+            final IndexRoutingTable indexRouting = clusterState.routingTable(entry.projectId()).index(shardId.getIndex());
+            final ShardRouting shardRouting;
+            if (indexRouting == null) {
+                shardRouting = null;
+            } else {
+                shardRouting = indexRouting.shard(shardId.id()).primaryShard();
+            }
+            final var newShardSnapshotStatus = initShardSnapshotStatus(
+                existingShardSnapshotStatus.generation(),
+                shardRouting,
+                nodeIdRemovalPredicate,
+                perNodeShardSnapshotCounter
+            );
+            if (newShardSnapshotStatus.state().completed()) {
+                // It can become complete if the shard is unassigned or deleted, i.e. state == MISSING.
+                // We cannot directly update its status here because there maybe another snapshot for
+                // the same shard that is QUEUED which must be updated as well, i.e. vertical update.
+                // So we submit the status update to let it be processed in a future cluster state update.
+                shardStatusUpdateConsumer.apply(entry.snapshot(), shardId, newShardSnapshotStatus);
+            } else if (newShardSnapshotStatus.equals(existingShardSnapshotStatus) == false) {
+                changedCallback.run();
+                if (newShardSnapshotStatus.state() == ShardState.INIT) {
+                    startedCallback.run();
                 }
                 shardsBuilder.put(shardId, newShardSnapshotStatus);
+                changed = true;
             }
         }
+        return changed;
     }
 
     /**
@@ -3555,7 +3561,7 @@ SnapshotsInProgress computeUpdatedState() {
             // due to snapshots running for a repository that is update to completion in this batch.
             // (2) Repos that have seen updates in this batch because updates releasing capacity may all belong to later snapshots
             // than the one has assigned-queued shards. These updates could be either for the same repo or a different repo.
-            for (var repo : existing.repos()) {
+            for (var repo : updated.repos()) {
                 if (perNodeShardSnapshotCounter.hasCapacityOnAnyNode() == false) {
                     break;
                 }
@@ -3690,18 +3696,22 @@ SnapshotsInProgress.Entry computeUpdatedSnapshotEntryFromShardUpdates() {
                             + " as well as "
                             + shardsBuilder;
 
-                    // Shard snapshots changed status for this entry, check within the snapshot to see whether any previously limited
-                    // shard snapshots can now start due to newly completed ones.
-                    maybeStartAssignedQueuedShardSnapshots(
-                        initialState,
-                        entry,
-                        nodeIdRemovalPredicate,
-                        shardsBuilder,
-                        perNodeShardSnapshotCounter,
-                        shardStatusUpdateConsumer,
-                        () -> changedCount++,
-                        () -> startedCount++
-                    );
+                    if (entry.hasAssignedQueuedShards() && perNodeShardSnapshotCounter.hasCapacityOnAnyNode()) {
+                        // Shard snapshots changed status for this entry, check within the snapshot to see whether any previously limited
+                        // shard snapshots can now start due to newly completed ones. This is only necessary if the entry has any
+                        // assigned-queued shards before the update. If the entry gets any new assigned-queued shards from processing the
+                        // update, they cannot be started anyway because they already reflect the latest node capacities.
+                        maybeStartAssignedQueuedShardSnapshots(
+                            initialState,
+                            entry,
+                            nodeIdRemovalPredicate,
+                            shardsBuilder,
+                            perNodeShardSnapshotCounter,
+                            shardStatusUpdateConsumer,
+                            () -> changedCount++,
+                            () -> startedCount++
+                        );
+                    }
                     return entry.withShardStates(shardsBuilder.build());
                 } else if (clonesBuilder != null) {
                     return entry.withClones(clonesBuilder.build());