[CI] Fix for SearchCancellationIT (#102774)

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)
+@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST)
-@LuceneTestCase.AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/102257")
 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) {
+            plugins.forEach(plugin -> plugin.setBeforeExecution(() -> {}));
-                plugin.setBeforeExecution(() -> {});
+            searchShardBlockingPlugins.forEach(plugin -> plugin.setRunOnNewReaderContext((ReaderContext c) -> {}));
-            }
         }
     }
 

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<Runnable> 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<String, Object> 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 {
+            waitForLock(10, TimeUnit.SECONDS);
-                assertBusy(() -> assertFalse(shouldBlock.get()));
-            } catch (Exception e) {
-                throw new RuntimeException(e);
-            }
             return true;
         }
 
@@ -226,15 +245,9 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase {
             if (runnable != null) {
                 runnable.run();
             }
-            if (shouldBlock.get()) {
+            logIfBlocked("Blocking in reduce");
-                LogManager.getLogger(AbstractSearchCancellationTestCase.class).info("Blocking in reduce");
-            }
             hits.incrementAndGet();
-            try {
+            waitForLock(10, TimeUnit.SECONDS);
-                assertBusy(() -> assertFalse(shouldBlock.get()));
-            } catch (Exception e) {
-                throw new RuntimeException(e);
-            }
             return 42;
         }
 
@@ -243,15 +256,9 @@ public class AbstractSearchCancellationTestCase extends ESIntegTestCase {
             if (runnable != null) {
                 runnable.run();
             }
-            if (shouldBlock.get()) {
+            logIfBlocked("Blocking in map");
-                LogManager.getLogger(AbstractSearchCancellationTestCase.class).info("Blocking in map");
-            }
             hits.incrementAndGet();
-            try {
+            waitForLock(10, TimeUnit.SECONDS);
-                assertBusy(() -> assertFalse(shouldBlock.get()));
-            } catch (Exception e) {
-                throw new RuntimeException(e);
-            }
             return 1;
         }