Move AbstractBlazeQueryEnvironment to a factory class, and have BlazeModule optionally expose a custom factory implementation.

--
MOS_MIGRATED_REVID=117546934

6 files changed, 118 insertions, 48 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 27a0bde0ab..19430a4736 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
@@ -23,12 +23,7 @@ import com.google.devtools.build.lib.events.ErrorSensingEventHandler;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.packages.DependencyFilter;
-import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
-import com.google.devtools.build.lib.pkgcache.PackageProvider;
-import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
-import com.google.devtools.build.lib.pkgcache.TargetPatternEvaluator;
-import com.google.devtools.build.lib.pkgcache.TransitivePackageLoader;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
@@ -37,7 +32,6 @@ 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.QueryUtil.AggregateAllCallback;
 import com.google.devtools.build.lib.util.Preconditions;
-import com.google.devtools.build.skyframe.WalkableGraph.WalkableGraphFactory;
 
 import java.util.Collection;
 import java.util.HashMap;
@@ -48,8 +42,6 @@ import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
 
-import javax.annotation.Nullable;
-
 /**
  * {@link QueryEnvironment} that can evaluate queries to produce a result, and implements as much
  * of QueryEnvironment as possible while remaining mostly agnostic as to the objects being stored.
@@ -66,7 +58,7 @@ public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironme
   private final Set<Setting> settings;
   private final List<QueryFunction> extraFunctions;
 
- private static final Logger LOG = Logger.getLogger(AbstractBlazeQueryEnvironment.class.getName());
+  private static final Logger LOG = Logger.getLogger(AbstractBlazeQueryEnvironment.class.getName());
 
   protected AbstractBlazeQueryEnvironment(boolean keepGoing,
       boolean strictScope,
@@ -98,38 +90,6 @@ public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironme
     return specifiedFilter;
   }
 
-  public static AbstractBlazeQueryEnvironment<Target> newQueryEnvironment(
-      TransitivePackageLoader transitivePackageLoader, WalkableGraphFactory graphFactory,
-      PackageProvider packageProvider,
-      TargetPatternEvaluator targetPatternEvaluator, boolean keepGoing, boolean orderedResults,
-      List<String> universeScope, int loadingPhaseThreads,
-      EventHandler eventHandler, Set<Setting> settings, Iterable<QueryFunction> functions,
-      @Nullable PathPackageLocator packagePath) {
-    return newQueryEnvironment(transitivePackageLoader, graphFactory, packageProvider,
-        targetPatternEvaluator, keepGoing, /*strictScope=*/true, orderedResults,
-        universeScope, loadingPhaseThreads, Rule.ALL_LABELS, eventHandler, settings, functions,
-        packagePath);
-  }
-
-  public static AbstractBlazeQueryEnvironment<Target> newQueryEnvironment(
-      TransitivePackageLoader transitivePackageLoader, WalkableGraphFactory graphFactory,
-      PackageProvider packageProvider,
-      TargetPatternEvaluator targetPatternEvaluator, boolean keepGoing, boolean strictScope,
-      boolean orderedResults, List<String> universeScope, int loadingPhaseThreads,
-      Predicate<Label> labelFilter,
-      EventHandler eventHandler, Set<Setting> settings, Iterable<QueryFunction> functions,
-      @Nullable PathPackageLocator packagePath) {
-    Preconditions.checkNotNull(universeScope);
-    return orderedResults || universeScope.isEmpty() || packagePath == null
-        ? new BlazeQueryEnvironment(transitivePackageLoader, packageProvider,
-        targetPatternEvaluator, keepGoing, strictScope, loadingPhaseThreads,
-        labelFilter, eventHandler, settings, functions)
-        : new SkyQueryEnvironment(
-            keepGoing, strictScope, loadingPhaseThreads, labelFilter, eventHandler, settings,
-            functions, targetPatternEvaluator.getOffset(), graphFactory, universeScope,
-            packagePath);
-  }
-
   /**
    * Evaluate the specified query expression in this environment.
    *
diff --git a/src/main/java/com/google/devtools/build/lib/query2/QueryEnvironmentFactory.java b/src/main/java/com/google/devtools/build/lib/query2/QueryEnvironmentFactory.java
new file mode 100644
index 0000000000..7ecdd5a0ee
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/QueryEnvironmentFactory.java
@@ -0,0 +1,62 @@
+// Copyright 2015 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.query2;
+
+import com.google.common.base.Predicate;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.events.EventHandler;
+import com.google.devtools.build.lib.packages.Target;
+import com.google.devtools.build.lib.pkgcache.PackageProvider;
+import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
+import com.google.devtools.build.lib.pkgcache.TargetPatternEvaluator;
+import com.google.devtools.build.lib.pkgcache.TransitivePackageLoader;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Setting;
+import com.google.devtools.build.lib.util.Preconditions;
+import com.google.devtools.build.skyframe.WalkableGraph.WalkableGraphFactory;
+
+import java.util.List;
+import java.util.Set;
+
+import javax.annotation.Nullable;
+
+/** A factory that creates instances of {@code AbstractBlazeQueryEnvironment<Target>}. */
+public class QueryEnvironmentFactory {
+  /** Creates an appropriate {@link AbstractBlazeQueryEnvironment} based on the given options. */
+  public AbstractBlazeQueryEnvironment<Target> create(
+      TransitivePackageLoader transitivePackageLoader, WalkableGraphFactory graphFactory,
+      PackageProvider packageProvider,
+      TargetPatternEvaluator targetPatternEvaluator, boolean keepGoing, boolean strictScope,
+      boolean orderedResults, List<String> universeScope, int loadingPhaseThreads,
+      Predicate<Label> labelFilter,
+      EventHandler eventHandler, Set<Setting> settings, Iterable<QueryFunction> functions,
+      @Nullable PathPackageLocator packagePath) {
+    Preconditions.checkNotNull(universeScope);
+    if (canUseSkyQuery(orderedResults, universeScope, packagePath)) {
+      return new SkyQueryEnvironment(keepGoing, strictScope, loadingPhaseThreads, labelFilter,
+          eventHandler, settings, functions, targetPatternEvaluator.getOffset(), graphFactory,
+          universeScope, packagePath);
+    } else {
+      return new BlazeQueryEnvironment(transitivePackageLoader, packageProvider,
+          targetPatternEvaluator, keepGoing, strictScope, loadingPhaseThreads, labelFilter,
+          eventHandler, settings, functions);
+    }
+  }
+
+  protected static boolean canUseSkyQuery(boolean orderedResults, List<String> universeScope,
+      @Nullable PathPackageLocator packagePath) {
+    return !orderedResults && !universeScope.isEmpty() && packagePath != null;
+  }
+}
+
diff --git a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
index 25d65ea0f4..b61d6ac59b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
@@ -47,7 +47,8 @@ import com.google.devtools.build.lib.pkgcache.FilteringPolicies;
 import com.google.devtools.build.lib.pkgcache.FilteringPolicy;
 import com.google.devtools.build.lib.pkgcache.PackageProvider;
 import com.google.devtools.build.lib.pkgcache.TargetPatternEvaluator;
-import com.google.devtools.build.lib.query2.AbstractBlazeQueryEnvironment;
+import com.google.devtools.build.lib.query2.BlazeQueryEnvironment;
+import com.google.devtools.build.lib.query2.QueryEnvironmentFactory;
 import com.google.devtools.build.lib.query2.engine.DigraphQueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.Setting;
@@ -93,6 +94,8 @@ import javax.annotation.Nullable;
  * An implementation of the 'genquery' rule.
  */
 public class GenQuery implements RuleConfiguredTargetFactory {
+  private static final QueryEnvironmentFactory QUERY_ENVIRONMENT_FACTORY =
+      new QueryEnvironmentFactory();
   public static final Precomputed<ImmutableList<OutputFormatter>> QUERY_OUTPUT_FORMATTERS =
       new Precomputed<>(SkyKey.create(SkyFunctions.PRECOMPUTED, "query_output_formatters"));
 
@@ -254,6 +257,7 @@ public class GenQuery implements RuleConfiguredTargetFactory {
     return doQuery(queryOptions, packageProvider, labelFilter, evaluator, query, ruleContext);
   }
 
+  @SuppressWarnings("unchecked")
   @Nullable
   private byte[] doQuery(QueryOptions queryOptions, PackageProvider packageProvider,
                          Predicate<Label> labelFilter, TargetPatternEvaluator evaluator,
@@ -284,8 +288,8 @@ public class GenQuery implements RuleConfiguredTargetFactory {
       // All the packages are already loaded at this point, so there is no need
       // to start up many threads. 4 are started up to make good use of multiple
       // cores.
-      queryResult = (DigraphQueryEvalResult<Target>) AbstractBlazeQueryEnvironment
-          .newQueryEnvironment(
+      BlazeQueryEnvironment queryEnvironment = (BlazeQueryEnvironment) QUERY_ENVIRONMENT_FACTORY
+          .create(
               /*transitivePackageLoader=*/null, /*graph=*/null, packageProvider,
               evaluator,
               /*keepGoing=*/false,
@@ -297,7 +301,8 @@ public class GenQuery implements RuleConfiguredTargetFactory {
               getEventHandler(ruleContext),
               settings,
               ImmutableList.<QueryFunction>of(),
-              /*packagePath=*/null).evaluateQuery(query, targets);
+              /*packagePath=*/null);
+      queryResult = (DigraphQueryEvalResult<Target>) queryEnvironment.evaluateQuery(query, targets);
     } catch (SkyframeRestartQueryException e) {
       // Do not emit errors for skyframe restarts. They make output of the ConfiguredTargetFunction
       // inconsistent from run to run, and make detecting legitimate errors more difficult.
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java
index 66739ad603..ff4839ba41 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeModule.java
@@ -30,6 +30,7 @@ import com.google.devtools.build.lib.exec.OutputService;
 import com.google.devtools.build.lib.packages.NoSuchThingException;
 import com.google.devtools.build.lib.packages.PackageFactory;
 import com.google.devtools.build.lib.packages.Preprocessor;
+import com.google.devtools.build.lib.query2.QueryEnvironmentFactory;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryFunction;
 import com.google.devtools.build.lib.query2.output.OutputFormatter;
 import com.google.devtools.build.lib.rules.test.CoverageReportActionFactory;
@@ -361,6 +362,15 @@ public abstract class BlazeModule {
   }
 
   /**
+   * Returns a factory for creating {@link AbstractBlazeQueryEnvironment} objects.
+   * If the module does not provide any {@link QueryEnvironmentFactory}, it should return null. Note
+   * that only one factory per Bazel/Blaze runtime is allowed.
+   */
+  public QueryEnvironmentFactory getQueryEnvironmentFactory() {
+    return null;
+  }
+
+  /**
    * Returns a factory for creating {@link SkyframeExecutor} objects. If the module does not
    * provide any SkyframeExecutorFactory, it returns null. Note that only one factory per
    * Bazel/Blaze runtime is allowed.
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
index 0ae29db126..237277a851 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
@@ -55,6 +55,8 @@ import com.google.devtools.build.lib.profiler.ProfilePhase;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.Profiler.ProfiledTaskKinds;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.query2.AbstractBlazeQueryEnvironment;
+import com.google.devtools.build.lib.query2.QueryEnvironmentFactory;
 import com.google.devtools.build.lib.query2.output.OutputFormatter;
 import com.google.devtools.build.lib.rules.test.CoverageReportActionFactory;
 import com.google.devtools.build.lib.runtime.commands.BuildCommand;
@@ -174,6 +176,7 @@ public final class BlazeRuntime {
   // Workspace state (currently exactly one workspace per server)
   private final BlazeDirectories directories;
   private final SkyframeExecutor skyframeExecutor;
+  private final QueryEnvironmentFactory queryEnvironmentFactory;
   /** The action cache is loaded lazily on the first build command. */
   private ActionCache actionCache;
   /** The execution time range of the previous build command in this server, if any. */
@@ -183,6 +186,7 @@ public final class BlazeRuntime {
   private BlazeRuntime(BlazeDirectories directories,
       WorkspaceStatusAction.Factory workspaceStatusActionFactory,
       final SkyframeExecutor skyframeExecutor,
+      QueryEnvironmentFactory queryEnvironmentFactory,
       PackageFactory pkgFactory, ConfiguredRuleClassProvider ruleClassProvider,
       ConfigurationFactory configurationFactory, Clock clock,
       OptionsProvider startupOptionsProvider, Iterable<BlazeModule> blazeModules,
@@ -206,6 +210,7 @@ public final class BlazeRuntime {
     this.timestampGranularityMonitor = Preconditions.checkNotNull(timestampGranularityMonitor);
     this.startupOptionsProvider = startupOptionsProvider;
     this.eventBusExceptionHandler = eventBusExceptionHandler;
+    this.queryEnvironmentFactory = queryEnvironmentFactory;
 
     // Workspace state
     this.directories = directories;
@@ -472,6 +477,14 @@ public final class BlazeRuntime {
   }
 
   /**
+   * Returns the {@link QueryEnvironmentFactory} that should be used to create a
+   * {@link AbstractBlazeQueryEnvironment}, whenever one is needed.
+   */
+  public QueryEnvironmentFactory getQueryEnvironmentFactory() {
+    return queryEnvironmentFactory;
+  }
+
+  /**
    * Returns the package factory.
    */
   public PackageFactory getPackageFactory() {
@@ -1333,6 +1346,7 @@ public final class BlazeRuntime {
 
       Preprocessor.Factory.Supplier preprocessorFactorySupplier = null;
       SkyframeExecutorFactory skyframeExecutorFactory = null;
+      QueryEnvironmentFactory queryEnvironmentFactory = null;
       for (BlazeModule module : blazeModules) {
         module.blazeStartup(startupOptionsProvider,
             BlazeVersionInfo.instance(), instanceId, directories, clock);
@@ -1350,10 +1364,21 @@ public final class BlazeRuntime {
               skyframeExecutorFactory);
           skyframeExecutorFactory = skyFactory;
         }
+        QueryEnvironmentFactory queryEnvFactory = module.getQueryEnvironmentFactory();
+        if (queryEnvFactory != null) {
+          Preconditions.checkState(queryEnvironmentFactory == null,
+              "At most one query environment factory supported. But found two: %s and %s",
+              queryEnvFactory,
+              queryEnvironmentFactory);
+          queryEnvironmentFactory = queryEnvFactory;
+        }
       }
       if (skyframeExecutorFactory == null) {
         skyframeExecutorFactory = new SequencedSkyframeExecutorFactory();
       }
+      if (queryEnvironmentFactory == null) {
+        queryEnvironmentFactory = new QueryEnvironmentFactory();
+      }
       if (preprocessorFactorySupplier == null) {
         preprocessorFactorySupplier = Preprocessor.Factory.Supplier.NullSupplier.INSTANCE;
       }
@@ -1472,7 +1497,7 @@ public final class BlazeRuntime {
       invocationPolicy = createInvocationPolicyFromModules(invocationPolicy, blazeModules);
 
       return new BlazeRuntime(directories, workspaceStatusActionFactory, skyframeExecutor,
-          pkgFactory, ruleClassProvider, configurationFactory,
+          queryEnvironmentFactory, pkgFactory, ruleClassProvider, configurationFactory,
           clock, startupOptionsProvider, ImmutableList.copyOf(blazeModules),
           timestampMonitor, eventBusExceptionHandler, binTools, projectFileProvider,
           invocationPolicy, commands);
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 32159f413a..824970a29c 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
@@ -13,6 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.runtime.commands;
 
+import static com.google.devtools.build.lib.packages.Rule.ALL_LABELS;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
@@ -250,12 +252,18 @@ public final class QueryCommand implements BlazeCommand {
     for (BlazeModule module : env.getRuntime().getBlazeModules()) {
       functions.addAll(module.getQueryFunctions());
     }
-    return AbstractBlazeQueryEnvironment.newQueryEnvironment(
+    return env.getRuntime().getQueryEnvironmentFactory().create(
         env.getPackageManager().newTransitiveLoader(),
         env.getSkyframeExecutor(),
         env.getPackageManager(),
         env.newTargetPatternEvaluator(),
-        keepGoing, orderedResults, universeScope, loadingPhaseThreads, env.getReporter(),
+        keepGoing,
+        /*strictScope=*/ true,
+        orderedResults,
+        universeScope,
+        loadingPhaseThreads,
+        /*labelFilter=*/ ALL_LABELS,
+        env.getReporter(),
         settings,
         functions.build(),
         env.getPackageManager().getPackagePath());