From bc0751d3928902fd83fd124de91a23000fa75b2c Mon Sep 17 00:00:00 2001 From: Panagiotis Bailis Date: Fri, 1 Dec 2023 12:18:49 +0200 Subject: [PATCH] [CI] Fix for SearchCancellationIT (#102774) --- .../search/SearchCancellationIT.java | 11 ++-- .../AbstractSearchCancellationTestCase.java | 57 +++++++++++-------- 2 files changed, 36 insertions(+), 32 deletions(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java index ad610954e86b..c41984b468b1 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java @@ -8,7 +8,6 @@ package org.elasticsearch.search; -import org.apache.lucene.tests.util.LuceneTestCase; import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.action.ActionFuture; import org.elasticsearch.action.search.MultiSearchResponse; @@ -50,8 +49,7 @@ import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.notNullValue; -@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.SUITE) -@LuceneTestCase.AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/102257") +@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST) public class SearchCancellationIT extends AbstractSearchCancellationTestCase { @Override @@ -288,12 +286,11 @@ public class SearchCancellationIT extends AbstractSearchCancellationTestCase { assertTrue("All SearchShardTasks should then be cancelled", shardQueryTask.isCancelled()); } }, 30, TimeUnit.SECONDS); - shardTaskLatch.countDown(); // unblock the shardTasks, allowing the test to conclude. } finally { + shardTaskLatch.countDown(); // unblock the shardTasks, allowing the test to conclude. searchThread.join(); - for (ScriptedBlockPlugin plugin : plugins) { - plugin.setBeforeExecution(() -> {}); - } + plugins.forEach(plugin -> plugin.setBeforeExecution(() -> {})); + searchShardBlockingPlugins.forEach(plugin -> plugin.setRunOnNewReaderContext((ReaderContext c) -> {})); } } diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractSearchCancellationTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractSearchCancellationTestCase.java index 7049954dc43f..ba0921972778 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/AbstractSearchCancellationTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/test/AbstractSearchCancellationTestCase.java @@ -35,7 +35,8 @@ import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -152,7 +153,7 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase { private final AtomicInteger hits = new AtomicInteger(); - private final AtomicBoolean shouldBlock = new AtomicBoolean(true); + private final Semaphore shouldBlock = new Semaphore(Integer.MAX_VALUE); private final AtomicReference beforeExecution = new AtomicReference<>(); @@ -161,11 +162,16 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase { } public void disableBlock() { - shouldBlock.set(false); + shouldBlock.release(Integer.MAX_VALUE); } public void enableBlock() { - shouldBlock.set(true); + try { + shouldBlock.acquire(Integer.MAX_VALUE); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } } public void setBeforeExecution(Runnable runnable) { @@ -196,6 +202,23 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase { ); } + public void logIfBlocked(String logMessage) { + if (shouldBlock.tryAcquire(1) == false) { + LogManager.getLogger(AbstractSearchCancellationTestCase.class).info(logMessage); + } else { + shouldBlock.release(1); + } + } + + public void waitForLock(int timeout, TimeUnit timeUnit) { + try { + assertTrue(shouldBlock.tryAcquire(timeout, timeUnit)); + shouldBlock.release(1); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + private Object searchBlockScript(Map params) { final Runnable runnable = beforeExecution.get(); if (runnable != null) { @@ -204,11 +227,7 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase { LeafStoredFieldsLookup fieldsLookup = (LeafStoredFieldsLookup) params.get("_fields"); LogManager.getLogger(AbstractSearchCancellationTestCase.class).info("Blocking on the document {}", fieldsLookup.get("_id")); hits.incrementAndGet(); - try { - assertBusy(() -> assertFalse(shouldBlock.get())); - } catch (Exception e) { - throw new RuntimeException(e); - } + waitForLock(10, TimeUnit.SECONDS); return true; } @@ -226,15 +245,9 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase { if (runnable != null) { runnable.run(); } - if (shouldBlock.get()) { - LogManager.getLogger(AbstractSearchCancellationTestCase.class).info("Blocking in reduce"); - } + logIfBlocked("Blocking in reduce"); hits.incrementAndGet(); - try { - assertBusy(() -> assertFalse(shouldBlock.get())); - } catch (Exception e) { - throw new RuntimeException(e); - } + waitForLock(10, TimeUnit.SECONDS); return 42; } @@ -243,15 +256,9 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase { if (runnable != null) { runnable.run(); } - if (shouldBlock.get()) { - LogManager.getLogger(AbstractSearchCancellationTestCase.class).info("Blocking in map"); - } + logIfBlocked("Blocking in map"); hits.incrementAndGet(); - try { - assertBusy(() -> assertFalse(shouldBlock.get())); - } catch (Exception e) { - throw new RuntimeException(e); - } + waitForLock(10, TimeUnit.SECONDS); return 1; }