[Backport 2.x] Fix flaky TransportMultiSearchActionTests.testCancellation

server/src/test/java/org/opensearch/action/search/TransportMultiSearchActionTests.java

@@ -321,26 +321,24 @@ public TaskManager getTaskManager() {
         // and if there are more searches than is allowed create an error and remember that.
         int maxAllowedConcurrentSearches = 1; // Allow 1 search at a time.
         AtomicInteger counter = new AtomicInteger();
-        AtomicReference<AssertionError> errorHolder = new AtomicReference<>();
-        // randomize whether or not requests are executed asynchronously
         ExecutorService executorService = threadPool.executor(ThreadPool.Names.GENERIC);
-        final Set<SearchRequest> requests = Collections.newSetFromMap(Collections.synchronizedMap(new IdentityHashMap<>()));
-        CountDownLatch countDownLatch = new CountDownLatch(1);
+        CountDownLatch canceledLatch = new CountDownLatch(1);
         CancellableTask[] parentTask = new CancellableTask[1];
         NodeClient client = new NodeClient(settings, threadPool) {
             @Override
             public void search(final SearchRequest request, final ActionListener<SearchResponse> listener) {
                 if (parentTask[0] != null && parentTask[0].isCancelled()) {
                     fail("Should not execute search after parent task is cancelled");
                 }
-                try {
-                    countDownLatch.await(10, TimeUnit.MILLISECONDS);
-                } catch (InterruptedException e) {
-                    throw new RuntimeException(e);
-                }
 
-                requests.add(request);
                 executorService.execute(() -> {
+                    try {
+                        if (!canceledLatch.await(1, TimeUnit.SECONDS)) {
+                            fail("Latch should have counted down");
+                        }
+                    } catch (InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
                     counter.decrementAndGet();
                     listener.onResponse(
                         new SearchResponse(
@@ -399,7 +397,7 @@ public void onFailure(Task task, Exception e) {
                 }
             });
             parentTask[0].cancel("Giving up");
-            countDownLatch.countDown();
+            canceledLatch.countDown();
 
             assertNull(responses[0]);
             assertNull(exceptions[0]);