Stream result of TargetPattern#eval to a callback instead of returning it directly, and pass a Query callback in when resolving target patterns. This means that the targets a pattern resolves to can be processed incrementally.

This is the fifth step in a series to allow processing large sets of targets in query target patterns via streaming batches rather than all at once. This should improve performance for SkyQueryEnvironment for certain classes of large queries.

--
MOS_MIGRATED_REVID=111696713

3 files changed, 57 insertions, 46 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java b/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java
index 6a23a64386..6587424cc8 100644
--- a/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java
+++ b/src/main/java/com/google/devtools/build/lib/cmdline/TargetPattern.java
@@ -19,11 +19,9 @@ import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.cmdline.LabelValidator.BadLabelException;
 import com.google.devtools.build.lib.cmdline.LabelValidator.PackageAndTarget;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier.RepositoryName;
-import com.google.devtools.build.lib.collect.CompactHashSet;
 import com.google.devtools.build.lib.util.BatchCallback;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.util.StringUtilities;
@@ -34,7 +32,6 @@ import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 
 import javax.annotation.concurrent.Immutable;
 
@@ -146,18 +143,10 @@ public abstract class TargetPattern implements Serializable {
   /**
    * Evaluates the current target pattern and returns the result.
    */
-  public <T> ResolvedTargets<T> eval(TargetPatternResolver<T> resolver)
-      throws TargetParsingException, InterruptedException {
-    final Set<T> results = CompactHashSet.create();
-    BatchCallback<T, RuntimeException> callback =
-        new BatchCallback<T, RuntimeException>() {
-          @Override
-          public void process(Iterable<T> partialResult) {
-            Iterables.addAll(results, partialResult);
-          }
-        };
+  public <T, E extends Exception> void eval(
+      TargetPatternResolver<T> resolver, BatchCallback<T, E> callback)
+      throws TargetParsingException, E, InterruptedException {
     eval(resolver, ImmutableSet.<String>of(), callback);
-    return ResolvedTargets.<T>builder().addAll(results).build();
   }
 
   /**
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 b5bf4f83fa..f6181f29ea 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
@@ -61,6 +61,7 @@ import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.TargetPatternValue;
 import com.google.devtools.build.lib.skyframe.TargetPatternValue.TargetPatternKey;
 import com.google.devtools.build.lib.skyframe.TransitiveTraversalValue;
+import com.google.devtools.build.lib.util.BatchCallback;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.RootedPath;
@@ -77,7 +78,6 @@ import java.util.Collection;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -94,6 +94,9 @@ import javax.annotation.Nullable;
  * even if the full closure isn't needed.
  */
 public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
+  // 10k is likely a good balance between using batch efficiently and not blowing up memory.
+  // TODO(janakr): Unify with RecursivePackageProviderBackedTargetPatternResolver's constant.
+  private static final int BATCH_CALLBACK_SIZE = 10000;
 
   private WalkableGraph graph;
 
@@ -307,9 +310,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
   @Override
   public void eval(QueryExpression expr, Callback<Target> callback)
       throws QueryException, InterruptedException {
-    // 10k is likely a good balance between using batch efficiently and not blowing up memory.
-    BatchStreamedCallback aggregator = new BatchStreamedCallback(callback, 10000,
-        createUniquifier());
+    BatchStreamedCallback aggregator =
+        new BatchStreamedCallback(callback, BATCH_CALLBACK_SIZE, createUniquifier());
     expr.eval(this, aggregator);
     aggregator.processLastPending();
   }
@@ -324,14 +326,50 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     };
   }
 
+  /**
+   * Wraps a {@link Callback<Target>} with three additional filtering mechanisms. First, it
+   * validates the scope of the targets it's given before it passes them to the delegate Callback.
+   * Second, it removes {@link Target}s not in the graph (outside the universe scope). Third, it
+   * wraps the Callback in a {@link BatchStreamedCallback}, which aggregates results into batches of
+   * {@link #BATCH_CALLBACK_SIZE} and also deduplicates elements.
+   */
+  private class FilteringBatchingUniquifyingCallback
+      implements BatchCallback<Target, QueryException> {
+    private final BatchStreamedCallback batchStreamedCallback;
+
+    private FilteringBatchingUniquifyingCallback(Callback<Target> callback) {
+      this.batchStreamedCallback =
+          new BatchStreamedCallback(callback, BATCH_CALLBACK_SIZE, createUniquifier());
+    }
+
+    @Override
+    public void process(Iterable<Target> partialResult)
+        throws QueryException, InterruptedException {
+      Set<Target> targets = CompactHashSet.create();
+      for (Target target : partialResult) {
+        if (validateScope(target.getLabel(), strictScope)) {
+          targets.add(target);
+        }
+      }
+      batchStreamedCallback.process(filterTargetsNotInGraph(targets));
+    }
+
+    private void processLastPending() throws QueryException, InterruptedException {
+      batchStreamedCallback.processLastPending();
+    }
+  }
+
   @Override
   public void getTargetsMatchingPattern(
       QueryExpression owner, String pattern, Callback<Target> callback) throws QueryException {
-    Set<Target> targets = ImmutableSet.of();
     if (precomputedPatterns.containsKey(pattern)) {
       Set<Label> labels = precomputedPatterns.get(pattern);
       if (labels != null) {
-        targets = ImmutableSet.copyOf(makeTargetsFromLabels(labels));
+        try {
+          callback.process(ImmutableSet.copyOf(makeTargetsFromLabels(labels)));
+        } catch (InterruptedException e) {
+          throw new QueryException(owner, e.getMessage());
+        }
       } else {
         TargetParsingException exception;
         try {
@@ -368,27 +406,16 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
             new RecursivePackageProviderBackedTargetPatternResolver(
                 provider, eventHandler, targetPatternKey.getPolicy());
         TargetPattern parsedPattern = targetPatternKey.getParsedPattern();
-        targets = parsedPattern.eval(resolver).getTargets();
+        FilteringBatchingUniquifyingCallback wrapper =
+            new FilteringBatchingUniquifyingCallback(callback);
+        parsedPattern.eval(resolver, wrapper);
+        wrapper.processLastPending();
       } catch (TargetParsingException e) {
         reportBuildFileError(owner, e.getMessage());
       } catch (InterruptedException e) {
         throw new QueryException(owner, e.getMessage());
       }
     }
-
-    // Sets.filter would be more convenient here, but can't deal with exceptions.
-    Iterator<Target> targetIterator = targets.iterator();
-    while (targetIterator.hasNext()) {
-      Target target = targetIterator.next();
-      if (!validateScope(target.getLabel(), strictScope)) {
-        targetIterator.remove();
-      }
-    }
-    try {
-      callback.process(filterTargetsNotInGraph(targets));
-    } catch (InterruptedException e) {
-      throw new QueryException(owner, e.getMessage());
-    }
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java b/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java
index 925fe57aeb..bee90ad0e3 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java
@@ -13,24 +13,19 @@
 // limitations under the License.
 package com.google.devtools.build.lib.query2.engine;
 
+import com.google.devtools.build.lib.util.BatchCallback;
+
 /**
  * Query callback to be called by a {@link QueryExpression} when it has part of the computation
  * result. Assuming the {@code QueryEnvironment} supports it, it would allow the caller
  * to stream the results.
  */
-// TODO(janakr): have this inherit from com.google.devtools.build.lib.util.BatchCallback.
-public interface Callback<T> {
+public interface Callback<T> extends BatchCallback<T, QueryException> {
 
   /**
-   * Called by {@code QueryExpression} when it has been able to compute part of the result.
-   *
-   * <p>Note that this method can be called several times for a QueryExpression. Callers
-   * implementing this method should assume that multiple calls can happen.
-   *
-   * @param partialResult Part of the result. Note that from the caller's perspective, it is
-   * guaranteed that no repeated elements will be returned. However {@code QueryExpression}s calling
-   * the callback do not need to maintain this property, as the {@code QueryEnvironment} should
-   * handle the uniqueness.
+   * According to the {@link BatchCallback} interface, repeated elements may be passed in here.
+   * However, {@code QueryExpression}s calling the callback do not need to maintain this property,
+   * as the {@code QueryEnvironment} should filter out duplicates.
    */
   void process(Iterable<T> partialResult) throws QueryException, InterruptedException;
 }