fix race condition in SearchPhaseControllerTests#consumerTestCase

miamia1999 · Jul 28, 2020 · 4e4ed6e · 4e4ed6e
1 parent e0286e9
commit 4e4ed6e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/server/src/main/java/org/elasticsearch/action/search/QueryPhaseResultConsumer.java b/server/src/main/java/org/elasticsearch/action/search/QueryPhaseResultConsumer.java
@@ -277,13 +277,15 @@ private void onMergeFailure(Exception exc) {
                 }
                 failure.compareAndSet(null, exc);
                 MergeTask task = runningTask.get();
+                runningTask.compareAndSet(task, null);
+                List<MergeTask> toCancel = new ArrayList<>();
                 if (task != null) {
-                    runningTask.compareAndSet(task, null);
-                    task.cancel();
+                    toCancel.add(task);
                 }
-                queue.stream().forEach(MergeTask::cancel);
+                queue.stream().forEach(toCancel::add);
                 queue.clear();
                 mergeResult = null;
+                toCancel.stream().forEach(MergeTask::cancel);
             }
             onPartialMergeFailure.accept(exc);
         }

diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchPhaseControllerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchPhaseControllerTests.java
@@ -487,8 +487,6 @@ private void consumerTestCase(int numEmptyResponses) throws Exception {
         }
 
         SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
-        int numCompletedTasks = (int) (fixedExecutor.getCompletedTaskCount() - beforeCompletedTasks);
-        assertEquals(numCompletedTasks, reduce.numReducePhases-1);
         assertEquals(numTotalReducePhases, reduce.numReducePhases);
         assertEquals(numTotalReducePhases, reductions.size());
         assertAggReduction(request);