1 files changed, 12 insertions, 22 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java b/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
index 375e7e2eca..1948cb0c53 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
@@ -23,7 +23,6 @@ import com.google.common.collect.Iterables;
 import com.google.common.collect.ListMultimap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
-import com.google.common.collect.Streams;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.collect.compacthashset.CompactHashSet;
@@ -84,28 +83,28 @@ class ParallelSkyQueryUtils {
   static void getRBuildFilesParallel(
       SkyQueryEnvironment env,
       Collection<PathFragment> fileIdentifiers,
-      Callback<Target> callback,
-      MultisetSemaphore<PackageIdentifier> packageSemaphore)
-          throws QueryException, InterruptedException {
+      Callback<Target> callback) throws QueryException, InterruptedException {
     Uniquifier<SkyKey> keyUniquifier = env.createSkyKeyUniquifier();
     RBuildFilesVisitor visitor =
-        new RBuildFilesVisitor(env, keyUniquifier, callback, packageSemaphore);
+        new RBuildFilesVisitor(env, keyUniquifier, callback);
     visitor.visitAndWaitForCompletion(env.getSkyKeysForFileFragments(fileIdentifiers));
   }
 
   /** A helper class that computes 'rbuildfiles(<blah>)' via BFS. */
   private static class RBuildFilesVisitor extends ParallelVisitor<SkyKey, Target> {
+    // Each target in the full output of 'rbuildfiles' corresponds to BUILD file InputFile of a
+    // unique package. So the processResultsBatchSize we choose to pass to the ParallelVisitor ctor
+    // influences how many packages each leaf task doing processPartialResults will have to
+    // deal with at once. A value of 100 was chosen experimentally.
+    private static final int PROCESS_RESULTS_BATCH_SIZE = 100;
     private final SkyQueryEnvironment env;
-    private final MultisetSemaphore<PackageIdentifier> packageSemaphore;
 
     private RBuildFilesVisitor(
         SkyQueryEnvironment env,
         Uniquifier<SkyKey> uniquifier,
-        Callback<Target> callback,
-        MultisetSemaphore<PackageIdentifier> packageSemaphore) {
-      super(uniquifier, callback, VISIT_BATCH_SIZE);
+        Callback<Target> callback) {
+      super(uniquifier, callback, VISIT_BATCH_SIZE, PROCESS_RESULTS_BATCH_SIZE);
       this.env = env;
-      this.packageSemaphore = packageSemaphore;
     }
 
     @Override
@@ -135,17 +134,7 @@ class ParallelSkyQueryUtils {
     protected void processPartialResults(
         Iterable<SkyKey> keysToUseForResult, Callback<Target> callback)
         throws QueryException, InterruptedException {
-      Set<PackageIdentifier> pkgIdsNeededForResult =
-          Streams.stream(keysToUseForResult)
-              .map(SkyQueryEnvironment.PACKAGE_SKYKEY_TO_PACKAGE_IDENTIFIER)
-              .collect(toImmutableSet());
-      packageSemaphore.acquireAll(pkgIdsNeededForResult);
-      try {
-        callback.process(SkyQueryEnvironment.getBuildFilesForPackageValues(
-            env.graph.getSuccessfulValues(keysToUseForResult).values()));
-      } finally {
-        packageSemaphore.releaseAll(pkgIdsNeededForResult);
-      }
+      env.getBuildFileTargetsForPackageKeysAndProcessViaCallback(keysToUseForResult, callback);
     }
 
     @Override
@@ -168,6 +157,7 @@ class ParallelSkyQueryUtils {
    */
   private static class AllRdepsUnboundedVisitor
       extends ParallelVisitor<Pair<SkyKey, SkyKey>, Target> {
+    private static final int PROCESS_RESULTS_BATCH_SIZE = SkyQueryEnvironment.BATCH_CALLBACK_SIZE;
     private final SkyQueryEnvironment env;
     private final MultisetSemaphore<PackageIdentifier> packageSemaphore;
 
@@ -176,7 +166,7 @@ class ParallelSkyQueryUtils {
         Uniquifier<Pair<SkyKey, SkyKey>> uniquifier,
         Callback<Target> callback,
         MultisetSemaphore<PackageIdentifier> packageSemaphore) {
-      super(uniquifier, callback, VISIT_BATCH_SIZE);
+      super(uniquifier, callback, VISIT_BATCH_SIZE, PROCESS_RESULTS_BATCH_SIZE);
       this.env = env;
       this.packageSemaphore = packageSemaphore;
     }