From c6b2c656b13f8d2512cceb32802b9f7a5810c03a Mon Sep 17 00:00:00 2001
From: Michael Froh <froh@amazon.com>
Date: Wed, 29 Jan 2025 13:04:16 -0800
Subject: [PATCH] Fix flaky TransportMultiSearchActionTests.testCancellation

I recently added this test, but incorrectly placed a
CountDownLatch#await call on the test thread. With this change, we
actually kick off the request, return control to the testy thread,
cancel the request, then continue executing.

Signed-off-by: Michael Froh <froh@amazon.com>
---
 .../TransportMultiSearchActionTests.java      | 20 +++++++++----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/server/src/test/java/org/opensearch/action/search/TransportMultiSearchActionTests.java b/server/src/test/java/org/opensearch/action/search/TransportMultiSearchActionTests.java
index 45980e7137ce4..f969313b71833 100644
--- a/server/src/test/java/org/opensearch/action/search/TransportMultiSearchActionTests.java
+++ b/server/src/test/java/org/opensearch/action/search/TransportMultiSearchActionTests.java
@@ -321,11 +321,8 @@ 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
@@ -333,14 +330,15 @@ public void search(final SearchRequest request, final ActionListener<SearchRespo
                 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]);