diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java index 5f8bc57dcbe09..f11f3edb3bd55 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java @@ -262,12 +262,16 @@ public void testCancelFailedSearchWhenPartialResultDisallowed() throws Exception // When the search request executes, block all shards except 1. final List searchShardBlockingPlugins = initSearchShardBlockingPlugin(); AtomicBoolean letOneShardProceed = new AtomicBoolean(); + // Ensure we have at least one task waiting on the latch + CountDownLatch waitingTaskLatch = new CountDownLatch(1); CountDownLatch shardTaskLatch = new CountDownLatch(1); for (SearchShardBlockingPlugin plugin : searchShardBlockingPlugins) { plugin.setRunOnNewReaderContext((ReaderContext c) -> { if (letOneShardProceed.compareAndSet(false, true)) { // Let one shard continue. } else { + // Signal that we have a task waiting on the latch + waitingTaskLatch.countDown(); safeAwait(shardTaskLatch); // Block the other shards. } }); @@ -280,6 +284,9 @@ public void testCancelFailedSearchWhenPartialResultDisallowed() throws Exception plugin.disableBlock(); plugin.setBeforeExecution(() -> { if (oneThreadWillError.compareAndSet(false, true)) { + // wait for some task to get to the latch + safeAwait(waitingTaskLatch); + // then throw the exception throw new IllegalStateException("This will cancel the ContextIndexSearcher.search task"); } });