Run newShardSnapshotTask tasks concurrently (elastic#126478)

In elastic#88707 we changed the behaviour here to run the shard-snapshot
initialization tasks all in sequence. Yet these tasks do nontrivial work
since they may flush to acquire the relevant index commit, so with this
commit we go back to distributing them across the `SNAPSHOT` pool again.

Backport of elastic#126452 to `8.x`

Author: DaveCTurner

docs/changelog/126452.yaml

@@ -0,0 +1,5 @@
+pr: 126452
+summary: Run `newShardSnapshotTask` tasks concurrently
+area: Snapshot/Restore
+type: bug
+issues: []

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

@@ -9,16 +9,23 @@
 
 package org.elasticsearch.snapshots;
 
+import org.elasticsearch.action.admin.cluster.snapshots.create.CreateSnapshotRequest;
+import org.elasticsearch.action.admin.cluster.snapshots.create.TransportCreateSnapshotAction;
+import org.elasticsearch.cluster.SnapshotsInProgress;
 import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.snapshots.mockstore.MockRepository;
+import org.elasticsearch.test.ClusterServiceUtils;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.disruption.NetworkDisruption;
 import org.elasticsearch.test.transport.MockTransportService;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.threadpool.ThreadPoolStats;
 
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
@@ -89,4 +96,58 @@ public void testRetryPostingSnapshotStatusMessages() throws Exception {
             assertThat(snapshotInfo.successfulShards(), equalTo(shards));
         }, 30L, TimeUnit.SECONDS);
     }
+
+    public void testStartSnapshotsConcurrently() {
+        internalCluster().startMasterOnlyNode();
+        final var dataNode = internalCluster().startDataOnlyNode();
+
+        final var repoName = randomIdentifier();
+        createRepository(repoName, "fs");
+
+        final var threadPool = internalCluster().getInstance(ThreadPool.class, dataNode);
+        final var snapshotThreadCount = threadPool.info(ThreadPool.Names.SNAPSHOT).getMax();
+
+        final var indexName = randomIdentifier();
+        final var shardCount = between(1, snapshotThreadCount * 2);
+        assertAcked(prepareCreate(indexName, 0, indexSettingsNoReplicas(shardCount)));
+        indexRandomDocs(indexName, scaledRandomIntBetween(50, 100));
+
+        final var snapshotExecutor = threadPool.executor(ThreadPool.Names.SNAPSHOT);
+        final var barrier = new CyclicBarrier(snapshotThreadCount + 1);
+        for (int i = 0; i < snapshotThreadCount; i++) {
+            snapshotExecutor.submit(() -> {
+                safeAwait(barrier);
+                safeAwait(barrier);
+            });
+        }
+
+        // wait until the snapshot threads are all blocked
+        safeAwait(barrier);
+
+        safeGet(
+            client().execute(
+                TransportCreateSnapshotAction.TYPE,
+                new CreateSnapshotRequest(TEST_REQUEST_TIMEOUT, repoName, randomIdentifier())
+            )
+        );
+
+        // one task for each snapshot thread (throttled) or shard (if fewer), plus one for runSyncTasksEagerly()
+        assertEquals(Math.min(snapshotThreadCount, shardCount) + 1, getSnapshotQueueLength(threadPool));
+
+        // release all the snapshot threads
+        safeAwait(barrier);
+
+        // wait for completion
+        safeAwait(ClusterServiceUtils.addMasterTemporaryStateListener(cs -> SnapshotsInProgress.get(cs).isEmpty()));
+    }
+
+    private static int getSnapshotQueueLength(ThreadPool threadPool) {
+        for (ThreadPoolStats.Stats stats : threadPool.stats().stats()) {
+            if (stats.name().equals(ThreadPool.Names.SNAPSHOT)) {
+                return stats.queue();
+            }
+        }
+
+        throw new AssertionError("threadpool stats for [" + ThreadPool.Names.SNAPSHOT + "] not found");
+    }
 }

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

@@ -29,6 +29,7 @@
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.concurrent.ThrottledTaskRunner;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.index.seqno.SequenceNumbers;
@@ -54,7 +55,6 @@
 import org.elasticsearch.transport.TransportService;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -87,6 +87,9 @@ public final class SnapshotShardsService extends AbstractLifecycleComponent impl
     // A map of snapshots to the shardIds that we already reported to the master as failed
     private final ResultDeduplicator<UpdateIndexShardSnapshotStatusRequest, Void> remoteFailedRequestDeduplicator;
 
+    // Runs the tasks that start each shard snapshot (e.g. acquiring the index commit)
+    private final ThrottledTaskRunner startShardSnapshotTaskRunner;
+
     // Runs the tasks that promptly notify shards of aborted snapshots so that resources can be released ASAP
     private final ThrottledTaskRunner notifyOnAbortTaskRunner;
 
@@ -114,6 +117,11 @@ public SnapshotShardsService(
             threadPool.info(ThreadPool.Names.SNAPSHOT).getMax(),
             threadPool.generic()
         );
+        this.startShardSnapshotTaskRunner = new ThrottledTaskRunner(
+            "start-shard-snapshots",
+            threadPool.info(ThreadPool.Names.SNAPSHOT).getMax(),
+            threadPool.executor(ThreadPool.Names.SNAPSHOT)
+        );
     }
 
     @Override
@@ -304,7 +312,6 @@ private void startNewShardSnapshots(String localNodeId, SnapshotsInProgress.Entr
 
         final var newSnapshotShards = shardSnapshots.computeIfAbsent(snapshot, s -> new HashMap<>());
 
-        final List<Runnable> shardSnapshotTasks = new ArrayList<>(shardsToStart.size());
         for (final Map.Entry<ShardId, ShardGeneration> shardEntry : shardsToStart.entrySet()) {
             final ShardId shardId = shardEntry.getKey();
             final IndexShardSnapshotStatus snapshotStatus = IndexShardSnapshotStatus.newInitializing(shardEntry.getValue());
@@ -316,10 +323,36 @@ private void startNewShardSnapshots(String localNodeId, SnapshotsInProgress.Entr
                 : "Found non-null, non-numeric shard generation ["
                     + snapshotStatus.generation()
                     + "] for snapshot with old-format compatibility";
-            shardSnapshotTasks.add(newShardSnapshotTask(shardId, snapshot, indexId, snapshotStatus, entry.version(), entry.startTime()));
+            final var shardSnapshotTask = newShardSnapshotTask(
+                shardId,
+                snapshot,
+                indexId,
+                snapshotStatus,
+                entry.version(),
+                entry.startTime()
+            );
+            startShardSnapshotTaskRunner.enqueueTask(new ActionListener<>() {
+                @Override
+                public void onResponse(Releasable releasable) {
+                    try (releasable) {
+                        shardSnapshotTask.run();
+                    }
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    final var wrapperException = new IllegalStateException(
+                        "impossible failure starting shard snapshot for " + shardId + " in " + snapshot,
+                        e
+                    );
+                    logger.error(wrapperException.getMessage(), wrapperException);
+                    assert false : wrapperException; // impossible
+                }
+            });
         }
 
-        threadPool.executor(ThreadPool.Names.SNAPSHOT).execute(() -> shardSnapshotTasks.forEach(Runnable::run));
+        // apply some backpressure by reserving one SNAPSHOT thread for the startup work
+        startShardSnapshotTaskRunner.runSyncTasksEagerly(threadPool.executor(ThreadPool.Names.SNAPSHOT));
     }
 
     private void pauseShardSnapshotsForNodeRemoval(String localNodeId, SnapshotsInProgress.Entry entry) {