Address race condition in ParallelVisitor

A task could potentially enqueue items and stop running by the time between checking if the processing queue is empty and the check for pending tasks.

PiperOrigin-RevId: 248765394

Author: shreyax

src/main/java/com/google/devtools/build/lib/query2/ParallelVisitor.java

@@ -257,7 +257,7 @@ private void visitAndWaitForCompletion() throws QueryException, InterruptedExcep
       // 1. Errors (QueryException or InterruptedException) occurred and visitations should fail
       //    fast.
       // 2. There is no pending visit in the queue and no pending task running.
-      while (!mustJobsBeStopped() && (!processingQueue.isEmpty() || getTaskCount() > 0)) {
+      while (!mustJobsBeStopped() && moreWorkToDo()) {
         // To achieve maximum efficiency, queue is drained in either of the following two
         // conditions:
         //
@@ -292,6 +292,15 @@ private void visitAndWaitForCompletion() throws QueryException, InterruptedExcep
       awaitTerminationAndPropagateErrorsIfAny();
     }
 
+    private boolean moreWorkToDo() {
+      // Note that we must check the task count first -- checking the processing queue first has the
+      // following race condition:
+      // (1) Check processing queue and observe that it is empty
+      // (2) A remaining task adds to the processing queue and shuts down
+      // (3) We check the task count and observe it is empty
+      return getTaskCount() > 0 || !processingQueue.isEmpty();
+    }
+
     private void awaitTerminationAndPropagateErrorsIfAny()
         throws QueryException, InterruptedException {
       try {