[Test] Wait for cancellation before assertions (elastic#134532)

Resolves: elastic#134277

Author: ywangd

muted-tests.yml

@@ -543,9 +543,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.expression.function.scalar.score.DecayTests
   method: "testEvaluateBlockWithNulls {TestCase=<integer>, <integer>, <integer>, <_source> #2}"
   issue: https://github.com/elastic/elasticsearch/issues/134679
-- class: org.elasticsearch.action.support.nodes.TransportNodesActionTests
-  method: testConcurrentlyCompletionAndCancellation
-  issue: https://github.com/elastic/elasticsearch/issues/134277
 - class: org.elasticsearch.search.vectors.CachingEnableFilterQueryTests
   method: testTermQuery
   issue: https://github.com/elastic/elasticsearch/issues/134723

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

@@ -58,6 +58,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
@@ -377,7 +378,13 @@ protected void newResponseAsync(
     public void testConcurrentlyCompletionAndCancellation() throws InterruptedException {
         final var action = getTestTransportNodesAction();
 
-        final CancellableTask cancellableTask = new CancellableTask(randomLong(), "transport", "action", "", null, emptyMap());
+        final CountDownLatch onCancelledLatch = new CountDownLatch(1);
+        final CancellableTask cancellableTask = new CancellableTask(randomLong(), "transport", "action", "", null, emptyMap()) {
+            @Override
+            protected void onCancelled() {
+                onCancelledLatch.countDown();
+            }
+        };
 
         final PlainActionFuture<TestNodesResponse> future = new PlainActionFuture<>();
         action.execute(cancellableTask, new TestNodesRequest(), future);
@@ -414,6 +421,11 @@ public void testConcurrentlyCompletionAndCancellation() throws InterruptedExcept
             assertNotNull("expect task cancellation exception, but got\n" + ExceptionsHelper.stackTrace(e), taskCancelledException);
             assertThat(e.getMessage(), containsString("task cancelled [simulated]"));
             assertTrue(cancellableTask.isCancelled());
+            // Wait for the latch, the listener for releasing node responses is called before it.
+            // We need to wait for the latch because the cancellation may be detected in CancellableFanOut#onCompletion with
+            // the responseHandled flag being true. The flag is set by the cancellation listener which is still in process of
+            // draining existing responses.
+            safeAwait(onCancelledLatch);
             // All previously captured responses are released due to cancellation
             assertTrue(nodeResponses.stream().allMatch(r -> r.hasReferences() == false));
             // Wait for the last response to be gathered and assert it is also released by either the concurrent cancellation or