From ffed7170911130d592174cdd8c375f6ede8afcd8 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Wed, 29 Jan 2025 20:34:18 -0500
Subject: [PATCH] Fix flaky TransportMultiSearchActionTests.testCancellation
 (#17193) (#17196)

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.


(cherry picked from commit 6f644e1c12de709b2136b999c8cb112a27c62100)

Signed-off-by: Michael Froh <froh@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.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]);