BlazeModule.workspaceInit: also pass the BlazeRuntime to the call

In order for BlazeModule.workspaceInit to be self-contained, also pass in the
BlazeRuntime; we have use cases where this context is relevant, and there's
currently no other way to get a reference to the BlazeRuntime.

PiperOrigin-RevId: 158861142

6 files changed, 17 insertions, 30 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BazelDiffAwarenessModule.java b/src/main/java/com/google/devtools/build/lib/bazel/BazelDiffAwarenessModule.java
index b5a85d3c4f..9d4719ef2d 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/BazelDiffAwarenessModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/BazelDiffAwarenessModule.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.bazel;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.runtime.BlazeModule;
+import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
 import com.google.devtools.build.lib.skyframe.DiffAwareness;
 import com.google.devtools.build.lib.skyframe.LocalDiffAwareness;
@@ -26,7 +27,8 @@ import com.google.devtools.common.options.OptionsBase;
  */
 public class BazelDiffAwarenessModule extends BlazeModule {
   @Override
-  public void workspaceInit(BlazeDirectories directories, WorkspaceBuilder builder) {
+  public void workspaceInit(
+      BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
     // Order here is important - LocalDiffAwareness creation always succeeds, so it must be last.
     builder.addDiffAwarenessFactory(new LocalDiffAwareness.Factory(ImmutableList.<String>of()));
   }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
index 53baececcf..23eec44981 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
@@ -61,6 +61,7 @@ import com.google.devtools.build.lib.rules.repository.RepositoryDirectoryValue;
 import com.google.devtools.build.lib.rules.repository.RepositoryFunction;
 import com.google.devtools.build.lib.rules.repository.RepositoryLoaderFunction;
 import com.google.devtools.build.lib.runtime.BlazeModule;
+import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import com.google.devtools.build.lib.runtime.Command;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.runtime.ServerBuilder;
@@ -153,7 +154,8 @@ public class BazelRepositoryModule extends BlazeModule {
   }
 
   @Override
-  public void workspaceInit(BlazeDirectories directories, WorkspaceBuilder builder) {
+  public void workspaceInit(
+      BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
     builder.addCustomDirtinessChecker(REPOSITORY_VALUE_CHECKER);
     // Create the repository function everything flows through.
     builder.addSkyFunction(SkyFunctions.REPOSITORY, new RepositoryLoaderFunction());
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/BazelWorkspaceStatusModule.java b/src/main/java/com/google/devtools/build/lib/bazel/BazelWorkspaceStatusModule.java
index 3e079e7701..186c42973d 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/BazelWorkspaceStatusModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/BazelWorkspaceStatusModule.java
@@ -41,6 +41,7 @@ import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.exec.ExecutorBuilder;
 import com.google.devtools.build.lib.runtime.BlazeModule;
+import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import com.google.devtools.build.lib.runtime.Command;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.runtime.GotOptionsEvent;
@@ -394,7 +395,8 @@ public class BazelWorkspaceStatusModule extends BlazeModule {
   }
 
   @Override
-  public void workspaceInit(BlazeDirectories directories, WorkspaceBuilder builder) {
+  public void workspaceInit(
+      BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
     builder.setWorkspaceStatusActionFactory(new BazelStatusActionFactory());
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRulesModule.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRulesModule.java
index 38a79ae224..a8daf4ad91 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRulesModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRulesModule.java
@@ -14,17 +14,14 @@
 
 package com.google.devtools.build.lib.bazel.rules;
 
-import com.google.common.base.Supplier;
-import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
 import com.google.devtools.build.lib.bazel.rules.cpp.BazelCppRuleClasses;
-import com.google.devtools.build.lib.query2.output.OutputFormatter;
 import com.google.devtools.build.lib.rules.cpp.FdoSupportFunction;
 import com.google.devtools.build.lib.rules.cpp.FdoSupportValue;
 import com.google.devtools.build.lib.rules.genquery.GenQuery;
 import com.google.devtools.build.lib.runtime.BlazeModule;
-import com.google.devtools.build.lib.runtime.CommandEnvironment;
+import com.google.devtools.build.lib.runtime.BlazeRuntime;
 import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.build.lib.util.ResourceFileLoader;
@@ -34,20 +31,6 @@ import java.io.IOException;
  * Module implementing the rule set of Bazel.
  */
 public class BazelRulesModule extends BlazeModule {
-
-  private CommandEnvironment env;
-
-  @Override
-  public void beforeCommand(CommandEnvironment env) {
-    this.env = env;
-    env.getEventBus().register(this);
-  }
-
-  @Override
-  public void afterCommand() {
-    this.env = null;
-  }
-
   @Override
   public void initializeRuleClasses(ConfiguredRuleClassProvider.Builder builder) {
     builder.setToolsRepository(BazelRuleClassProvider.TOOLS_REPOSITORY);
@@ -65,15 +48,11 @@ public class BazelRulesModule extends BlazeModule {
   }
 
   @Override
-  public void workspaceInit(BlazeDirectories directories, WorkspaceBuilder builder) {
+  public void workspaceInit(
+      BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
     builder.addSkyFunction(FdoSupportValue.SKYFUNCTION, new FdoSupportFunction());
     builder.addPrecomputedValue(PrecomputedValue.injected(
         GenQuery.QUERY_OUTPUT_FORMATTERS,
-        new Supplier<ImmutableList<OutputFormatter>>() {
-          @Override
-          public ImmutableList<OutputFormatter> get() {
-            return env.getRuntime().getQueryOutputFormatters();
-          }
-        }));
+        runtime.getQueryOutputFormatters()));
   }
 }
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 2d72ce307b..40b420e9ac 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
@@ -133,10 +133,12 @@ public abstract class BlazeModule {
    * and only if the server initialization was successful. Modules can override this method to
    * affect how the workspace is configured.
    *
+   * @param runtime the blaze runtime
    * @param directories the workspace directories
    * @param builder the workspace builder
    */
-  public void workspaceInit(BlazeDirectories directories, WorkspaceBuilder builder) {
+  public void workspaceInit(
+      BlazeRuntime runtime, BlazeDirectories directories, WorkspaceBuilder builder) {
   }
 
   /**
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 a43ec2b707..94c38debd7 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
@@ -199,7 +199,7 @@ public final class BlazeRuntime {
     Preconditions.checkState(this.workspace == null);
     WorkspaceBuilder builder = new WorkspaceBuilder(directories, binTools);
     for (BlazeModule module : blazeModules) {
-      module.workspaceInit(directories, builder);
+      module.workspaceInit(this, directories, builder);
     }
     this.workspace = builder.build(
         this, packageFactory, ruleClassProvider, getProductName(), eventBusExceptionHandler);