(1) Allow QueryExpressionMapper#map to throw a QueryException.

(2) Refactor the simple query concurrency support by introducing QueryExpression#canEvalConcurrently and removing the primitive BinaryOperatorExpression#evalConcurrently and the hardcoded SkyQueryEnvironment logic around it. This way each QueryExpression can safely manage its own concurrency. A followup CL will ensure that concurrent evaluation occurs for as much of the query expression as possible, rather than just for the top-level query expression node.
(3) Make a few query internals public.

--
MOS_MIGRATED_REVID=126324637

11 files changed, 97 insertions, 130 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 0292f08458..c2d6e24816 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
@@ -148,7 +148,8 @@ public abstract class AbstractBlazeQueryEnvironment<T>
   @Override
   public abstract void close();
 
-  public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
+  public QueryExpression transformParsedQuery(QueryExpression queryExpression)
+      throws QueryException {
     return queryExpression;
   }
 
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 1865381c9c..ce5aa14181 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
@@ -55,7 +55,6 @@ import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
 import com.google.devtools.build.lib.pkgcache.TargetPatternEvaluator;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.query2.engine.AllRdepsFunction;
-import com.google.devtools.build.lib.query2.engine.BinaryOperatorExpression;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.FunctionExpression;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
@@ -135,7 +134,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   private final int loadingPhaseThreads;
   private final WalkableGraphFactory graphFactory;
   private final List<String> universeScope;
-  private final String parserPrefix;
+  protected final String parserPrefix;
   private final PathPackageLocator pkgPath;
 
   // Note that the executor returned by Executors.newFixedThreadPool doesn't start any threads
@@ -227,42 +226,50 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     }
   }
 
+  /**
+   * A {@link QueryExpressionMapper} that transforms each occurrence of an expression of the form
+   * {@literal 'rdeps(<universeScope>, <T>)'} to {@literal 'allrdeps(<T>)'}. The latter is more
+   * efficient.
+   */
+  protected static class RdepsToAllRdepsQueryExpressionMapper extends QueryExpressionMapper {
+    protected final TargetPattern.Parser targetPatternParser;
+    private final String absoluteUniverseScopePattern;
+
+    protected RdepsToAllRdepsQueryExpressionMapper(
+        TargetPattern.Parser targetPatternParser,
+        String universeScopePattern) {
+      this.targetPatternParser = targetPatternParser;
+      this.absoluteUniverseScopePattern = targetPatternParser.absolutize(universeScopePattern);
+    }
+
+    @Override
+    public QueryExpression map(FunctionExpression functionExpression) throws QueryException {
+      if (functionExpression.getFunction().getName().equals(new RdepsFunction().getName())) {
+        List<Argument> args = functionExpression.getArgs();
+        QueryExpression universeExpression = args.get(0).getExpression();
+        if (universeExpression instanceof TargetLiteral) {
+          TargetLiteral literalUniverseExpression = (TargetLiteral) universeExpression;
+          String absolutizedUniverseExpression =
+              targetPatternParser.absolutize(literalUniverseExpression.getPattern());
+          if (absolutizedUniverseExpression.equals(absoluteUniverseScopePattern)) {
+            List<Argument> argsTail = args.subList(1, functionExpression.getArgs().size());
+            return new FunctionExpression(new AllRdepsFunction(), argsTail);
+          }
+        }
+      }
+      return super.map(functionExpression);
+    }
+  }
+
   @Override
   public void close() {
     ExecutorUtil.interruptibleShutdown(threadPool);
   }
 
   @Override
-  public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
-    // Transform each occurrence of an expressions of the form 'rdeps(<universeScope>, <T>)' to
-    // 'allrdeps(<T>)'. The latter is more efficient.
-    if (universeScope.size() != 1) {
-      return queryExpression;
-    }
-    final TargetPattern.Parser targetPatternParser = new TargetPattern.Parser(parserPrefix);
-    String universeScopePattern = Iterables.getOnlyElement(universeScope);
-    final String absoluteUniverseScopePattern =
-        targetPatternParser.absolutize(universeScopePattern);
-    QueryExpressionMapper rdepsToAllRDepsMapper = new QueryExpressionMapper() {
-      @Override
-      public QueryExpression map(FunctionExpression functionExpression) {
-        if (functionExpression.getFunction().getName().equals(new RdepsFunction().getName())) {
-          List<Argument> args = functionExpression.getArgs();
-          QueryExpression universeExpression = args.get(0).getExpression();
-          if (universeExpression instanceof TargetLiteral) {
-            TargetLiteral literalUniverseExpression = (TargetLiteral) universeExpression;
-            String absolutizedUniverseExpression =
-                targetPatternParser.absolutize(literalUniverseExpression.getPattern());
-            if (absolutizedUniverseExpression.equals(absoluteUniverseScopePattern)) {
-              List<Argument> argsTail = args.subList(1, functionExpression.getArgs().size());
-              return new FunctionExpression(new AllRdepsFunction(), argsTail);
-            }
-          }
-        }
-        return super.map(functionExpression);
-      }
-    };
-    QueryExpression transformedQueryExpression = queryExpression.getMapped(rdepsToAllRDepsMapper);
+  public final QueryExpression transformParsedQuery(QueryExpression queryExpression)
+      throws QueryException {
+    QueryExpression transformedQueryExpression = getTransformedQueryExpression(queryExpression);
     LOG.info(String.format(
         "transformed query [%s] to [%s]",
         Ascii.truncate(
@@ -272,6 +279,17 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     return transformedQueryExpression;
   }
 
+  protected QueryExpression getTransformedQueryExpression(QueryExpression queryExpression)
+      throws QueryException {
+    if (universeScope.size() != 1) {
+      return queryExpression;
+    }
+    TargetPattern.Parser targetPatternParser = new TargetPattern.Parser(parserPrefix);
+    String universeScopePattern = Iterables.getOnlyElement(universeScope);
+    return queryExpression.getMapped(
+        new RdepsToAllRdepsQueryExpressionMapper(targetPatternParser, universeScopePattern));
+  }
+
   @Override
   public QueryEvalResult evaluateQuery(QueryExpression expr, Callback<Target> callback)
       throws QueryException, InterruptedException {
@@ -303,7 +321,7 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
         };
     try (final AutoProfiler p = AutoProfiler.logged("evaluating query", LOG)) {
       try {
-        if (canEvalConcurrently(expr)) {
+        if (expr.canEvalConcurrently()) {
           expr.evalConcurrently(this, callbackWithEmptyCheck, threadPool);
         } else {
           expr.eval(this, callbackWithEmptyCheck);
@@ -329,26 +347,6 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     return new QueryEvalResult(!eventHandler.hasErrors(), empty.get());
   }
 
-  // TODO(mschaller): This method and its use above are a quick temporary fix to a threadsafety
-  // problem that can happen when the operands of a BinaryOperatorExpression contain LetExpressions.
-  // Namely, concurrent reads and writes to AbstractBlazeQueryEnvironment#letBindings may fail or
-  // produce the wrong results.
-  // For now, this limits concurrent query expression evaluation to BinaryOperatorExpressions with
-  // TargetLiteral operands.
-  private static boolean canEvalConcurrently(QueryExpression expr) {
-    if (!(expr instanceof BinaryOperatorExpression)) {
-      return false;
-    }
-    BinaryOperatorExpression binaryExpr = (BinaryOperatorExpression) expr;
-    ImmutableList<QueryExpression> operands = binaryExpr.getOperands();
-    for (QueryExpression operand : operands) {
-      if (!(operand instanceof TargetLiteral)) {
-        return false;
-      }
-    }
-    return true;
-  }
-
   private Map<Target, Collection<Target>> makeTargetsMap(Map<SkyKey, Iterable<SkyKey>> input) {
     ImmutableMap.Builder<Target, Collection<Target>> result = ImmutableMap.builder();
 
@@ -612,6 +610,18 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     }
   }
 
+  public Map<PackageIdentifier, Package> bulkGetPackages(Iterable<PackageIdentifier> pkgIds) {
+    Set<SkyKey> pkgKeys = ImmutableSet.copyOf(PackageValue.keys(pkgIds));
+    ImmutableMap.Builder<PackageIdentifier, Package> pkgResults = ImmutableMap.builder();
+    Map<SkyKey, SkyValue> packages = graph.getSuccessfulValues(pkgKeys);
+    for (Map.Entry<SkyKey, SkyValue> pkgEntry : packages.entrySet()) {
+      PackageIdentifier pkgId = (PackageIdentifier) pkgEntry.getKey().argument();
+      PackageValue pkgValue = (PackageValue) pkgEntry.getValue();
+      pkgResults.put(pkgId, Preconditions.checkNotNull(pkgValue.getPackage(), pkgId));
+    }
+    return pkgResults.build();
+  }
+
   @Override
   public void buildTransitiveClosure(QueryExpression caller, Set<Target> targets, int maxDepth)
       throws QueryException {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java
index c1d74c9873..241aa522c7 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/BinaryOperatorExpression.java
@@ -14,19 +14,12 @@
 package com.google.devtools.build.lib.query2.engine;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.util.concurrent.Futures;
-import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.ListeningExecutorService;
-import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.lib.query2.engine.Lexer.TokenKind;
 import com.google.devtools.build.lib.util.Preconditions;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * A binary algebraic set operation.
@@ -45,14 +38,13 @@ public class BinaryOperatorExpression extends QueryExpression {
   private final Lexer.TokenKind operator; // ::= INTERSECT/CARET | UNION/PLUS | EXCEPT/MINUS
   private final ImmutableList<QueryExpression> operands;
 
-  BinaryOperatorExpression(Lexer.TokenKind operator,
-                           List<QueryExpression> operands) {
+  public BinaryOperatorExpression(Lexer.TokenKind operator, List<QueryExpression> operands) {
     Preconditions.checkState(operands.size() > 1);
     this.operator = operator;
     this.operands = ImmutableList.copyOf(operands);
   }
 
-  Lexer.TokenKind getOperator() {
+  public Lexer.TokenKind getOperator() {
     return operator;
   }
 
@@ -63,55 +55,10 @@ public class BinaryOperatorExpression extends QueryExpression {
   @Override
   public <T> void eval(QueryEnvironment<T> env, Callback<T> callback)
       throws QueryException, InterruptedException {
-    evalConcurrently(env, callback, MoreExecutors.newDirectExecutorService());
-  }
 
-  @Override
-  public <T> void evalConcurrently(
-      final QueryEnvironment<T> env,
-      final Callback<T> callback,
-      ListeningExecutorService executorService)
-      throws QueryException, InterruptedException {
     if (operator == TokenKind.PLUS || operator == TokenKind.UNION) {
-      final AtomicReference<InterruptedException> interruptRef = new AtomicReference<>();
-      final AtomicReference<QueryException> queryExceptionRef = new AtomicReference<>();
-      ArrayList<ListenableFuture<?>> futures = new ArrayList<>(operands.size());
-      for (final QueryExpression operand : operands) {
-        // When executorService has an implementation that evaluates runnables in a non-serial
-        // order, like a fixedSizeThreadPool, the following code does not guarantee that operands'
-        // targets are emitted via the callback in the operands' order. And that's OK!
-        // BinaryOperatorExpression is a set operation. The query documentation states
-        // that set operations don't introduce any ordering constraints of their own.
-        //
-        // Ordering constraints for other kinds of expressions are enforced by the query
-        // environment.
-        futures.add(
-            executorService.submit(
-                new Runnable() {
-                  @Override
-                  public void run() {
-                    try {
-                      env.eval(operand, callback);
-                    } catch (QueryException e) {
-                      queryExceptionRef.compareAndSet(null, e);
-                    } catch (InterruptedException e) {
-                      interruptRef.compareAndSet(null, e);
-                    }
-                  }
-                }));
-      }
-      try {
-        Futures.allAsList(futures).get();
-      } catch (ExecutionException e) {
-        throw new IllegalStateException(e);
-      }
-      InterruptedException interruptedExceptionIfAny = interruptRef.get();
-      if (interruptedExceptionIfAny != null) {
-        throw interruptedExceptionIfAny;
-      }
-      QueryException queryException = queryExceptionRef.get();
-      if (queryException != null) {
-        throw queryException;
+      for (QueryExpression operand : operands) {
+        env.eval(operand, callback);
       }
       return;
     }
@@ -145,7 +92,7 @@ public class BinaryOperatorExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+  public QueryExpression getMapped(QueryExpressionMapper mapper) throws QueryException {
     return mapper.map(this);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
index f6fcc27ef5..790abf9372 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
@@ -60,7 +60,7 @@ public class FunctionExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+  public QueryExpression getMapped(QueryExpressionMapper mapper) throws QueryException {
     return mapper.map(this);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java
index a1cc7a7cbf..b5805d4544 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/LetExpression.java
@@ -89,7 +89,7 @@ class LetExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+  public QueryExpression getMapped(QueryExpressionMapper mapper) throws QueryException {
     return mapper.map(this);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/Lexer.java b/src/main/java/com/google/devtools/build/lib/query2/engine/Lexer.java
index d259e9543f..ef5343b9f1 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/Lexer.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/Lexer.java
@@ -29,7 +29,7 @@ import java.util.Set;
  * No string escapes are allowed ("\").  Given the domain, that's not currently
  * a problem.
  */
-final class Lexer {
+public final class Lexer {
 
   /**
    * Discriminator for different kinds of tokens.
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java
index e93b397299..73688ec84d 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpression.java
@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.query2.engine;
 
 import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.devtools.build.lib.util.Preconditions;
 
 import java.util.Collection;
 
@@ -73,9 +74,8 @@ public abstract class QueryExpression {
       throws QueryException, InterruptedException;
 
   /**
-   * Evaluates this query in the specified environment, as in {@link
-   * #eval(QueryEnvironment, Callback)}. If the query expression supports concurrent evaluation, it
-   * may employ {@code executorService}.
+   * If {@code canEvalConcurrently()}, evaluates this query in the specified environment, as in
+   * {@link #eval(QueryEnvironment, Callback)}, employing {@code executorService}.
    *
    * <p>The caller must ensure that both {@code env} and {@code callback} are effectively
    * threadsafe. The query expression may call their methods from multiple threads.
@@ -83,7 +83,16 @@ public abstract class QueryExpression {
   public <T> void evalConcurrently(
       QueryEnvironment<T> env, Callback<T> callback, ListeningExecutorService executorService)
       throws QueryException, InterruptedException {
-    this.eval(env, callback);
+    Preconditions.checkState(canEvalConcurrently());
+    eval(env, callback);
+  }
+
+  /**
+   * Whether the query expression can be evaluated concurrently. If so, {@link #evalConcurrently}
+   * should be preferred over {@link #eval}.
+   */
+  public boolean canEvalConcurrently() {
+    return false;
   }
 
   /**
@@ -93,7 +102,7 @@ public abstract class QueryExpression {
   public abstract void collectTargetPatterns(Collection<String> literals);
 
   /* Implementations should just be {@code return mapper.map(this)}. */
-  public abstract QueryExpression getMapped(QueryExpressionMapper mapper);
+  public abstract QueryExpression getMapped(QueryExpressionMapper mapper) throws QueryException;
 
   /**
    * Returns this query expression pretty-printed.
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java
index f04ed1ec4b..896e80596f 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryExpressionMapper.java
@@ -26,11 +26,12 @@ import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Argument;
  * override these methods in order to implement an arbitrary transformation.
  */
 public class QueryExpressionMapper {
-  public QueryExpression map(TargetLiteral targetLiteral) {
+  public QueryExpression map(TargetLiteral targetLiteral) throws QueryException {
     return targetLiteral;
   }
 
-  public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression) {
+  public QueryExpression map(BinaryOperatorExpression binaryOperatorExpression)
+      throws QueryException {
     boolean changed = false;
     ImmutableList.Builder<QueryExpression> mappedOperandsBuilder = ImmutableList.builder();
     for (QueryExpression operand : binaryOperatorExpression.getOperands()) {
@@ -46,7 +47,7 @@ public class QueryExpressionMapper {
         : binaryOperatorExpression;
   }
 
-  public QueryExpression map(FunctionExpression functionExpression) {
+  public QueryExpression map(FunctionExpression functionExpression) throws QueryException {
     boolean changed = false;
     ImmutableList.Builder<Argument> mappedArgumentBuilder = ImmutableList.builder();
     for (Argument argument : functionExpression.getArgs()) {
@@ -70,7 +71,7 @@ public class QueryExpressionMapper {
         : functionExpression;
   }
 
-  public QueryExpression map(LetExpression letExpression) {
+  public QueryExpression map(LetExpression letExpression) throws QueryException {
     boolean changed = false;
     QueryExpression mappedVarExpr = letExpression.getVarExpr().getMapped(this);
     if (mappedVarExpr != letExpression.getVarExpr()) {
@@ -85,7 +86,7 @@ public class QueryExpressionMapper {
         : letExpression;
   }
 
-  public QueryExpression map(SetExpression setExpression) {
+  public QueryExpression map(SetExpression setExpression) throws QueryException {
     return setExpression;
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java
index 1369d8ace6..8297829c63 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/SetExpression.java
@@ -61,7 +61,7 @@ class SetExpression extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+  public QueryExpression getMapped(QueryExpressionMapper mapper) throws QueryException {
     return mapper.map(this);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java b/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java
index e7d8c00d73..b3f9b812fd 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/TargetLiteral.java
@@ -32,7 +32,7 @@ public final class TargetLiteral extends QueryExpression {
 
   private final String pattern;
 
-  TargetLiteral(String pattern) {
+  public TargetLiteral(String pattern) {
     this.pattern = Preconditions.checkNotNull(pattern);
   }
 
@@ -67,7 +67,7 @@ public final class TargetLiteral extends QueryExpression {
   }
 
   @Override
-  public QueryExpression getMapped(QueryExpressionMapper mapper) {
+  public QueryExpression getMapped(QueryExpressionMapper mapper) throws QueryException {
     return mapper.map(this);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
index 3dd0f4d7c8..1ec51f01b2 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
@@ -142,17 +142,16 @@ public final class QueryCommand implements BlazeCommand {
             queryOptions.universeScope,
             queryOptions.loadingPhaseThreads,
             settings)) {
-
       // 1. Parse and transform query:
       QueryExpression expr;
       try {
         expr = QueryExpression.parse(query, queryEnv);
+        expr = queryEnv.transformParsedQuery(expr);
       } catch (QueryException e) {
         env.getReporter()
             .handle(Event.error(null, "Error while parsing '" + query + "': " + e.getMessage()));
         return ExitCode.COMMAND_LINE_ERROR;
       }
-      expr = queryEnv.transformParsedQuery(expr);
 
       PrintStream output = null;
       OutputFormatterCallback<Target> callback;