Unify AbstractBlazeQueryEnvironment#evaluateQuery with its subclass overrides. Also, have AbstractBlazeQueryEnvironment#evaluateQuery take an OutputFormatterCallback instance rather than a Callback instance. This is more sensible since the latter is only intended to be used intra-query, while the former is intended for usage in end-to-end query evaluation. This lets us slightly simplify QueryCommand, by shifting the responsibility for managing the OutputFormatterCallback to AbstractBlazeQueryEnvironment#evaluateQuery.

--
MOS_MIGRATED_REVID=134827588

11 files changed, 210 insertions, 149 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/commands/FetchCommand.java b/src/main/java/com/google/devtools/build/lib/bazel/commands/FetchCommand.java
index d9172ec8ff..6f6d48a5ea 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/commands/FetchCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/commands/FetchCommand.java
@@ -21,7 +21,7 @@ import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.pkgcache.PackageCacheOptions;
 import com.google.devtools.build.lib.query2.AbstractBlazeQueryEnvironment;
-import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Setting;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
@@ -35,6 +35,7 @@ import com.google.devtools.build.lib.util.AbruptExitException;
 import com.google.devtools.build.lib.util.ExitCode;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsProvider;
+import java.io.IOException;
 
 /**
  * Fetches external repositories. Which is so fetch.
@@ -118,10 +119,9 @@ public final class FetchCommand implements BlazeCommand {
 
     // 2. Evaluate expression:
     try {
-      queryEnv.evaluateQuery(expr, new Callback<Target>() {
+      queryEnv.evaluateQuery(expr, new OutputFormatterCallback<Target>() {
         @Override
-        public void process(Iterable<Target> partialResult)
-            throws QueryException, InterruptedException {
+        public void processOutput(Iterable<Target> partialResult) {
           // Throw away the result.
         }
       });
@@ -129,6 +129,9 @@ public final class FetchCommand implements BlazeCommand {
       // Keep consistent with reportBuildFileError()
       env.getReporter().handle(Event.error(e.getMessage()));
       return ExitCode.COMMAND_LINE_ERROR;
+    } catch (IOException e) {
+      // Should be impossible since our OutputFormatterCallback doesn't throw IOException.
+      throw new IllegalStateException(e);
     }
 
     env.getReporter().handle(
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 1962cdfaf4..76a7f36fa7 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
@@ -25,15 +25,18 @@ import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.packages.DependencyFilter;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
-import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionEvalListener;
+import com.google.devtools.build.lib.query2.engine.QueryUtil;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.AggregateAllCallback;
+import com.google.devtools.build.lib.query2.engine.ThreadSafeCallback;
 import com.google.devtools.build.lib.query2.engine.VariableContext;
 import com.google.devtools.build.lib.util.Preconditions;
+import java.io.IOException;
 import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -93,7 +96,20 @@ public abstract class AbstractBlazeQueryEnvironment<T>
   }
 
   /**
-   * Evaluate the specified query expression in this environment.
+   * Used by {@link #evaluateQuery} to evaluate the given {@code expr}. The caller,
+   * {@link #evaluateQuery}, not {@link #evalTopLevelInternal}, is responsible for managing
+   * {@code callback}.
+   */
+  protected void evalTopLevelInternal(QueryExpression expr, OutputFormatterCallback<T> callback)
+      throws QueryException, InterruptedException {
+    eval(expr, VariableContext.<T>empty(), callback);
+  }
+
+  /**
+   * Evaluate the specified query expression in this environment, streaming results to the given
+   * {@code callback}. {@code callback.start()} will be called before query evaluation and
+   * {@code callback.close()} will be unconditionally called at the end of query evaluation
+   * (i.e. regardless of whether it was successful).
    *
    * @return a {@link QueryEvalResult} object that contains the resulting set of targets and a bit
    *   to indicate whether errors occurred during evaluation; note that the
@@ -101,12 +117,12 @@ public abstract class AbstractBlazeQueryEnvironment<T>
    * @throws QueryException if the evaluation failed and {@code --nokeep_going} was in
    *   effect
    */
-  public QueryEvalResult evaluateQuery(QueryExpression expr, final Callback<T> callback)
-      throws QueryException, InterruptedException {
-
-    final AtomicBoolean empty = new AtomicBoolean(true);
+  public QueryEvalResult evaluateQuery(
+      QueryExpression expr,
+      final OutputFormatterCallback<T> callback)
+          throws QueryException, InterruptedException, IOException {
+    EmptinessSensingCallback<T> emptySensingCallback = createEmptinessSensingCallback(callback);
     try (final AutoProfiler p = AutoProfiler.logged("evaluating query", LOG)) {
-
       // In the --nokeep_going case, errors are reported in the order in which the patterns are
       // specified; using a linked hash set here makes sure that the left-most error is reported.
       Set<String> targetPatternSet = new LinkedHashSet<>();
@@ -117,17 +133,24 @@ public abstract class AbstractBlazeQueryEnvironment<T>
         // Unfortunately, by evaluating the patterns in parallel, we lose some location information.
         throw new QueryException(expr, e.getMessage());
       }
+      IOException ioExn = null;
       try {
-        this.eval(expr, VariableContext.<T>empty(), new Callback<T>() {
-          @Override
-          public void process(Iterable<T> partialResult)
-              throws QueryException, InterruptedException {
-            empty.compareAndSet(true, Iterables.isEmpty(partialResult));
-            callback.process(partialResult);
-          }
-        });
+        callback.start();
+        evalTopLevelInternal(expr, emptySensingCallback);
       } catch (QueryException e) {
         throw new QueryException(e, expr);
+      } catch (InterruptedException e) {
+        throw e;
+      } finally {
+        try {
+          callback.close();
+        } catch (IOException e) {
+          // Only throw this IOException if we weren't about to throw a different exception.
+          ioExn = e;
+        }
+      }
+      if (ioExn != null) {
+        throw ioExn;
       }
     }
 
@@ -143,16 +166,62 @@ public abstract class AbstractBlazeQueryEnvironment<T>
       }
     }
 
-    return new QueryEvalResult(!eventHandler.hasErrors(), empty.get());
+    return new QueryEvalResult(!eventHandler.hasErrors(), emptySensingCallback.isEmpty());
+  }
+
+  private static <T> EmptinessSensingCallback<T> createEmptinessSensingCallback(
+      OutputFormatterCallback<T> callback) {
+    return (callback instanceof ThreadSafeCallback)
+        ? new ThreadSafeEmptinessSensingCallback<>(callback)
+        : new EmptinessSensingCallback<>(callback);
+  }
+
+  private static class EmptinessSensingCallback<T> extends OutputFormatterCallback<T> {
+    private final OutputFormatterCallback<T> callback;
+    private final AtomicBoolean empty = new AtomicBoolean(true);
+
+    private EmptinessSensingCallback(OutputFormatterCallback<T> callback) {
+      this.callback = callback;
+    }
+
+    @Override
+    public void start() throws IOException {
+      callback.start();
+    }
+
+    @Override
+    public void processOutput(Iterable<T> partialResult)
+        throws IOException, InterruptedException {
+      empty.compareAndSet(true, Iterables.isEmpty(partialResult));
+      callback.processOutput(partialResult);
+    }
+
+    @Override
+    public void close() throws InterruptedException, IOException {
+      callback.close();
+    }
+
+    boolean isEmpty() {
+      return empty.get();
+    }
+  }
+
+  private static class ThreadSafeEmptinessSensingCallback<T>
+      extends EmptinessSensingCallback<T> implements ThreadSafeCallback<T> {
+    private ThreadSafeEmptinessSensingCallback(OutputFormatterCallback<T> callback) {
+      super(callback);
+      Preconditions.checkState(callback instanceof ThreadSafeCallback);
+    }
   }
 
   public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
     return queryExpression;
   }
 
-  public QueryEvalResult evaluateQuery(String query, Callback<T> callback)
-      throws QueryException, InterruptedException {
-    return evaluateQuery(QueryExpression.parse(query, this), callback);
+  public QueryEvalResult evaluateQuery(String query, OutputFormatterCallback<T> callback)
+      throws QueryException, InterruptedException, IOException {
+    return evaluateQuery(
+        QueryExpression.parse(query, this), callback);
   }
 
   @Override
@@ -189,7 +258,7 @@ public abstract class AbstractBlazeQueryEnvironment<T>
       // Will skip the target and keep going if -k is specified.
       reportBuildFileError(caller, e.getMessage());
     }
-    AggregateAllCallback<T> aggregatingCallback = new AggregateAllCallback<>();
+    AggregateAllCallback<T> aggregatingCallback = QueryUtil.newAggregateAllCallback();
     getTargetsMatchingPattern(caller, pattern, aggregatingCallback);
     return aggregatingCallback.getResult();
   }
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 2b8c3e602f..cbff9d9628 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
@@ -16,7 +16,6 @@ package com.google.devtools.build.lib.query2;
 import com.google.common.base.Function;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
@@ -38,10 +37,12 @@ import com.google.devtools.build.lib.pkgcache.TargetProvider;
 import com.google.devtools.build.lib.pkgcache.TransitivePackageLoader;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.DigraphQueryEvalResult;
+import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionEvalListener;
+import com.google.devtools.build.lib.query2.engine.QueryUtil;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.AbstractUniquifier;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.AggregateAllCallback;
 import com.google.devtools.build.lib.query2.engine.SkyframeRestartQueryException;
@@ -49,6 +50,7 @@ import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.lib.query2.engine.VariableContext;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
@@ -58,7 +60,6 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * The environment of a Blaze query. Not thread-safe.
@@ -110,23 +111,15 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   }
 
   @Override
-  public DigraphQueryEvalResult<Target> evaluateQuery(QueryExpression expr,
-      final Callback<Target> callback) throws QueryException, InterruptedException {
-    // Some errors are reported as QueryExceptions and others as ERROR events (if --keep_going). The
-    // result is set to have an error iff there were errors emitted during the query, so we reset
-    // errors here.
+  public DigraphQueryEvalResult<Target> evaluateQuery(
+      QueryExpression expr,
+      final OutputFormatterCallback<Target> callback)
+          throws QueryException, InterruptedException, IOException {
     eventHandler.resetErrors();
-    final AtomicBoolean empty = new AtomicBoolean(true);
     resolvedTargetPatterns.clear();
-    QueryEvalResult queryEvalResult = super.evaluateQuery(expr, new Callback<Target>() {
-      @Override
-      public void process(Iterable<Target> partialResult)
-          throws QueryException, InterruptedException {
-        empty.compareAndSet(true, Iterables.isEmpty(partialResult));
-        callback.process(partialResult);
-      }
-    });
-    return new DigraphQueryEvalResult<>(queryEvalResult.getSuccess(), empty.get(), graph);
+    QueryEvalResult queryEvalResult = super.evaluateQuery(expr, callback);
+    return new DigraphQueryEvalResult<>(
+        queryEvalResult.getSuccess(), queryEvalResult.isEmpty(), graph);
   }
 
   @Override
@@ -283,7 +276,7 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   @Override
   public void eval(QueryExpression expr, VariableContext<Target> context, Callback<Target> callback)
       throws QueryException, InterruptedException {
-    AggregateAllCallback<Target> aggregator = new AggregateAllCallback<>();
+    AggregateAllCallback<Target> aggregator = QueryUtil.newAggregateAllCallback();
     expr.eval(this, context, aggregator);
     callback.process(aggregator.getResult());
   }
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 2913d1e249..b920787f2b 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
@@ -52,6 +52,7 @@ import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.query2.engine.AllRdepsFunction;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.FunctionExpression;
+import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
@@ -87,6 +88,7 @@ import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.WalkableGraph;
 import com.google.devtools.build.skyframe.WalkableGraph.WalkableGraphFactory;
+import java.io.IOException;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -102,7 +104,6 @@ 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;
 
@@ -309,8 +310,24 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   }
 
   @Override
-  public QueryEvalResult evaluateQuery(QueryExpression expr, Callback<Target> callback)
-      throws QueryException, InterruptedException {
+  protected void evalTopLevelInternal(
+      QueryExpression expr, OutputFormatterCallback<Target> callback)
+          throws QueryException, InterruptedException {
+    try {
+      super.evalTopLevelInternal(expr, callback);
+    } 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);
+    }
+  }
+
+  @Override
+  public QueryEvalResult evaluateQuery(
+      QueryExpression expr, OutputFormatterCallback<Target> callback)
+          throws QueryException, InterruptedException, IOException {
     // Some errors are reported as QueryExceptions and others as ERROR events (if --keep_going). The
     // result is set to have an error iff there were errors emitted during the query, so we reset
     // errors here.
@@ -324,47 +341,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     //
     // This flushes the batched callback prior to constructing the QueryEvalResult in the unlikely
     // case of a race between the original callback and the eventHandler.
-    final BatchStreamedCallback aggregator =
-        new BatchStreamedCallback(callback, BATCH_CALLBACK_SIZE);
-
-    final AtomicBoolean empty = new AtomicBoolean(true);
-    ThreadSafeCallback<Target> callbackWithEmptyCheck =
-        new ThreadSafeCallback<Target>() {
-          @Override
-          public void process(Iterable<Target> partialResult)
-              throws QueryException, InterruptedException {
-            empty.compareAndSet(true, Iterables.isEmpty(partialResult));
-            aggregator.process(partialResult);
-          }
-        };
-    try (final AutoProfiler p = AutoProfiler.logged("evaluating query", LOG)) {
-      try {
-        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();
-    }
-
-    if (eventHandler.hasErrors()) {
-      if (!keepGoing) {
-        // This case represents loading-phase errors reported during evaluation
-        // of target patterns that don't cause evaluation to fail per se.
-        throw new QueryException(
-            "Evaluation of query \"" + expr + "\" failed due to BUILD file errors");
-      } else {
-        eventHandler.handle(
-            Event.warn("--keep_going specified, ignoring errors.  " + "Results may be inaccurate"));
-      }
-    }
-
-    return new QueryEvalResult(!eventHandler.hasErrors(), empty.get());
+    BatchStreamedCallback batchCallback = new BatchStreamedCallback(callback, BATCH_CALLBACK_SIZE);
+    return super.evaluateQuery(expr, batchCallback);
   }
 
   private Map<Target, Collection<Target>> makeTargetsMap(Map<SkyKey, Iterable<SkyKey>> input)
@@ -1025,38 +1003,50 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
    * call the wrapped {@code callback} concurrently.
    */
   @ThreadSafe
-  private static class BatchStreamedCallback implements ThreadSafeCallback<Target> {
+  private static class BatchStreamedCallback
+      extends OutputFormatterCallback<Target> implements ThreadSafeCallback<Target> {
 
-    private final Callback<Target> callback;
+    private final OutputFormatterCallback<Target> callback;
     private final ThreadSafeUniquifier<Target> uniquifier =
         new ThreadSafeTargetUniquifier(DEFAULT_THREAD_COUNT);
     private final Object pendingLock = new Object();
     private List<Target> pending = new ArrayList<>();
     private int batchThreshold;
 
-    private BatchStreamedCallback(Callback<Target> callback, int batchThreshold) {
+    private BatchStreamedCallback(OutputFormatterCallback<Target> callback, int batchThreshold) {
       this.callback = callback;
       this.batchThreshold = batchThreshold;
     }
 
     @Override
-    public void process(Iterable<Target> partialResult)
-        throws QueryException, InterruptedException {
+    public void start() throws IOException {
+      callback.start();
+    }
+
+    @Override
+    public void processOutput(Iterable<Target> partialResult)
+        throws IOException, InterruptedException {
       ImmutableList<Target> uniquifiedTargets = uniquifier.unique(partialResult);
       synchronized (pendingLock) {
         Preconditions.checkNotNull(pending, "Reuse of the callback is not allowed");
         pending.addAll(uniquifiedTargets);
         if (pending.size() >= batchThreshold) {
-          callback.process(pending);
+          callback.processOutput(pending);
           pending = new ArrayList<>();
         }
       }
     }
 
-    private void processLastPending() throws QueryException, InterruptedException {
+    @Override
+    public void close() throws IOException, InterruptedException {
+      processLastPending();
+      callback.close();
+    }
+
+    private void processLastPending() throws IOException, InterruptedException {
       synchronized (pendingLock) {
         if (!pending.isEmpty()) {
-          callback.process(pending);
+          callback.processOutput(pending);
           pending = null;
         }
       }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/OutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/engine/OutputFormatterCallback.java
index 546b0d51bb..c0d80f8c6c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/OutputFormatterCallback.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/OutputFormatterCallback.java
@@ -14,15 +14,12 @@
 package com.google.devtools.build.lib.query2.engine;
 
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadCompatible;
-
-import java.io.Closeable;
 import java.io.IOException;
-
 import javax.annotation.Nullable;
 
 /** A callback that can receive a finish event when there are no more partial results */
 @ThreadCompatible
-public abstract class OutputFormatterCallback<T> implements Callback<T>, Closeable {
+public abstract class OutputFormatterCallback<T> implements Callback<T> {
 
   private IOException ioException;
 
@@ -31,15 +28,16 @@ public abstract class OutputFormatterCallback<T> implements Callback<T>, Closeab
    *
    * <p>It should be used for opening resources or sending a header to the output.
    */
-  public void start() throws IOException { }
+  public void start() throws IOException {
+  }
 
   /**
-   * Same as start but for closing resources or writting a footer.
+   * Same as start but for closing resources or writing a footer.
    *
    * <p>Will be called even in the case of an error.
    */
-  @Override
-  public void close() throws IOException{}
+  public void close() throws InterruptedException, IOException {
+  }
 
   /**
    * Note that {@link Callback} interface does not throw IOExceptions. What this implementation does
@@ -57,7 +55,7 @@ public abstract class OutputFormatterCallback<T> implements Callback<T>, Closeab
     }
   }
 
-  protected abstract void processOutput(Iterable<T> partialResult)
+  public abstract void processOutput(Iterable<T> partialResult)
       throws IOException, InterruptedException;
 
   @Nullable
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
index 4447488a1a..2d7be2ed74 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
@@ -26,37 +26,64 @@ public final class QueryUtil {
 
   private QueryUtil() { }
 
-  /** A callback that can aggregate all the partial results in one set */
-  public static class AggregateAllCallback<T> implements Callback<T> {
+  /** A {@link Callback} that can aggregate all the partial results into one set. */
+  public interface AggregateAllCallback<T> extends Callback<T> {
+    Set<T> getResult();
+  }
+
+  /** A {@link OutputFormatterCallback} that can aggregate all the partial results into one set. */
+  public abstract static class AggregateAllOutputFormatterCallback<T>
+      extends OutputFormatterCallback<T> implements AggregateAllCallback<T>  {
+  }
 
-    private final CompactHashSet<T> result = CompactHashSet.create();
+  private static class AggregateAllOutputFormatterCallbackImpl<T>
+      extends AggregateAllOutputFormatterCallback<T> {
+    private final Set<T> result = CompactHashSet.create();
 
     @Override
-    public void process(Iterable<T> partialResult) throws QueryException, InterruptedException {
+    public final void processOutput(Iterable<T> partialResult) {
       Iterables.addAll(result, partialResult);
     }
 
+    @Override
     public Set<T> getResult() {
       return result;
     }
+  }
 
-    @Override
-    public String toString() {
-      return "Aggregate all: " + result;
-    }
+  /**
+   * Returns a fresh {@link AggregateAllOutputFormatterCallback} that can aggregate all the partial
+   * results into one set.
+   *
+   * <p>Intended to be used by top-level evaluation of {@link QueryExpression}s; contrast with
+   * {@link #newAggregateAllCallback}.
+   */
+  public static <T> AggregateAllOutputFormatterCallback<T>
+      newAggregateAllOutputFormatterCallback() {
+    return new AggregateAllOutputFormatterCallbackImpl<>();
+  }
+
+  /**
+   * Returns a fresh {@link AggregateAllCallback}.
+   *
+   * <p>Intended to be used by {@link QueryExpression} implementations; contrast with
+   * {@link #newAggregateAllOutputFormatterCallback}.
+   */
+  public static <T> AggregateAllCallback<T> newAggregateAllCallback() {
+    return new AggregateAllOutputFormatterCallbackImpl<>();
   }
 
   /**
    * Fully evaluate a {@code QueryExpression} and return a set with all the results.
    *
-   * <p>Should ony be used by QueryExpressions when it is the only way of achieving correctness.
+   * <p>Should only be used by QueryExpressions when it is the only way of achieving correctness.
    */
   public static <T> Set<T> evalAll(
       QueryEnvironment<T> env, VariableContext<T> context, QueryExpression expr)
           throws QueryException, InterruptedException {
-    AggregateAllCallback<T> callback = new AggregateAllCallback<>();
+    AggregateAllCallback<T> callback = newAggregateAllCallback();
     env.eval(expr, context, callback);
-    return callback.result;
+    return callback.getResult();
   }
 
   /** A trivial {@link Uniquifier} base class. */
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
index d7c6eb3bb5..5727546356 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
@@ -269,7 +269,7 @@ public abstract class OutputFormatter implements Serializable {
         OutputStream out, final QueryOptions options) {
       return new TextOutputFormatterCallback<Target>(out) {
         @Override
-        protected void processOutput(Iterable<Target> partialResult) {
+        public void processOutput(Iterable<Target> partialResult) {
           for (Target target : partialResult) {
             if (showKind) {
               printStream.print(target.getTargetKind());
@@ -319,7 +319,7 @@ public abstract class OutputFormatter implements Serializable {
         private final Set<String> packageNames = Sets.newTreeSet();
 
         @Override
-        protected void processOutput(Iterable<Target> partialResult) {
+        public void processOutput(Iterable<Target> partialResult) {
 
           for (Target target : partialResult) {
             packageNames.add(target.getLabel().getPackageName());
@@ -363,7 +363,7 @@ public abstract class OutputFormatter implements Serializable {
       return new TextOutputFormatterCallback<Target>(out) {
 
         @Override
-        protected void processOutput(Iterable<Target> partialResult) {
+        public void processOutput(Iterable<Target> partialResult) {
           final String lineTerm = options.getLineTerminator();
           for (Target target : partialResult) {
             Location location = target.getLocation();
@@ -443,7 +443,7 @@ public abstract class OutputFormatter implements Serializable {
         }
 
         @Override
-        protected void processOutput(Iterable<Target> partialResult) {
+        public void processOutput(Iterable<Target> partialResult) {
 
           for (Target target : partialResult) {
             Rule rule = target.getAssociatedRule();
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
index 8b388bbd61..3852ce294b 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
@@ -100,7 +100,7 @@ public class ProtoOutputFormatter extends AbstractUnorderedFormatter {
       }
 
       @Override
-      protected void processOutput(Iterable<Target> partialResult)
+      public void processOutput(Iterable<Target> partialResult)
           throws IOException, InterruptedException {
 
         for (Target target : partialResult) {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
index 5727723604..2d7348fd30 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
@@ -90,7 +90,7 @@ class XmlOutputFormatter extends AbstractUnorderedFormatter {
       }
 
       @Override
-      protected void processOutput(Iterable<Target> partialResult)
+      public void processOutput(Iterable<Target> partialResult)
           throws IOException, InterruptedException {
         for (Target target : partialResult) {
           queryElem.appendChild(createTargetElement(doc, target));
diff --git a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
index 3620b8fbda..1fa0af0fcc 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
@@ -56,7 +56,8 @@ import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunctio
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Setting;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionEvalListener;
-import com.google.devtools.build.lib.query2.engine.QueryUtil.AggregateAllCallback;
+import com.google.devtools.build.lib.query2.engine.QueryUtil;
+import com.google.devtools.build.lib.query2.engine.QueryUtil.AggregateAllOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.SkyframeRestartQueryException;
 import com.google.devtools.build.lib.query2.output.OutputFormatter;
 import com.google.devtools.build.lib.query2.output.QueryOptions;
@@ -261,7 +262,8 @@ public class GenQuery implements RuleConfiguredTargetFactory {
 
     DigraphQueryEvalResult<Target> queryResult;
     OutputFormatter formatter;
-    AggregateAllCallback<Target> targets = new AggregateAllCallback<>();
+    AggregateAllOutputFormatterCallback<Target> targets =
+        QueryUtil.newAggregateAllOutputFormatterCallback();
     try {
       Set<Setting> settings = queryOptions.toSettings();
 
@@ -305,6 +307,8 @@ public class GenQuery implements RuleConfiguredTargetFactory {
     } catch (QueryException e) {
       ruleContext.ruleError("query failed: " + e.getMessage());
       return null;
+    } catch (IOException e) {
+      throw new RuntimeException(e);
     }
 
     ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
index a84f078429..6c689d2baa 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
@@ -18,9 +18,7 @@ import static com.google.devtools.build.lib.packages.Rule.ALL_LABELS;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.analysis.NoBuildEvent;
-import com.google.devtools.build.lib.collect.CompactHashSet;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.pkgcache.PackageCacheOptions;
@@ -31,6 +29,8 @@ import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionEvalListener;
+import com.google.devtools.build.lib.query2.engine.QueryUtil;
+import com.google.devtools.build.lib.query2.engine.QueryUtil.AggregateAllOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.output.OutputFormatter;
 import com.google.devtools.build.lib.query2.output.OutputFormatter.StreamedFormatter;
 import com.google.devtools.build.lib.query2.output.QueryOptions;
@@ -163,11 +163,10 @@ public final class QueryCommand implements BlazeCommand {
           queryOptions.aspectDeps.createResolver(env.getPackageManager(), env.getReporter()));
       callback = streamedFormatter.createStreamCallback(out, queryOptions, queryEnv);
     } else {
-      callback = new AggregateAllOutputFormatterCallback<>();
+      callback = QueryUtil.newAggregateAllOutputFormatterCallback();
     }
     boolean catastrophe = true;
     try {
-      callback.start();
       result = queryEnv.evaluateQuery(expr, callback);
       catastrophe = false;
     } catch (QueryException e) {
@@ -195,13 +194,6 @@ public final class QueryCommand implements BlazeCommand {
     } finally {
       if (!catastrophe) {
         try {
-          callback.close();
-        } catch (IOException e) {
-          env.getReporter().handle(Event.error("I/O error: " + e.getMessage()));
-          return ExitCode.LOCAL_ENVIRONMENTAL_ERROR;
-        }
-
-        try {
           out.flush();
         } catch (IOException e) {
           env.getReporter().handle(
@@ -218,7 +210,7 @@ public final class QueryCommand implements BlazeCommand {
       // 3. Output results:
       try {
         Set<Target> targets =
-            ((AggregateAllOutputFormatterCallback<Target>) callback).getOutput();
+            ((AggregateAllOutputFormatterCallback<Target>) callback).getResult();
         QueryOutputUtils.output(
             queryOptions,
             result,
@@ -290,19 +282,4 @@ public final class QueryCommand implements BlazeCommand {
             QueryExpressionEvalListener.NullListener.<Target>instance(),
             env.getPackageManager().getPackagePath());
   }
-
-  private static class AggregateAllOutputFormatterCallback<T> extends OutputFormatterCallback<T> {
-
-    private Set<T> output = CompactHashSet.create();
-
-    @Override
-    protected final void processOutput(Iterable<T> partialResult)
-        throws IOException, InterruptedException {
-      Iterables.addAll(output, partialResult);
-    }
-
-    public Set<T> getOutput() {
-      return output;
-    }
-  }
 }