Remove LegacyCTRAL from SnapshotsService (#86178)

Relates #83784

Author: DaveCTurner

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

@@ -3012,15 +3012,9 @@ public boolean assertAllListenersResolved() {
      *
      * Package private to allow for tests.
      */
-    static final ClusterStateTaskExecutor<ShardSnapshotUpdate> SHARD_STATE_EXECUTOR = (currentState, taskContexts) -> {
-        for (var taskContext : taskContexts) {
-            taskContext.success(new ClusterStateTaskExecutor.LegacyClusterTaskResultActionListener(taskContext.getTask(), currentState));
-        }
-        return new SnapshotShardsUpdateContext(
-            currentState,
-            () -> taskContexts.stream().map(ClusterStateTaskExecutor.TaskContext::getTask).iterator()
-        ).computeUpdatedState();
-    };
+    static final ClusterStateTaskExecutor<ShardSnapshotUpdate> SHARD_STATE_EXECUTOR = (
+        currentState,
+        taskContexts) -> new SnapshotShardsUpdateContext(currentState, taskContexts).computeUpdatedState();
 
     private static boolean isQueued(@Nullable ShardSnapshotStatus status) {
         return status != null && status.state() == ShardState.QUEUED;
@@ -3041,17 +3035,25 @@ private static final class SnapshotShardsUpdateContext {
         // current cluster state
         private final ClusterState currentState;
 
+        // task contexts to be completed on success
+        private final List<ClusterStateTaskExecutor.TaskContext<ShardSnapshotUpdate>> taskContexts;
+
         // updates outstanding to be applied to existing snapshot entries
         private final Map<String, List<ShardSnapshotUpdate>> updatesByRepo;
 
         // updates that were used to update an existing in-progress shard snapshot
         private final Set<ShardSnapshotUpdate> executedUpdates = new HashSet<>();
 
-        SnapshotShardsUpdateContext(ClusterState currentState, Iterable<ShardSnapshotUpdate> updates) {
+        SnapshotShardsUpdateContext(
+            ClusterState currentState,
+            List<ClusterStateTaskExecutor.TaskContext<ShardSnapshotUpdate>> taskContexts
+        ) {
             this.currentState = currentState;
+            this.taskContexts = taskContexts;
             updatesByRepo = new HashMap<>();
-            for (ShardSnapshotUpdate update : updates) {
-                updatesByRepo.computeIfAbsent(update.snapshot.getRepository(), r -> new ArrayList<>()).add(update);
+            for (final var taskContext : taskContexts) {
+                updatesByRepo.computeIfAbsent(taskContext.getTask().snapshot.getRepository(), r -> new ArrayList<>())
+                    .add(taskContext.getTask());
             }
         }
 
@@ -3071,6 +3073,11 @@ ClusterState computeUpdatedState() {
                 updated = updated.withUpdatedEntriesForRepo(repoName, newEntries);
             }
 
+            final var result = new ShardSnapshotUpdateResult(currentState.metadata(), updated);
+            for (final var taskContext : taskContexts) {
+                taskContext.success(taskContext.getTask().listener.map(ignored -> result));
+            }
+
             if (changedCount > 0) {
                 logger.trace(
                     "changed cluster state triggered by [{}] snapshot state updates and resulted in starting " + "[{}] shard snapshots",
@@ -3079,6 +3086,7 @@ ClusterState computeUpdatedState() {
                 );
                 return ClusterState.builder(currentState).putCustom(SnapshotsInProgress.TYPE, updated).build();
             }
+            assert existing == updated;
             return currentState;
         }
 
@@ -3312,6 +3320,11 @@ private ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> shardsBuilder() {
         }
     }
 
+    /**
+     * The result of a {@link ShardSnapshotUpdate}, capturing the info needed to finalize the relevant snapshot if appropriate.
+     */
+    record ShardSnapshotUpdateResult(Metadata metadata, SnapshotsInProgress snapshotsInProgress) {}
+
     /**
      * An update to the snapshot state of a shard.
      *
@@ -3323,14 +3336,14 @@ static final class ShardSnapshotUpdate implements ClusterStateTaskListener {
         private final ShardId shardId;
         private final RepositoryShardId repoShardId;
         private final ShardSnapshotStatus updatedState;
-        private final ActionListener<ClusterState> listener;
+        private final ActionListener<ShardSnapshotUpdateResult> listener;
 
         ShardSnapshotUpdate(
             Snapshot snapshot,
             ShardId shardId,
             RepositoryShardId repoShardId,
             ShardSnapshotStatus updatedState,
-            ActionListener<ClusterState> listener
+            ActionListener<ShardSnapshotUpdateResult> listener
         ) {
             assert shardId != null ^ repoShardId != null;
             this.snapshot = snapshot;
@@ -3351,7 +3364,7 @@ public void onFailure(Exception e) {
 
         @Override
         public void clusterStateProcessed(ClusterState oldState, ClusterState newState) {
-            listener.onResponse(newState);
+            assert false : "never called";
         }
 
         @Override
@@ -3396,29 +3409,23 @@ private void innerUpdateSnapshotState(
         ShardSnapshotStatus updatedState,
         ActionListener<Void> listener
     ) {
-        var update = new ShardSnapshotUpdate(
-            snapshot,
-            shardId,
-            repoShardId,
-            updatedState,
-            listener.delegateFailure((delegate, newState) -> {
-                try {
-                    delegate.onResponse(null);
-                } finally {
-                    // Maybe this state update completed the snapshot. If we are not already ending it because of a concurrent
-                    // state update we check if its state is completed and end it if it is.
-                    final SnapshotsInProgress snapshotsInProgress = newState.custom(SnapshotsInProgress.TYPE, SnapshotsInProgress.EMPTY);
-                    if (endingSnapshots.contains(snapshot) == false) {
-                        final SnapshotsInProgress.Entry updatedEntry = snapshotsInProgress.snapshot(snapshot);
-                        // If the entry is still in the cluster state and is completed, try finalizing the snapshot in the repo
-                        if (updatedEntry != null && updatedEntry.state().completed()) {
-                            endSnapshot(updatedEntry, newState.metadata(), null);
-                        }
+        var update = new ShardSnapshotUpdate(snapshot, shardId, repoShardId, updatedState, listener.delegateFailure((delegate, result) -> {
+            try {
+                delegate.onResponse(null);
+            } finally {
+                // Maybe this state update completed the snapshot. If we are not already ending it because of a concurrent
+                // state update we check if its state is completed and end it if it is.
+                final SnapshotsInProgress snapshotsInProgress = result.snapshotsInProgress();
+                if (endingSnapshots.contains(snapshot) == false) {
+                    final SnapshotsInProgress.Entry updatedEntry = snapshotsInProgress.snapshot(snapshot);
+                    // If the entry is still in the cluster state and is completed, try finalizing the snapshot in the repo
+                    if (updatedEntry != null && updatedEntry.state().completed()) {
+                        endSnapshot(updatedEntry, result.metadata(), null);
                     }
-                    startExecutableClones(snapshotsInProgress, snapshot.getRepository());
                 }
-            })
-        );
+                startExecutableClones(snapshotsInProgress, snapshot.getRepository());
+            }
+        }));
         logger.trace("received updated snapshot restore state [{}]", update);
         clusterService.submitStateUpdateTask(
             "update snapshot state",

test/framework/src/main/java/org/elasticsearch/cluster/service/ClusterStateTaskExecutorUtils.java

@@ -128,5 +128,10 @@ public void success(ActionListener<ClusterState> publishListener) {
             assert task instanceof ClusterStateAckListener == false;
             this.succeeded = true;
         }
+
+        @Override
+        public String toString() {
+            return "TestTaskContext[" + task + "]";
+        }
     }
 }