When transforming labels into targets in SkyQueryEnvironment, stream the result to a callback instead of returning it directly. This means that the targets a precomputed pattern resolves to can be processed incrementally.

This is the sixth and hopefully final 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=111697983

1 files changed, 48 insertions, 22 deletions

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 f6181f29ea..7df803396a 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
@@ -14,7 +14,6 @@
 package com.google.devtools.build.lib.query2;
 
 import com.google.common.base.Function;
-import com.google.common.base.Functions;
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ArrayListMultimap;
@@ -185,10 +184,10 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
   private Map<Target, Collection<Target>> makeTargetsMap(Map<SkyKey, Iterable<SkyKey>> input) {
     ImmutableMap.Builder<Target, Collection<Target>> result = ImmutableMap.builder();
     
-    Map<SkyKey, Target> allTargets = makeTargetsWithAssociations(
-        Sets.newHashSet(Iterables.concat(input.values())));
+    Map<SkyKey, Target> allTargets =
+        makeTargetsFromSkyKeys(Sets.newHashSet(Iterables.concat(input.values())));
 
-    for (Map.Entry<SkyKey, Target> entry : makeTargetsWithAssociations(input.keySet()).entrySet()) {
+    for (Map.Entry<SkyKey, Target> entry : makeTargetsFromSkyKeys(input.keySet()).entrySet()) {
       Iterable<SkyKey> skyKeys = input.get(entry.getKey());
       Set<Target> targets = CompactHashSet.createWithExpectedSize(Iterables.size(skyKeys));
       for (SkyKey key : skyKeys) {
@@ -316,8 +315,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     aggregator.processLastPending();
   }
 
-  @Override
-  public Uniquifier<Target> createUniquifier() {
+  private static Uniquifier<Target> uniquifier() {
     return new AbstractUniquifier<Target, Label>() {
       @Override
       protected Label extractKey(Target target) {
@@ -326,6 +324,11 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     };
   }
 
+  @Override
+  public Uniquifier<Target> createUniquifier() {
+    return uniquifier();
+  }
+
   /**
    * 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.
@@ -366,7 +369,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
       Set<Label> labels = precomputedPatterns.get(pattern);
       if (labels != null) {
         try {
-          callback.process(ImmutableSet.copyOf(makeTargetsFromLabels(labels)));
+          makeTargetsFromLabels(labels, callback);
         } catch (InterruptedException e) {
           throw new QueryException(owner, e.getMessage());
         }
@@ -626,19 +629,39 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
     }
   };
 
-  private Map<SkyKey, Target> makeTargetsWithAssociations(Iterable<SkyKey> keys) {
-    return makeTargetsWithAssociations(keys, SKYKEY_TO_LABEL);
-  }
-
-  private Collection<Target> makeTargetsFromLabels(Iterable<Label> labels) {
-    return makeTargetsWithAssociations(labels, Functions.<Label>identity()).values();
+  private void makeTargetsFromLabels(Iterable<Label> labels, Callback<Target> callback)
+      throws QueryException, InterruptedException {
+    Multimap<SkyKey, Label> packageKeyToLabelMap = ArrayListMultimap.create();
+    for (Label label : labels) {
+      try {
+        packageKeyToLabelMap.put(getPackageKeyAndValidateLabel(label), label);
+      } catch (QueryException e) {
+        // Skip disallowed labels.
+      }
+    }
+    Uniquifier<Target> uniquifier = createUniquifier();
+    for (List<SkyKey> packageKeys :
+        Iterables.partition(packageKeyToLabelMap.keySet(), BATCH_CALLBACK_SIZE)) {
+      Map<SkyKey, SkyValue> packageMap = graph.getSuccessfulValues(packageKeys);
+      // Conservatively say all our targets are in different packages.
+      List<Target> targets = new ArrayList<>(packageMap.size());
+      for (Map.Entry<SkyKey, SkyValue> entry : packageMap.entrySet()) {
+        for (Label label : packageKeyToLabelMap.get(entry.getKey())) {
+          try {
+            targets.add(((PackageValue) entry.getValue()).getPackage().getTarget(label.getName()));
+          } catch (NoSuchTargetException e) {
+            // Skip missing target.
+          }
+        }
+      }
+      callback.process(uniquifier.unique(targets));
+    }
   }
 
-  private <E> Map<E, Target> makeTargetsWithAssociations(Iterable<E> keys,
-      Function<E, Label> toLabel) {
-    Multimap<SkyKey, E> packageKeyToTargetKeyMap = ArrayListMultimap.create();
-    for (E key : keys) {
-      Label label = toLabel.apply(key);
+  private Map<SkyKey, Target> makeTargetsFromSkyKeys(Iterable<SkyKey> keys) {
+    Multimap<SkyKey, SkyKey> packageKeyToTargetKeyMap = ArrayListMultimap.create();
+    for (SkyKey key : keys) {
+      Label label = SKYKEY_TO_LABEL.apply(key);
       if (label == null) {
         continue;
       }
@@ -648,13 +671,16 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
         // Skip disallowed labels.
       }
     }
-    ImmutableMap.Builder<E, Target> result = ImmutableMap.builder();
+    ImmutableMap.Builder<SkyKey, Target> result = ImmutableMap.builder();
     Map<SkyKey, SkyValue> packageMap = graph.getSuccessfulValues(packageKeyToTargetKeyMap.keySet());
     for (Map.Entry<SkyKey, SkyValue> entry : packageMap.entrySet()) {
-      for (E targetKey : packageKeyToTargetKeyMap.get(entry.getKey())) {
+      for (SkyKey targetKey : packageKeyToTargetKeyMap.get(entry.getKey())) {
         try {
-          result.put(targetKey, ((PackageValue) entry.getValue()).getPackage()
-              .getTarget((toLabel.apply(targetKey)).getName()));
+          result.put(
+              targetKey,
+              ((PackageValue) entry.getValue())
+                  .getPackage()
+                  .getTarget((SKYKEY_TO_LABEL.apply(targetKey)).getName()));
         } catch (NoSuchTargetException e) {
           // Skip missing target.
         }