TaskExecutor should not fork unnecessarily

lucene/core/src/java/org/apache/lucene/search/TaskExecutor.java

@@ -30,27 +30,21 @@
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.RunnableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.ThreadInterruptedException;
 
 /**
  * Executor wrapper responsible for the execution of concurrent tasks. Used to parallelize search
  * across segments as well as query rewrite in some cases. Exposes a single {@link
  * #invokeAll(Collection)} method that takes a collection of {@link Callable}s and executes them
- * concurrently/ Once all tasks are submitted to the executor, it blocks and wait for all tasks to
- * be completed, and then returns a list with the obtained results. Ensures that the underlying
- * executor is only used for top-level {@link #invokeAll(Collection)} calls, and not for potential
- * {@link #invokeAll(Collection)} calls made from one of the tasks. This is to prevent deadlock with
- * certain types of pool based executors (e.g. {@link java.util.concurrent.ThreadPoolExecutor}).
+ * concurrently. Once all but one task have been submitted to the executor, it tries to run as many
+ * tasks as possible on the calling thread, then waits for all tasks that have been executed in
+ * parallel on the executor to be completed and then returns a list with the obtained results.
  *
  * @lucene.experimental
  */
 public final class TaskExecutor {
-  // a static thread local is ok as long as we use a counter, which accounts for multiple
-  // searchers holding a different TaskExecutor all backed by the same executor
-  private static final ThreadLocal<Integer> numberOfRunningTasksInCurrentThread =
-      ThreadLocal.withInitial(() -> 0);
-
   private final Executor executor;
 
   /**
@@ -84,26 +78,21 @@ public String toString() {
   /**
    * Holds all the sub-tasks that a certain operation gets split into as it gets parallelized and
    * exposes the ability to invoke such tasks and wait for them all to complete their execution and
-   * provide their results. Ensures that each task does not get parallelized further: this is
-   * important to avoid a deadlock in situations where one executor thread waits on other executor
-   * threads to complete before it can progress. This happens in situations where for instance
-   * {@link Query#createWeight(IndexSearcher, ScoreMode, float)} is called as part of searching each
-   * slice, like {@link TopFieldCollector#populateScores(ScoreDoc[], IndexSearcher, Query)} does.
-   * Additionally, if one task throws an exception, all other tasks from the same group are
-   * cancelled, to avoid needless computation as their results would not be exposed anyways. Creates
-   * one {@link FutureTask} for each {@link Callable} provided
+   * provide their results. Additionally, if one task throws an exception, all other tasks from the
+   * same group are cancelled, to avoid needless computation as their results would not be exposed
+   * anyways. Creates one {@link FutureTask} for each {@link Callable} provided
    *
    * @param <T> the return type of all the callables
    */
   private static final class TaskGroup<T> {
-    private final Collection<RunnableFuture<T>> futures;
+    private final List<RunnableFuture<T>> futures;
 
     TaskGroup(Collection<Callable<T>> callables) {
       List<RunnableFuture<T>> tasks = new ArrayList<>(callables.size());
       for (Callable<T> callable : callables) {
         tasks.add(createTask(callable));
       }
-      this.futures = Collections.unmodifiableCollection(tasks);
+      this.futures = Collections.unmodifiableList(tasks);
     }
 
     RunnableFuture<T> createTask(Callable<T> callable) {
@@ -112,15 +101,10 @@ RunnableFuture<T> createTask(Callable<T> callable) {
           () -> {
             if (startedOrCancelled.compareAndSet(false, true)) {
               try {
-                Integer counter = numberOfRunningTasksInCurrentThread.get();
-                numberOfRunningTasksInCurrentThread.set(counter + 1);
                 return callable.call();
               } catch (Throwable t) {
                 cancelAll();
                 throw t;
-              } finally {
-                Integer counter = numberOfRunningTasksInCurrentThread.get();
-                numberOfRunningTasksInCurrentThread.set(counter - 1);
               }
             }
             // task is cancelled hence it has no results to return. That's fine: they would be
@@ -144,32 +128,45 @@ public boolean cancel(boolean mayInterruptIfRunning) {
     }
 
     List<T> invokeAll(Executor executor) throws IOException {
-      boolean runOnCallerThread = numberOfRunningTasksInCurrentThread.get() > 0;
-      for (Runnable runnable : futures) {
-        if (runOnCallerThread) {
-          runnable.run();
-        } else {
-          executor.execute(runnable);
+      final int count = futures.size();
+      // taskId provides the first index of an un-executed task in #futures
+      final AtomicInteger taskId = new AtomicInteger(0);
+      // we fork execution count - 1 tasks to execute at least one task on the current thread to
+      // minimize needless forking and blocking of the current thread
+      if (count > 1) {
+        final Runnable work =
+            () -> {
+              int id = taskId.getAndIncrement();
+              if (id < count) {
+                futures.get(id).run();
+              }
+            };
+        for (int j = 0; j < count - 1; j++) {
+          executor.execute(work);
+        }
+      }
+      // try to execute as many tasks as possible on the current thread to minimize context
+      // switching in case of long running concurrent
+      // tasks as well as dead-locking if the current thread is part of #executor for executors that
+      // have limited or no parallelism
+      int id;
+      while ((id = taskId.getAndIncrement()) < count) {
+        futures.get(id).run();
+        if (id >= count - 1) {
+          // save redundant CAS in case this was the last task
+          break;
         }
       }
       Throwable exc = null;
-      List<T> results = new ArrayList<>(futures.size());
-      for (Future<T> future : futures) {
+      List<T> results = new ArrayList<>(count);
+      for (int i = 0; i < count; i++) {
+        Future<T> future = futures.get(i);
         try {
           results.add(future.get());
         } catch (InterruptedException e) {
-          var newException = new ThreadInterruptedException(e);
-          if (exc == null) {
-            exc = newException;
-          } else {
-            exc.addSuppressed(newException);
-          }
+          exc = IOUtils.useOrSuppress(exc, new ThreadInterruptedException(e));
         } catch (ExecutionException e) {
-          if (exc == null) {
-            exc = e.getCause();
-          } else {
-            exc.addSuppressed(e.getCause());
-          }
+          exc = IOUtils.useOrSuppress(exc, e.getCause());
         }
       }
       assert assertAllFuturesCompleted() : "Some tasks are still running?";

lucene/core/src/test/org/apache/lucene/search/TestIndexSearcher.java

@@ -286,7 +286,7 @@ protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
           }
         };
     searcher.search(new MatchAllDocsQuery(), 10);
-    assertEquals(leaves.size(), numExecutions.get());
+    assertEquals(leaves.size(), numExecutions.get() + 1);
   }
 
   public void testNullExecutorNonNullTaskExecutor() {

lucene/core/src/test/org/apache/lucene/search/TestTaskExecutor.java

@@ -234,11 +234,8 @@ public void testInvokeAllDoesNotLeaveTasksBehind() {
     TaskExecutor taskExecutor =
         new TaskExecutor(
             command -> {
-              executorService.execute(
-                  () -> {
-                    tasksStarted.incrementAndGet();
-                    command.run();
-                  });
+              tasksStarted.incrementAndGet();
+              command.run();
             });
     AtomicInteger tasksExecuted = new AtomicInteger(0);
     List<Callable<Void>> callables = new ArrayList<>();
@@ -258,7 +255,8 @@ public void testInvokeAllDoesNotLeaveTasksBehind() {
     expectThrows(RuntimeException.class, () -> taskExecutor.invokeAll(callables));
     assertEquals(1, tasksExecuted.get());
     // the callables are technically all run, but the cancelled ones will be no-op
-    assertEquals(100, tasksStarted.get());
+    // add one for the task the gets executed on the current thread
+    assertEquals(100, tasksStarted.get() + 1);
   }
 
   /**
@@ -308,7 +306,7 @@ public void testInvokeAllCatchesMultipleExceptions() {
   }
 
   public void testCancelTasksOnException() {
-    TaskExecutor taskExecutor = new TaskExecutor(executorService);
+    TaskExecutor taskExecutor = new TaskExecutor(Runnable::run);
     final int numTasks = random().nextInt(10, 50);
     final int throwingTask = random().nextInt(numTasks);
     boolean error = random().nextBoolean();