Resolve target patterns on the fly in SkyQueryEnvironment. Cache only the label sets that are precomputed in the graph.

This is the fourth step in a series to allow processing large sets of targets in query target patterns via streaming batches rather than all at once. This may make SkyQueryEnvironment slower when evaluating queries with repeated target patterns, or many target patterns that would benefit from graph lookups that were batched across all patterns. But that is not currently a bottleneck we're concerned about.

--
MOS_MIGRATED_REVID=111626483

3 files changed, 81 insertions, 62 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
index c63e00dd37..004adec534 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
@@ -58,7 +58,6 @@ import javax.annotation.Nullable;
 public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironment<T> {
   protected final ErrorSensingEventHandler eventHandler;
   private final Map<String, Set<T>> letBindings = new HashMap<>();
-  protected final Map<String, Set<Target>> resolvedTargetPatterns = new HashMap<>();
   protected final boolean keepGoing;
   protected final boolean strictScope;
 
@@ -143,14 +142,13 @@ public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironme
 
     final AtomicBoolean empty = new AtomicBoolean(true);
     try (final AutoProfiler p = AutoProfiler.logged("evaluating query", LOG)) {
-      resolvedTargetPatterns.clear();
 
       // In the --nokeep_going case, errors are reported in the order in which the patterns are
       // specified; using a linked hash set here makes sure that the left-most error is reported.
       Set<String> targetPatternSet = new LinkedHashSet<>();
       expr.collectTargetPatterns(targetPatternSet);
       try {
-        resolvedTargetPatterns.putAll(preloadOrThrow(expr, targetPatternSet));
+        preloadOrThrow(expr, targetPatternSet);
       } catch (TargetParsingException e) {
         // Unfortunately, by evaluating the patterns in parallel, we lose some location information.
         throw new QueryException(expr, e.getMessage());
@@ -226,21 +224,23 @@ public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironme
 
   public Set<T> evalTargetPattern(QueryExpression caller, String pattern)
       throws QueryException {
-    if (!resolvedTargetPatterns.containsKey(pattern)) {
-      try {
-        resolvedTargetPatterns.putAll(preloadOrThrow(caller, ImmutableList.of(pattern)));
-      } catch (TargetParsingException e) {
-        // Will skip the target and keep going if -k is specified.
-        reportBuildFileError(caller, e.getMessage());
-      }
+    try {
+      preloadOrThrow(caller, ImmutableList.of(pattern));
+    } catch (TargetParsingException e) {
+      // Will skip the target and keep going if -k is specified.
+      reportBuildFileError(caller, e.getMessage());
     }
     AggregateAllCallback<T> aggregatingCallback = new AggregateAllCallback<>();
     getTargetsMatchingPattern(caller, pattern, aggregatingCallback);
     return aggregatingCallback.getResult();
   }
 
-  protected abstract Map<String, Set<Target>> preloadOrThrow(
-      QueryExpression caller, Collection<String> patterns)
+  /**
+   * Perform any work that should be done ahead of time to resolve the target patterns in the
+   * query. Implementations may choose to cache the results of resolving the patterns, cache
+   * intermediate work, or not cache and resolve patterns on the fly.
+   */
+  protected abstract void preloadOrThrow(QueryExpression caller, Collection<String> patterns)
       throws QueryException, TargetParsingException;
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
index 2f889aeb32..cde183c790 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/BlazeQueryEnvironment.java
@@ -50,6 +50,7 @@ import com.google.devtools.build.lib.vfs.PathFragment;
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
@@ -64,6 +65,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
 public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
 
   private static final int MAX_DEPTH_FULL_SCAN_LIMIT = 20;
+  private final Map<String, Set<Target>> resolvedTargetPatterns = new HashMap<>();
   private final TargetPatternEvaluator targetPatternEvaluator;
   private final TransitivePackageLoader transitivePackageLoader;
   private final TargetProvider targetProvider;
@@ -112,6 +114,7 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     // errors here.
     eventHandler.resetErrors();
     final AtomicBoolean empty = new AtomicBoolean(true);
+    resolvedTargetPatterns.clear();
     QueryEvalResult queryEvalResult = super.evaluateQuery(expr, new Callback<Target>() {
       @Override
       public void process(Iterable<Target> partialResult)
@@ -414,17 +417,20 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
       };
 
   @Override
-  protected Map<String, Set<Target>> preloadOrThrow(
-      QueryExpression caller, Collection<String> patterns) throws TargetParsingException {
-    try {
-      // Note that this may throw a RuntimeException if deps are missing in Skyframe and this is
-      // being called from within a SkyFunction.
-      return Maps.transformValues(
-          targetPatternEvaluator.preloadTargetPatterns(eventHandler, patterns, keepGoing),
-          RESOLVED_TARGETS_TO_TARGETS);
-    } catch (InterruptedException e) {
-      // TODO(bazel-team): Propagate the InterruptedException from here [skyframe-loading].
-      throw new TargetParsingException("interrupted");
+  protected void preloadOrThrow(QueryExpression caller, Collection<String> patterns)
+      throws TargetParsingException {
+    if (!resolvedTargetPatterns.keySet().containsAll(patterns)) {
+      try {
+        // Note that this may throw a RuntimeException if deps are missing in Skyframe and this is
+        // being called from within a SkyFunction.
+        resolvedTargetPatterns.putAll(
+            Maps.transformValues(
+                targetPatternEvaluator.preloadTargetPatterns(eventHandler, patterns, keepGoing),
+                RESOLVED_TARGETS_TO_TARGETS));
+      } catch (InterruptedException e) {
+        // TODO(bazel-team): Propagate the InterruptedException from here [skyframe-loading].
+        throw new TargetParsingException("interrupted");
+      }
     }
   }
 
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 749a563529..b5bf4f83fa 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
@@ -99,6 +99,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
 
   private ImmutableList<TargetPatternKey> universeTargetPatternKeys;
 
+  private final Map<String, Set<Label>> precomputedPatterns = new HashMap<>();
   private final BlazeTargetAccessor accessor = new BlazeTargetAccessor(this);
   private final int loadingPhaseThreads;
   private final WalkableGraphFactory graphFactory;
@@ -326,7 +327,54 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
   @Override
   public void getTargetsMatchingPattern(
       QueryExpression owner, String pattern, Callback<Target> callback) throws QueryException {
-    Set<Target> targets = new LinkedHashSet<>(resolvedTargetPatterns.get(pattern));
+    Set<Target> targets = ImmutableSet.of();
+    if (precomputedPatterns.containsKey(pattern)) {
+      Set<Label> labels = precomputedPatterns.get(pattern);
+      if (labels != null) {
+        targets = ImmutableSet.copyOf(makeTargetsFromLabels(labels));
+      } else {
+        TargetParsingException exception;
+        try {
+          // Because the graph was always initialized via a keep_going build, we know that the
+          // exception stored here must be a TargetParsingException. Thus the comment in
+          // SkyframeTargetPatternEvaluator#parseTargetPatternKeys describing the situation in which
+          // the exception acceptance must be looser does not apply here.
+          exception =
+              (TargetParsingException)
+                  Preconditions.checkNotNull(
+                      graph.getException(
+                          TargetPatternValue.key(
+                              pattern,
+                              TargetPatternEvaluator.DEFAULT_FILTERING_POLICY,
+                              parserPrefix)),
+                      pattern);
+        } catch (TargetParsingException e) {
+          exception = e;
+        }
+        reportBuildFileError(owner, exception.getMessage());
+      }
+    } else {
+      // If the graph doesn't contain a value for this target pattern, try to directly evaluate
+      // it, by making use of packages already present in the graph.
+      try {
+        TargetPatternKey targetPatternKey =
+            ((TargetPatternKey)
+                TargetPatternValue.key(
+                        pattern, TargetPatternEvaluator.DEFAULT_FILTERING_POLICY, parserPrefix)
+                    .argument());
+        GraphBackedRecursivePackageProvider provider =
+            new GraphBackedRecursivePackageProvider(graph, universeTargetPatternKeys, pkgPath);
+        RecursivePackageProviderBackedTargetPatternResolver resolver =
+            new RecursivePackageProviderBackedTargetPatternResolver(
+                provider, eventHandler, targetPatternKey.getPolicy());
+        TargetPattern parsedPattern = targetPatternKey.getParsedPattern();
+        targets = parsedPattern.eval(resolver).getTargets();
+      } 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();
@@ -337,7 +385,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
       }
     }
     try {
-      callback.process(targets);
+      callback.process(filterTargetsNotInGraph(targets));
     } catch (InterruptedException e) {
       throw new QueryException(owner, e.getMessage());
     }
@@ -492,13 +540,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
   }
 
   @Override
-  protected Map<String, Set<Target>> preloadOrThrow(
-      QueryExpression caller, Collection<String> patterns)
+  protected void preloadOrThrow(QueryExpression caller, Collection<String> patterns)
       throws QueryException, TargetParsingException {
-    GraphBackedRecursivePackageProvider provider =
-        new GraphBackedRecursivePackageProvider(graph, universeTargetPatternKeys, pkgPath);
-    Map<String, Set<Target>> result = Maps.newHashMapWithExpectedSize(patterns.size());
-
     Map<String, SkyKey> keys = new HashMap<>(patterns.size());
 
     for (String pattern : patterns) {
@@ -509,13 +552,11 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
                 pattern, TargetPatternEvaluator.DEFAULT_FILTERING_POLICY, parserPrefix));
       } catch (TargetParsingException e) {
         reportBuildFileError(caller, e.getMessage());
-        result.put(pattern, ImmutableSet.<Target>of());
       }
     }
     // Get all the patterns in one batch call
     Map<SkyKey, SkyValue> existingPatterns = graph.getSuccessfulValues(keys.values());
 
-    Map<String, Set<Target>> patternsWithTargetsToFilter = new HashMap<>();
     for (String pattern : patterns) {
       SkyKey patternKey = keys.get(pattern);
       if (patternKey == null) {
@@ -527,8 +568,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
         // The graph already contains a value or exception for this target pattern, so we use it.
         TargetPatternValue value = (TargetPatternValue) existingPatterns.get(patternKey);
         if (value != null) {
-          result.put(
-              pattern, ImmutableSet.copyOf(makeTargetsFromLabels(value.getTargets().getTargets())));
+          precomputedPatterns.put(pattern, value.getTargets().getTargets());
         } else {
           // Because the graph was always initialized via a keep_going build, we know that the
           // exception stored here must be a TargetParsingException. Thus the comment in
@@ -538,39 +578,12 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
               (TargetParsingException)
                   Preconditions.checkNotNull(graph.getException(patternKey), pattern);
         }
-      } else {
-        // If the graph doesn't contain a value for this target pattern, try to directly evaluate
-        // it, by making use of packages already present in the graph.
-        TargetPatternValue.TargetPatternKey targetPatternKey =
-            ((TargetPatternValue.TargetPatternKey) patternKey.argument());
-
-        RecursivePackageProviderBackedTargetPatternResolver resolver =
-            new RecursivePackageProviderBackedTargetPatternResolver(provider, eventHandler,
-                targetPatternKey.getPolicy());
-        TargetPattern parsedPattern = targetPatternKey.getParsedPattern();
-        try {
-          patternsWithTargetsToFilter.put(pattern, parsedPattern.eval(resolver).getTargets());
-        } catch (TargetParsingException e) {
-          targetParsingException = e;
-        } catch (InterruptedException e) {
-          throw new QueryException(e.getMessage());
-        }
       }
 
       if (targetParsingException != null) {
         reportBuildFileError(caller, targetParsingException.getMessage());
-        result.put(pattern, ImmutableSet.<Target>of());
       }
     }
-    // filterTargetsNotInGraph does graph lookups. So we batch all the queries in one call.
-    Set<Target> targetsInGraph = filterTargetsNotInGraph(
-        ImmutableSet.copyOf(Iterables.concat(patternsWithTargetsToFilter.values())));
-
-    for (Entry<String, Set<Target>> pattern : patternsWithTargetsToFilter.entrySet()) {
-      result.put(pattern.getKey(), Sets.intersection(pattern.getValue(), targetsInGraph));
-    }
-
-    return result;
   }
 
   private static final Function<SkyKey, Label> SKYKEY_TO_LABEL = new Function<SkyKey, Label>() {