iter

Author: javanna

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

@@ -241,7 +241,11 @@ private static Executor wrapExecutor(Executor executor) {
             AtomicInteger segmentLevelTasks = new AtomicInteger(0);
             return command -> {
                 if (segmentLevelTasks.incrementAndGet() > tpe.getMaximumPoolSize()) {
-                    command.run();
+                    try {
+                        command.run();
+                    } finally {
+                        segmentLevelTasks.decrementAndGet();
+                    }
                 } else {
                     executor.execute(() -> {
                         try {

server/src/test/java/org/elasticsearch/search/DefaultSearchContextTests.java

@@ -1056,8 +1056,8 @@ public int size() {
         assertThat(executor.getCompletedTaskCount(), greaterThanOrEqualTo((long) numIters * executorPoolSize));
         // while we parallelize we also limit the number of tasks that each searcher submits
         assertThat(executor.getCompletedTaskCount(), lessThan((long) numIters * numSegmentTasks));
-        // *2 is just a wild guess to account for tasks that get executed while we are still submitting
-        assertThat(executor.getCompletedTaskCount(), lessThan((long) numIters * executorPoolSize * 2));
+        // *3 is just a wild guess to account for tasks that get executed while we are still submitting
+        assertThat(executor.getCompletedTaskCount(), lessThan((long) numIters * executorPoolSize * 3));
     }
 
     private void doTestSearchConcurrency(ThreadPoolExecutor executor, int numIters, int numSegmentTasks, AtomicInteger completedTasks)