1 files changed, 15 insertions, 13 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
index 39385ddd78..51fb19c50a 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
@@ -101,6 +101,7 @@ import java.util.Map.Entry;
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -133,11 +134,12 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   private final List<String> universeScope;
   protected final String parserPrefix;
   private final PathPackageLocator pkgPath;
-  private final ForkJoinPool forkJoinPool;
+  private final int queryEvaluationParallelismLevel;
 
   // The following fields are set in the #beforeEvaluateQuery method.
   protected WalkableGraph graph;
   private InterruptibleSupplier<ImmutableSet<PathFragment>> blacklistPatternsSupplier;
+  private ForkJoinPool forkJoinPool;
   private RecursivePackageProviderBackedTargetPatternResolver resolver;
 
   public SkyQueryEnvironment(
@@ -188,15 +190,13 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
         extraFunctions,
         evalListener);
     this.loadingPhaseThreads = loadingPhaseThreads;
-    // Note that ForkJoinPool doesn't start any thread until work is submitted to it.
-    this.forkJoinPool = NamedForkJoinPool.newNamedPool(
-        "QueryEnvironment", queryEvaluationParallelismLevel);
     this.graphFactory = graphFactory;
     this.pkgPath = pkgPath;
     this.universeScope = Preconditions.checkNotNull(universeScope);
     this.parserPrefix = parserPrefix;
-    Preconditions.checkState(!universeScope.isEmpty(),
-        "No queries can be performed with an empty universe");
+    Preconditions.checkState(
+        !universeScope.isEmpty(), "No queries can be performed with an empty universe");
+    this.queryEvaluationParallelismLevel = queryEvaluationParallelismLevel;
   }
 
   private void beforeEvaluateQuery() throws InterruptedException {
@@ -217,6 +217,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
             PrepareDepsOfPatternsFunction.getSkyKeys(universeKey, eventHandler));
     GraphBackedRecursivePackageProvider graphBackedRecursivePackageProvider =
         new GraphBackedRecursivePackageProvider(graph, universeTargetPatternKeys, pkgPath);
+    forkJoinPool =
+        NamedForkJoinPool.newNamedPool("QueryEnvironment", queryEvaluationParallelismLevel);
     resolver =
         new RecursivePackageProviderBackedTargetPatternResolver(
             graphBackedRecursivePackageProvider,
@@ -285,11 +287,6 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   }
 
   @Override
-  public void close() {
-    forkJoinPool.shutdownNow();
-  }
-
-  @Override
   public final QueryExpression transformParsedQuery(QueryExpression queryExpression) {
     QueryExpression transformedQueryExpression = getTransformedQueryExpression(queryExpression);
     LOG.info(String.format(
@@ -345,6 +342,12 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
         eval(expr, VariableContext.<Target>empty(), callbackWithEmptyCheck);
       } catch (QueryException e) {
         throw new QueryException(e, expr);
+      } finally {
+        // Force termination of remaining tasks - if evaluateQuery was successful there should be
+        // none, if it failed abruptly (e.g. was interrupted) we don't want to leave any dangling
+        // threads running tasks.
+        forkJoinPool.shutdownNow();
+        forkJoinPool.awaitQuiescence(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
       }
       aggregator.processLastPending();
     }
@@ -794,8 +797,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     if (originalFileFragment.equals(currentPathFragment)
         && originalFileFragment.equals(Label.EXTERNAL_PACKAGE_FILE_NAME)) {
       Preconditions.checkState(
-          Label.EXTERNAL_PACKAGE_FILE_NAME.getParentDirectory().equals(
-              PathFragment.EMPTY_FRAGMENT),
+          Label.EXTERNAL_PACKAGE_FILE_NAME.getParentDirectory().equals(PathFragment.EMPTY_FRAGMENT),
           Label.EXTERNAL_PACKAGE_FILE_NAME);
       return ImmutableList.of(
           PackageLookupValue.key(Label.EXTERNAL_PACKAGE_IDENTIFIER),