Make some clean-ups and assertions in query environments that may be …

…useful for attached bug.

PiperOrigin-RevId: 423178864

src/main/java/com/google/devtools/build/lib/query2/PostAnalysisQueryEnvironment.java

@@ -149,7 +149,7 @@ public QueryEvalResult evaluateQuery(
  QueryExpression expr, ThreadSafeOutputFormatterCallback<T> callback)
  throws QueryException, InterruptedException, IOException {
  beforeEvaluateQuery();
- return super.evaluateQuery(expr, callback);
+ return evaluateQueryInternal(expr, callback);
  }
 
  private void beforeEvaluateQuery() throws QueryException {

src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java

@@ -514,7 +514,7 @@ public QueryEvalResult evaluateQuery(
  BatchStreamedCallback batchCallback =
  new BatchStreamedCallback(
  callback, BATCH_CALLBACK_SIZE, createUniquifierForOuterBatchStreamedCallback(expr));
- return super.evaluateQuery(expr, batchCallback);
+ return evaluateQueryInternal(expr, batchCallback);
  }
 
  Map<SkyKey, Collection<Target>> targetifyValues(Map<SkyKey, ? extends Iterable<SkyKey>> input)
@@ -1080,8 +1080,7 @@ public void buildTransitiveClosure(
  }
 
  @Override
- protected void preloadOrThrow(QueryExpression caller, Collection<String> patterns)
- throws QueryException, TargetParsingException {
+ protected final void preloadOrThrow(QueryExpression caller, Collection<String> patterns) {
  // SkyQueryEnvironment directly evaluates target patterns in #getTarget and similar methods
  // using its graph, which is prepopulated using the universeScope (see #beforeEvaluateQuery),
  // so no preloading of target patterns is necessary.

src/main/java/com/google/devtools/build/lib/query2/common/AbstractBlazeQueryEnvironment.java

@@ -15,6 +15,7 @@
 
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.concurrent.TimeUnit.SECONDS;
 
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
@@ -73,6 +74,8 @@
  */
 public abstract class AbstractBlazeQueryEnvironment<T>
  implements QueryEnvironment<T>, AutoCloseable {
+ private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
+
  protected ErrorSensingEventHandler<DetailedExitCode> eventHandler;
  protected final boolean keepGoing;
  protected final boolean strictScope;
@@ -83,8 +86,6 @@ public abstract class AbstractBlazeQueryEnvironment<T>
  protected final Set<Setting> settings;
  protected final List<QueryFunction> extraFunctions;
 
- private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
-
  protected AbstractBlazeQueryEnvironment(
  boolean keepGoing,
  boolean strictScope,
@@ -131,6 +132,10 @@ protected QueryExpressionContext<T> createEmptyContext() {
  return QueryExpressionContext.empty();
  }
 
+ public abstract QueryEvalResult evaluateQuery(
+ QueryExpression expr, ThreadSafeOutputFormatterCallback<T> callback)
+ throws QueryException, IOException, InterruptedException;
+
  /**
  * Evaluate the specified query expression in this environment, streaming results to the given
  * {@code callback}. {@code callback.start()} will be called before query evaluation and {@code
@@ -143,7 +148,7 @@ protected QueryExpressionContext<T> createEmptyContext() {
  * @throws QueryException if the evaluation failed and {@code --nokeep_going} was in effect
  * @throws IOException for output formatter failures from {@code callback}
  */
- public QueryEvalResult evaluateQuery(
+ protected final QueryEvalResult evaluateQueryInternal(
  QueryExpression expr, ThreadSafeOutputFormatterCallback<T> callback)
  throws QueryException, InterruptedException, IOException {
  EmptinessSensingCallback<T> emptySensingCallback = new EmptinessSensingCallback<>(callback);
@@ -330,7 +335,7 @@ public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
  }
 
  @Override
- public void handleError(
+ public final void handleError(
  QueryExpression expression, String message, DetailedExitCode detailedExitCode)
  throws QueryException {
  if (!keepGoing) {
@@ -358,6 +363,7 @@ public Map<Label, Target> getTargets(Iterable<Label> labels)
  try {
  target = getTarget(label);
  } catch (TargetNotFoundException e) {
+ logger.atInfo().withCause(e).atMostEvery(1, SECONDS).log("Failure to load %s", label);
  continue;
  }
  resultBuilder.put(label, target);

src/main/java/com/google/devtools/build/lib/query2/query/BlazeQueryEnvironment.java

@@ -68,6 +68,9 @@
 
 /**
  * The environment of a Blaze query. Not thread-safe.
+ *
+ * <p>This environment is valid only for a single query, called via {@link #evaluateQuery}. Call
+ * only once!
  */
 public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target> {
  private static final int MAX_DEPTH_FULL_SCAN_LIMIT = 20;
@@ -84,6 +87,8 @@ public class BlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
 
  private final BlazeTargetAccessor accessor = new BlazeTargetAccessor(this);
 
+ private boolean doneQuery;
+
  /**
  * Note that the correct operation of this class critically depends on the Reporter being a
  * singleton object, shared by all cooperating classes contributing to Query.
@@ -130,8 +135,9 @@ public DigraphQueryEvalResult<Target> evaluateQuery(
  QueryExpression expr,
  ThreadSafeOutputFormatterCallback<Target> callback)
  throws QueryException, InterruptedException, IOException {
- resolvedTargetPatterns.clear();
- QueryEvalResult queryEvalResult = super.evaluateQuery(expr, callback);
+ Preconditions.checkState(!doneQuery, "Can only use environment for one query: %s", expr);
+ doneQuery = true;
+ QueryEvalResult queryEvalResult = evaluateQueryInternal(expr, callback);
  return new DigraphQueryEvalResult<>(
  queryEvalResult.getSuccess(),
  queryEvalResult.isEmpty(),
@@ -458,13 +464,16 @@ public ThreadSafeMutableSet<Target> getBuildFiles(
  @Override
  protected void preloadOrThrow(QueryExpression caller, Collection<String> patterns)
  throws TargetParsingException, InterruptedException {
- if (!resolvedTargetPatterns.keySet().containsAll(patterns)) {
- // Note that this may throw a RuntimeException if deps are missing in Skyframe and this is
- // being called from within a SkyFunction.
- resolvedTargetPatterns.putAll(
- targetPatternPreloader.preloadTargetPatterns(
- eventHandler, relativeWorkingDirectory, patterns, keepGoing));
- }
+ Preconditions.checkState(
+ resolvedTargetPatterns.isEmpty(),
+ "Already resolved patterns: %s %s",
+ patterns,
+ resolvedTargetPatterns);
+ // Note that this may throw a RuntimeException if deps are missing in Skyframe and this is
+ // being called from within a SkyFunction.
+ resolvedTargetPatterns.putAll(
+ targetPatternPreloader.preloadTargetPatterns(
+ eventHandler, relativeWorkingDirectory, patterns, keepGoing));
  }
 
  private static void addIfUniqueLabel(Node<Target> node, Set<Label> labels, Set<Target> nodes) {

src/main/java/com/google/devtools/build/lib/query2/query/GraphlessBlazeQueryEnvironment.java

@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.query2.query;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
@@ -36,6 +37,7 @@
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.MinDepthUniquifier;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.CustomFunctionQueryEnvironment;
+import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionContext;
@@ -44,8 +46,10 @@
 import com.google.devtools.build.lib.query2.engine.QueryUtil.ThreadSafeMutableKeyExtractorBackedSetImpl;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.UniquifierImpl;
 import com.google.devtools.build.lib.query2.engine.SkyframeRestartQueryException;
+import com.google.devtools.build.lib.query2.engine.ThreadSafeOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
@@ -67,6 +71,9 @@
  * com.google.devtools.build.lib.query2.engine.DigraphQueryEvalResult}, and can therefore not be
  * used with most existing {@link com.google.devtools.build.lib.query2.query.output.OutputFormatter}
  * implementations, many of which expect the latter.
+ *
+ * <p>This environment is valid only for a single query, called via {@link #evaluateQuery}. Call
+ * only once!
  */
 public class GraphlessBlazeQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
  implements CustomFunctionQueryEnvironment<Target> {
@@ -83,6 +90,8 @@ public class GraphlessBlazeQueryEnvironment extends AbstractBlazeQueryEnvironmen
 
  private final BlazeTargetAccessor accessor = new BlazeTargetAccessor(this);
 
+ private boolean doneQuery = false;
+
  /**
  * Note that the correct operation of this class critically depends on the Reporter being a
  * singleton object, shared by all cooperating classes contributing to Query.
@@ -144,6 +153,15 @@ public QueryTaskFuture<Void> eval(
  return expr.eval(this, context, callback);
  }
 
+ @Override
+ public QueryEvalResult evaluateQuery(
+ QueryExpression expr, ThreadSafeOutputFormatterCallback<Target> callback)
+ throws QueryException, IOException, InterruptedException {
+ Preconditions.checkState(!doneQuery, "Can only use environment for one query: %s", expr);
+ doneQuery = true;
+ return evaluateQueryInternal(expr, callback);
+ }
+
  @Override
  public void close() {
  // BlazeQueryEnvironment has no resources that need to be cleaned up.
@@ -457,13 +475,16 @@ public ThreadSafeMutableSet<Target> getBuildFiles(
  @Override
  protected void preloadOrThrow(QueryExpression caller, Collection<String> patterns)
  throws TargetParsingException, InterruptedException {
- if (!resolvedTargetPatterns.keySet().containsAll(patterns)) {
- // Note that this may throw a RuntimeException if deps are missing in Skyframe and this is
- // being called from within a SkyFunction.
- resolvedTargetPatterns.putAll(
- targetPatternPreloader.preloadTargetPatterns(
- eventHandler, relativeWorkingDirectory, patterns, keepGoing));
- }
+ Preconditions.checkState(
+ resolvedTargetPatterns.isEmpty(),
+ "Already resolved patterns: %s %s",
+ patterns,
+ resolvedTargetPatterns);
+ // Note that this may throw a RuntimeException if deps are missing in Skyframe and this is
+ // being called from within a SkyFunction.
+ resolvedTargetPatterns.putAll(
+ targetPatternPreloader.preloadTargetPatterns(
+ eventHandler, relativeWorkingDirectory, patterns, keepGoing));
  }
 
  @Override