newPublicationListener

Author: DaveCTurner

server/src/main/java/org/elasticsearch/cluster/action/shard/ShardStateAction.java

@@ -333,7 +333,7 @@ public ClusterTasksResult<FailedShardUpdateTask> execute(ClusterState currentSta
                         entry,
                         entry.getShardId().getIndex()
                     );
-                    batchResultBuilder.success(task, task);
+                    batchResultBuilder.success(task, task.newPublicationListener());
                 } else {
                     // The primary term is 0 if the shard failed itself. It is > 0 if a write was done on a primary but was failed to be
                     // replicated to the shard copy with the provided allocation id. In case where the shard failed itself, it's ok to just
@@ -393,7 +393,7 @@ public ClusterTasksResult<FailedShardUpdateTask> execute(ClusterState currentSta
                         } else {
                             // tasks that correspond to non-existent shards are marked as successful
                             logger.debug("{} ignoring shard failed task [{}] (shard does not exist anymore)", entry.getShardId(), entry);
-                            batchResultBuilder.success(task, task);
+                            batchResultBuilder.success(task, task.newPublicationListener());
                         }
                     } else {
                         // failing a shard also possibly marks it as stale (see IndexMetadataUpdater)
@@ -409,7 +409,7 @@ public ClusterTasksResult<FailedShardUpdateTask> execute(ClusterState currentSta
             try {
                 maybeUpdatedState = applyFailedShards(currentState, failedShardsToBeApplied, staleShardsToBeApplied);
                 for (var task : tasksToBeApplied) {
-                    batchResultBuilder.success(task, task);
+                    batchResultBuilder.success(task, task.newPublicationListener());
                 }
             } catch (Exception e) {
                 logger.warn(() -> new ParameterizedMessage("failed to apply failed shards {}", failedShardsToBeApplied), e);
@@ -535,16 +535,21 @@ public int hashCode() {
 
     public record FailedShardUpdateTask(FailedShardEntry entry, ActionListener<TransportResponse.Empty> listener)
         implements
-            ClusterStateTaskListener,
-            ActionListener<ClusterState> {
+            ClusterStateTaskListener {
 
-        /**
-         * This task is re-used as its own publication listener, so this method is called when it has been executed and the resulting
-         * publication completed successfully.
-         */
-        @Override
-        public void onResponse(ClusterState clusterState) {
-            listener.onResponse(TransportResponse.Empty.INSTANCE);
+        public ActionListener<ClusterState> newPublicationListener() {
+            return new ActionListener<>() {
+                @Override
+                public void onResponse(ClusterState clusterState) {
+                    listener.onResponse(TransportResponse.Empty.INSTANCE);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    // delegate to task's onFailure for logging
+                    FailedShardUpdateTask.this.onFailure(e);
+                }
+            };
         }
 
         @Override
@@ -649,7 +654,7 @@ public ClusterTasksResult<StartedShardUpdateTask> execute(ClusterState currentSt
                     // requests might still be in flight even after the shard has already been started or failed on the master. We just
                     // ignore these requests for now.
                     logger.debug("{} ignoring shard started task [{}] (shard does not exist anymore)", entry.shardId, entry);
-                    builder.success(task, task);
+                    builder.success(task, task.newPublicationListener());
                 } else {
                     if (matched.primary() && entry.primaryTerm > 0) {
                         final IndexMetadata indexMetadata = currentState.metadata().index(entry.shardId.getIndex());
@@ -670,7 +675,7 @@ public ClusterTasksResult<StartedShardUpdateTask> execute(ClusterState currentSt
                                 entry.primaryTerm,
                                 currentPrimaryTerm
                             );
-                            builder.success(task, task);
+                            builder.success(task, task.newPublicationListener());
                             continue;
                         }
                     }
@@ -683,7 +688,7 @@ public ClusterTasksResult<StartedShardUpdateTask> execute(ClusterState currentSt
                             entry,
                             matched
                         );
-                        builder.success(task, task);
+                        builder.success(task, task.newPublicationListener());
                     } else {
                         // remove duplicate actions as allocation service expects a clean list without duplicates
                         if (seenShardRoutings.contains(matched)) {
@@ -740,7 +745,7 @@ public ClusterTasksResult<StartedShardUpdateTask> execute(ClusterState currentSt
                 assert assertStartedIndicesHaveCompleteTimestampRanges(maybeUpdatedState);
 
                 for (var task : tasksToBeApplied) {
-                    builder.success(task, task);
+                    builder.success(task, task.newPublicationListener());
                 }
             } catch (Exception e) {
                 logger.warn(() -> new ParameterizedMessage("failed to apply started shards {}", shardRoutingsToBeApplied), e);
@@ -849,20 +854,25 @@ public int hashCode() {
 
     public record StartedShardUpdateTask(StartedShardEntry entry, ActionListener<TransportResponse.Empty> listener)
         implements
-            ClusterStateTaskListener,
-            ActionListener<ClusterState> {
+            ClusterStateTaskListener {
 
         public StartedShardEntry getEntry() {
             return entry;
         }
 
-        /**
-         * This task is re-used as its own publication listener, so this method is called when it has been executed and the resulting
-         * publication completed successfully.
-         */
-        @Override
-        public void onResponse(ClusterState clusterState) {
-            listener.onResponse(TransportResponse.Empty.INSTANCE);
+        public ActionListener<ClusterState> newPublicationListener() {
+            return new ActionListener<>() {
+                @Override
+                public void onResponse(ClusterState clusterState) {
+                    listener.onResponse(TransportResponse.Empty.INSTANCE);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    // delegate to task's onFailure for logging
+                    StartedShardUpdateTask.this.onFailure(e);
+                }
+            };
         }
 
         @Override