Make AbstractBlazeQueryEnvironment no longer implement AutoCloseable. Instead have SkyQueryEnvironment#evaluateQuery be responsible for handling cleanup of its internal ForkJoinPool.

In addition to being a more sensible way of organizing the code (imo, it makes sense for SkyQueryEnvironment to clean up after itself), this fixes several issues:
(i) If query evaluation is interrupted, the AbstractBlazeQueryEnvironment#close call at the end of the try-with-resources statement in QueryCommand would be blocking and non-urgent. N.B. This was not an issue with the current ForkJoinPool usage, but was an issue with the old ThreadPoolExecutor usage.
(ii) Because of how the code in QueryCommand was structured, OutputFormatterCallback#close would happen _before_ the AbstractBlazeQueryEnvironment#close call. If query evaluation is interrupted, threads executing query tasks may be invoking the callback after the callback had been shut down!

--
MOS_MIGRATED_REVID=134573395

3 files changed, 16 insertions, 24 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
index cd597a66cd..1962cdfaf4 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
@@ -46,7 +46,7 @@ import java.util.logging.Logger;
  * QueryEnvironment as possible while remaining mostly agnostic as to the objects being stored.
  */
 public abstract class AbstractBlazeQueryEnvironment<T>
-    implements QueryEnvironment<T>, AutoCloseable {
+    implements QueryEnvironment<T> {
   protected final ErrorSensingEventHandler eventHandler;
   protected final boolean keepGoing;
   protected final boolean strictScope;
@@ -146,9 +146,6 @@ public abstract class AbstractBlazeQueryEnvironment<T>
     return new QueryEvalResult(!eventHandler.hasErrors(), empty.get());
   }
 
-  @Override
-  public abstract void close();
-
   public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
     return queryExpression;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
index 1b1749fe08..2b8c3e602f 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
@@ -109,13 +109,6 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     this.labelVisitor = new LabelVisitor(packageProvider, dependencyFilter);
   }
 
-  /**
-   * Calling close is optional because {@link BlazeQueryEnvironment} has no resources that need
-   * manual management.
-   */
-  @Override
-  public void close() {}
-
   @Override
   public DigraphQueryEvalResult<Target> evaluateQuery(QueryExpression expr,
       final Callback<Target> callback) throws QueryException, InterruptedException {
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),