In AllRdepsUnboundedVisitor, group rdeps by package before partitioning them into batches. This way nodes corresponding to targets in the same package are likely to be in the same batch.

--
MOS_MIGRATED_REVID=135684088

2 files changed, 18 insertions, 3 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 ec7474dd4f..fb9841d9af 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
@@ -13,9 +13,11 @@
 // limitations under the License.
 package com.google.devtools.build.lib.query2;
 
+import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.collect.CompactHashSet;
 import com.google.devtools.build.lib.concurrent.ForkJoinQuiescingExecutor;
 import com.google.devtools.build.lib.concurrent.MoreFutures;
@@ -175,8 +177,21 @@ class ParallelSkyQueryUtils {
       // (i)  It only returns rdeps that have corresponding Targets.
       // (ii) It only returns rdeps whose corresponding Targets have a valid dependency edge to
       //      their direct dep.
-      Iterable<SkyKey> keysToVisit = SkyQueryEnvironment.makeTransitiveTraversalKeysStrict(
-          env.getReverseDepsOfTransitiveTraversalKeys(keys));
+      Iterable<Target> rdepTargets = env.getReverseDepsOfTransitiveTraversalKeys(keys);
+      // Group the targets by package - this way when computeImpl splits these targets into batches,
+      // targets in the same package are likely to be in the same batch.
+      ArrayListMultimap<PackageIdentifier, SkyKey> rdepKeysByPackage = ArrayListMultimap.create();
+      for (Target rdepTarget : rdepTargets) {
+        rdepKeysByPackage.put(
+            rdepTarget.getLabel().getPackageIdentifier(),
+            SkyQueryEnvironment.TARGET_TO_SKY_KEY.apply(rdepTarget));
+      }
+      // A couple notes here:
+      // (i)  ArrayListMultimap#values returns the values grouped by key, which is exactly what we
+      //      want.
+      // (ii) ArrayListMultimap#values returns a Collection view, so we make a copy to avoid
+      //      accidentally retaining the entire ArrayListMultimap object.
+      Iterable<SkyKey> keysToVisit = ImmutableList.copyOf(rdepKeysByPackage.values());
       return new Visit(
           /*keysToUseForResult=*/ keys,
           /*keysToVisit=*/ keysToVisit);
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 000261cc1c..825adafdaf 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
@@ -774,7 +774,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     return result.build();
   }
 
-  private static final Function<Target, SkyKey> TARGET_TO_SKY_KEY =
+  static final Function<Target, SkyKey> TARGET_TO_SKY_KEY =
       new Function<Target, SkyKey>() {
         @Override
         public SkyKey apply(Target target) {