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 6f644e1)

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>

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]);