Drain responses on completion for TransportNodesAction

docs/changelog/130303.yaml

@@ -0,0 +1,5 @@
+pr: 130303
+summary: Drain responses on completion for `TransportNodesAction`
+area: Distributed
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/action/support/nodes/TransportNodesAction.java

@@ -163,8 +163,13 @@ protected void onItemFailure(DiscoveryNode discoveryNode, Exception e) {
             protected CheckedConsumer<ActionListener<NodesResponse>, Exception> onCompletion() {
                 // ref releases all happen-before here so no need to be synchronized
                 return l -> {
-                    try (var ignored = Releasables.wrap(Iterators.map(responses.iterator(), r -> r::decRef))) {
-                        newResponseAsync(task, request, actionContext, responses, exceptions, l);
+                    final List<NodeResponse> completedResponses;
+                    synchronized (responses) {
+                        completedResponses = List.copyOf(responses);
+                        responses.clear();
+                    }
+                    try (var ignored = Releasables.wrap(Iterators.map(completedResponses.iterator(), r -> r::decRef))) {
+                        newResponseAsync(task, request, actionContext, completedResponses, exceptions, l);
                     }
                 };
             }

server/src/test/java/org/elasticsearch/action/support/nodes/TransportNodesActionTests.java

@@ -57,8 +57,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.function.ObjLongConsumer;
@@ -316,6 +318,62 @@ protected Object createActionContext(Task task, TestNodesRequest request) {
         assertTrue(cancellableTask.isCancelled()); // keep task alive
     }
 
+    public void testCompletionAndCancellationShouldMutualExclusivelyHandleResponses() {
+        final var barrier = new CyclicBarrier(2);
+        final var action = new TestTransportNodesAction(
+            clusterService,
+            transportService,
+            new ActionFilters(Set.of()),
+            TestNodeRequest::new,
+            THREAD_POOL.executor(ThreadPool.Names.GENERIC)
+        ) {
+            @Override
+            protected void newResponseAsync(
+                Task task,
+                TestNodesRequest request,
+                Void unused,
+                List<TestNodeResponse> testNodeResponses,
+                List<FailedNodeException> failures,
+                ActionListener<TestNodesResponse> listener
+            ) {
+                final var waited = new AtomicBoolean();
+                for (var response : testNodeResponses) {
+                    if (waited.compareAndSet(false, true)) {
+                        safeAwait(barrier);
+                        safeAwait(barrier);
+                    }
+                }
+                super.newResponseAsync(task, request, unused, testNodeResponses, failures, listener);
+            }
+        };
+
+        final CancellableTask cancellableTask = new CancellableTask(randomLong(), "transport", "action", "", null, emptyMap());
+        final PlainActionFuture<TestNodesResponse> future = new PlainActionFuture<>();
+        action.execute(cancellableTask, new TestNodesRequest(), future);
+
+        final List<TestNodeResponse> nodeResponses = new ArrayList<>();
+        for (var capturedRequest : transport.getCapturedRequestsAndClear()) {
+            final var response = new TestNodeResponse(capturedRequest.node());
+            nodeResponses.add(response);
+            try {
+                transport.getTransportResponseHandler(capturedRequest.requestId()).handleResponse(response);
+            } finally {
+                response.decRef();
+            }
+        }
+
+        // Wait for the overall response starts to processing the node responses in a loop
+        safeAwait(barrier);
+
+        // Cancel the task while the overall response is being processed
+        TaskCancelHelper.cancel(cancellableTask, "simulated");
+
+        // Let the process continue to process the node responses and it should be successful
+        safeAwait(barrier);
+        safeGet(future);
+        assertTrue(nodeResponses.stream().allMatch(r -> r.hasReferences() == false));
+    }
+
     @BeforeClass
     public static void startThreadPool() {
         THREAD_POOL = new TestThreadPool(TransportNodesActionTests.class.getSimpleName());