Introduce executor for concurrent search

docs/changelog/98204.yaml

@@ -0,0 +1,5 @@
+pr: 98204
+summary: Introduce executor for concurrent search
+area: Search
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/common/util/concurrent/BoundedExecutor.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util.concurrent;
+
+import org.apache.lucene.util.ThreadInterruptedException;
+
+import java.util.Objects;
+import java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.ThreadPoolExecutor;
+
+/**
+ * {@link Executor} implementation that wraps another executor and limits the number of concurrent tasks that are delegated to it.
+ * Meant to be used when the caller thread is executed as part of a separate thread pool that is subject to its own queueing and rejection
+ * policy. Simplifies queue and rejection handling in the secondary thread pool as it effectively relies on those of the caller thread pool.
+ *
+ * Not suitable for non-blocking scenarios as it assumes that the caller thread will block and wait for all its tasks to be
+ * completed anyway. In this scenario an additional wait when tasks are submitted does not introduce extra overhead.
+ *
+ * Note: execution permits are released before the delegate executor updates its internal state to mark each task completed.
+ * That means that we'll accept tasks before there are free threads to take them, hence there will be some overflow queueing required in
+ * the underlying thread pool executor. This is the reason why a queue is required despite the wait applied at submit.
+ */
+public class BoundedExecutor implements Executor {
+    private final Executor executor;
+    private final int bound;
+    private final Semaphore semaphore;
+
+    public BoundedExecutor(ThreadPoolExecutor executor) {
+        this(executor, executor.getMaximumPoolSize());
+    }
+
+    public BoundedExecutor(Executor executor, int bound) {
+        if (bound <= 0) {
+            throw new IllegalArgumentException("bound should be positive");
+        }
+        this.executor = executor;
+        this.bound = bound;
+        this.semaphore = new Semaphore(bound, true);
+    }
+
+    /**
+     * Returns the number of tasks that this executor will allow to be executed concurrently
+     */
+    public final int getBound() {
+        return bound;
+    }
+
+    @Override
+    public void execute(Runnable command) {
+        Objects.requireNonNull(command, "command cannot be null");
+        try {
+            semaphore.acquire();
+        } catch (InterruptedException e) {
+            throw new ThreadInterruptedException(e);
+        }
+        try {
+            executor.execute(() -> {
+                try {
+                    command.run();
+                } finally {
+                    semaphore.release();
+                }
+            });
+        } catch (RejectedExecutionException e) {
+            semaphore.release();
+        }
+    }
+}

server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.action.search.SearchShardTask;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.common.util.concurrent.BoundedExecutor;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.TimeValue;
@@ -139,7 +140,8 @@ final class DefaultSearchContext extends SearchContext {
         int minimumDocsPerSlice,
         FetchPhase fetchPhase,
         boolean lowLevelCancellation,
-        boolean parallelize
+        BoundedExecutor executor,
+        boolean forceSequentialCollection
     ) throws IOException {
         this.readerContext = readerContext;
         this.request = request;
@@ -157,11 +159,8 @@ final class DefaultSearchContext extends SearchContext {
             engineSearcher.getQueryCachingPolicy(),
             minimumDocsPerSlice,
             lowLevelCancellation,
-            // TODO not set the for now, this needs a special thread pool and can be enabled after its introduction
-            // parallelize
-            // ? (EsThreadPoolExecutor) this.indexService.getThreadPool().executor(ThreadPool.Names.CONCURRENT_COLLECTION_TBD)
-            // : null,
-            null
+            executor,
+            forceSequentialCollection
         );
         releasables.addAll(List.of(engineSearcher, searcher));
 

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -39,6 +39,7 @@
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.common.util.concurrent.BoundedExecutor;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.Releasable;
@@ -130,6 +131,7 @@
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -284,6 +286,8 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
 
     private final Tracer tracer;
 
+    private final BoundedExecutor concurrentSearchExecutor;
+
     public SearchService(
         ClusterService clusterService,
         IndicesService indicesService,
@@ -346,16 +350,14 @@ public SearchService(
 
         enableConcurrentCollection = SEARCH_CONCURRENCY_ENABLED.get(settings);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(SEARCH_CONCURRENCY_ENABLED, this::setEnableConcurrentCollection);
+
+        concurrentSearchExecutor = new BoundedExecutor((ThreadPoolExecutor) threadPool.executor(Names.SEARCH_WORKER));
     }
 
     private void setEnableConcurrentCollection(boolean concurrentCollection) {
         this.enableConcurrentCollection = concurrentCollection;
     }
 
-    boolean isConcurrentCollectionEnabled() {
-        return this.enableConcurrentCollection;
-    }
-
     private static void validateKeepAlives(TimeValue defaultKeepAlive, TimeValue maxKeepAlive) {
         if (defaultKeepAlive.millis() > maxKeepAlive.millis()) {
             throw new IllegalArgumentException(
@@ -1039,7 +1041,7 @@ public DefaultSearchContext createSearchContext(ShardSearchRequest request, Time
         final Engine.SearcherSupplier reader = indexShard.acquireSearcherSupplier();
         final ShardSearchContextId id = new ShardSearchContextId(sessionId, idGenerator.incrementAndGet());
         try (ReaderContext readerContext = new ReaderContext(id, indexService, indexShard, reader, -1L, true)) {
-            DefaultSearchContext searchContext = createSearchContext(readerContext, request, timeout, null);
+            DefaultSearchContext searchContext = createSearchContext(readerContext, request, timeout, ResultsType.NONE);
             searchContext.addReleasable(readerContext.markAsUsed(0L));
             return searchContext;
         }
@@ -1069,7 +1071,8 @@ private DefaultSearchContext createSearchContext(
                 minimumDocsPerSlice,
                 fetchPhase,
                 lowLevelCancellation,
-                this.enableConcurrentCollection && concurrentSearchEnabled(resultsType, request.source())
+                this.enableConcurrentCollection ? concurrentSearchExecutor : null,
+                supportsConcurrency(resultsType, request.source()) == false
             );
             // we clone the query shard context here just for rewriting otherwise we
             // might end up with incorrect state since we are using now() or script services
@@ -1089,7 +1092,7 @@ this.enableConcurrentCollection && concurrentSearchEnabled(resultsType, request.
         return searchContext;
     }
 
-    static boolean concurrentSearchEnabled(ResultsType resultsType, SearchSourceBuilder source) {
+    static boolean supportsConcurrency(ResultsType resultsType, SearchSourceBuilder source) {
         if (resultsType == ResultsType.DFS) {
             return true; // only enable concurrent collection for DFS phase for now
         }