Stream results of targets below directory to a callback rather than returning a ResolvedTargets set.

This is the first 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 be a functional no-op.

--
MOS_MIGRATED_REVID=111609309

6 files changed, 112 insertions, 39 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 3b93221acf..60515979ff 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,9 +19,12 @@ 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;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -31,6 +34,7 @@ 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;
 
@@ -445,9 +449,22 @@ public abstract class TargetPattern implements Serializable {
     public <T> ResolvedTargets<T> eval(TargetPatternResolver<T> resolver,
         ImmutableSet<String> excludedSubdirectories)
         throws TargetParsingException, InterruptedException {
-      return resolver.findTargetsBeneathDirectory(
-          directory.getRepository(), getOriginalPattern(),
-          directory.getPackageFragment().getPathString(), rulesOnly, excludedSubdirectories);
+      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);
+            }
+          };
+      resolver.findTargetsBeneathDirectory(
+          directory.getRepository(),
+          getOriginalPattern(),
+          directory.getPackageFragment().getPathString(),
+          rulesOnly,
+          excludedSubdirectories,
+          callback);
+      return ResolvedTargets.<T>builder().addAll(results).build();
     }
 
     @Override
diff --git a/src/main/java/com/google/devtools/build/lib/cmdline/TargetPatternResolver.java b/src/main/java/com/google/devtools/build/lib/cmdline/TargetPatternResolver.java
index 6e273b43ff..d98f78f7ef 100644
--- a/src/main/java/com/google/devtools/build/lib/cmdline/TargetPatternResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/cmdline/TargetPatternResolver.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.cmdline;
 
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier.RepositoryName;
+import com.google.devtools.build.lib.util.BatchCallback;
 
 /**
  * A callback interface that is used during the process of converting target patterns (such as
@@ -57,11 +58,12 @@ public interface TargetPatternResolver<T> {
       throws TargetParsingException, InterruptedException;
 
   /**
-   * Returns the set containing the targets found below the given {@code directory}. Conceptually,
-   * this method should look for all packages that start with the {@code directory} (as a proper
-   * prefix directory, i.e., "foo/ba" is not a proper prefix of "foo/bar/"), and then collect all
-   * targets in each such package (subject to {@code rulesOnly}) as if calling {@link
-   * #getTargetsInPackage}. The specified directory is not necessarily a valid package name.
+   * Computes the set containing the targets found below the given {@code directory}, passing it in
+   * batches to {@code callback}. Conceptually, this method should look for all packages that start
+   * with the {@code directory} (as a proper prefix directory, i.e., "foo/ba" is not a proper prefix
+   * of "foo/bar/"), and then collect all targets in each such package (subject to
+   * {@code rulesOnly}) as if calling {@link #getTargetsInPackage}. The specified directory is not
+   * necessarily a valid package name.
    *
    * <p>Note that the {@code directory} can be empty, which corresponds to the "//..." pattern.
    * Implementations may choose not to support this case and throw an {@link
@@ -76,11 +78,17 @@ public interface TargetPatternResolver<T> {
    * @param rulesOnly whether to return rules only
    * @param excludedSubdirectories a set of transitive subdirectories beneath {@code directory}
    *    to ignore
+   * @param callback the callback to receive the result, possibly in multiple batches.
    * @throws TargetParsingException under implementation-specific failure conditions
    */
-  ResolvedTargets<T> findTargetsBeneathDirectory(RepositoryName repository, String originalPattern,
-      String directory, boolean rulesOnly, ImmutableSet<String> excludedSubdirectories)
-      throws TargetParsingException, InterruptedException;
+  <E extends Exception> void findTargetsBeneathDirectory(
+      RepositoryName repository,
+      String originalPattern,
+      String directory,
+      boolean rulesOnly,
+      ImmutableSet<String> excludedSubdirectories,
+      BatchCallback<T, E> callback)
+      throws TargetParsingException, E, InterruptedException;
 
   /**
    * Returns true, if and only if the given package identifier corresponds to a package, i.e., a
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 ef9c3cde18..925fe57aeb 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
@@ -18,6 +18,7 @@ package com.google.devtools.build.lib.query2.engine;
  * 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> {
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PrepareDepsOfPatternFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/PrepareDepsOfPatternFunction.java
index 889cbbb0e5..5b02ecc041 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PrepareDepsOfPatternFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PrepareDepsOfPatternFunction.java
@@ -33,6 +33,7 @@ import com.google.devtools.build.lib.pkgcache.FilteringPolicy;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
 import com.google.devtools.build.lib.pkgcache.TargetPatternResolverUtil;
 import com.google.devtools.build.lib.skyframe.EnvironmentBackedRecursivePackageProvider.MissingDepException;
+import com.google.devtools.build.lib.util.BatchCallback;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -207,10 +208,14 @@ public class PrepareDepsOfPatternFunction implements SkyFunction {
     }
 
     @Override
-    public ResolvedTargets<Void> findTargetsBeneathDirectory(RepositoryName repository,
-        String originalPattern, String directory, boolean rulesOnly,
-        ImmutableSet<String> excludedSubdirectories)
-        throws TargetParsingException, InterruptedException {
+    public <E extends Exception> void findTargetsBeneathDirectory(
+        RepositoryName repository,
+        String originalPattern,
+        String directory,
+        boolean rulesOnly,
+        ImmutableSet<String> excludedSubdirectories,
+        BatchCallback<Void, E> callback)
+        throws TargetParsingException, E, InterruptedException {
       FilteringPolicy policy =
           rulesOnly ? FilteringPolicies.RULES_ONLY : FilteringPolicies.NO_FILTER;
       ImmutableSet<PathFragment> excludedPathFragments =
@@ -239,7 +244,6 @@ public class PrepareDepsOfPatternFunction implements SkyFunction {
           throw new MissingDepException();
         }
       }
-      return ResolvedTargets.empty();
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java b/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java
index d3f758f028..2641056e5a 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/RecursivePackageProviderBackedTargetPatternResolver.java
@@ -33,8 +33,11 @@ import com.google.devtools.build.lib.pkgcache.FilteringPolicies;
 import com.google.devtools.build.lib.pkgcache.FilteringPolicy;
 import com.google.devtools.build.lib.pkgcache.RecursivePackageProvider;
 import com.google.devtools.build.lib.pkgcache.TargetPatternResolverUtil;
+import com.google.devtools.build.lib.util.BatchCallback;
 import com.google.devtools.build.lib.vfs.PathFragment;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -43,6 +46,9 @@ import java.util.Map;
 public class RecursivePackageProviderBackedTargetPatternResolver
     implements TargetPatternResolver<Target> {
 
+  // TODO(janakr): Move this to a more generic place and unify with SkyQueryEnvironment's value?
+  private static final int MAX_PACKAGES_BULK_GET = 10000;
+
   private final RecursivePackageProvider recursivePackageProvider;
   private final EventHandler eventHandler;
   private final FilteringPolicy policy;
@@ -154,17 +160,20 @@ public class RecursivePackageProviderBackedTargetPatternResolver
   }
 
   @Override
-  public ResolvedTargets<Target> findTargetsBeneathDirectory(final RepositoryName repository,
-      String originalPattern, String directory, boolean rulesOnly,
-      ImmutableSet<String> excludedSubdirectories)
-      throws TargetParsingException, InterruptedException {
+  public <E extends Exception> void findTargetsBeneathDirectory(
+      final RepositoryName repository,
+      String originalPattern,
+      String directory,
+      boolean rulesOnly,
+      ImmutableSet<String> excludedSubdirectories,
+      BatchCallback<Target, E> callback)
+      throws TargetParsingException, E, InterruptedException {
     FilteringPolicy actualPolicy = rulesOnly
         ? FilteringPolicies.and(FilteringPolicies.RULES_ONLY, policy)
         : policy;
     ImmutableSet<PathFragment> excludedPathFragments =
         TargetPatternResolverUtil.getPathFragments(excludedSubdirectories);
     PathFragment pathFragment = TargetPatternResolverUtil.getPathFragment(directory);
-    ResolvedTargets.Builder<Target> targetBuilder = ResolvedTargets.builder();
     Iterable<PathFragment> packagesUnderDirectory =
         recursivePackageProvider.getPackagesUnderDirectory(
             repository, pathFragment, excludedPathFragments);
@@ -176,28 +185,30 @@ public class RecursivePackageProviderBackedTargetPatternResolver
                 return PackageIdentifier.create(repository, path);
               }
             });
-    for (ResolvedTargets<Target> targets : bulkGetTargetsInPackage(originalPattern, pkgIds,
-            FilteringPolicies.NO_FILTER).values()) {
-      targetBuilder.merge(targets);
+    boolean foundTarget = false;
+    // For very large sets of packages, we may not want to process all of them at once, so we split
+    // into batches.
+    for (Iterable<PackageIdentifier> pkgIdBatch :
+        Iterables.partition(pkgIds, MAX_PACKAGES_BULK_GET)) {
+      for (ResolvedTargets<Target> targets :
+          bulkGetTargetsInPackage(originalPattern, pkgIdBatch, FilteringPolicies.NO_FILTER)
+              .values()) {
+        List<Target> filteredTargets = new ArrayList<>(targets.getTargets().size());
+        for (Target target : targets.getTargets()) {
+          // Perform the no-targets-found check before applying the filtering policy so we only
+          // return the error if the input directory's subtree really contains no targets.
+          foundTarget = true;
+          if (actualPolicy.shouldRetain(target, false)) {
+            filteredTargets.add(target);
+          }
+        }
+        callback.process(filteredTargets);
+      }
     }
 
-    // Perform the no-targets-found check before applying the filtering policy so we only return the
-    // error if the input directory's subtree really contains no targets.
-    if (targetBuilder.isEmpty()) {
+    if (!foundTarget) {
       throw new TargetParsingException("no targets found beneath '" + pathFragment + "'");
     }
-    ResolvedTargets<Target> prefilteredTargets = targetBuilder.build();
-     
-    ResolvedTargets.Builder<Target> filteredBuilder = ResolvedTargets.builder();
-    if (prefilteredTargets.hasError()) {
-      filteredBuilder.setError();
-    }
-    for (Target target : prefilteredTargets.getTargets()) {
-      if (actualPolicy.shouldRetain(target, false)) {
-        filteredBuilder.add(target);
-      }
-    }
-    return filteredBuilder.build();
   }
 }
 
diff --git a/src/main/java/com/google/devtools/build/lib/util/BatchCallback.java b/src/main/java/com/google/devtools/build/lib/util/BatchCallback.java
new file mode 100644
index 0000000000..70cc5bff50
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/util/BatchCallback.java
@@ -0,0 +1,32 @@
+// Copyright 2016 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.util;
+
+/**
+ * Callback to be invoked when part of a result has been computed. Allows a client interested in
+ * the result to process it as it is computed, for instance by streaming it, if it is too big to
+ * fit in memory.
+ */
+public interface BatchCallback<T, E extends Exception> {
+  /**
+   * Called when part of a result has been computed.
+   *
+   * <p>Note that this method can be called several times for a single {@code BatchCallback}.
+   * Implementations should assume that multiple calls can happen.
+   *
+   * @param partialResult Part of the result. May contain duplicates, either in the same call or
+   * across calls.
+   */
+  void process(Iterable<T> partialResult) throws E, InterruptedException;
+}