Simplify BlazeModule.beforeCommand

Don't pass the Command annotation explicitly, but add it to CommandEnvironment
instead; most modules don't need it in the first place, so it was a lot of
boilerplate for not much. Also change it so that the command is passed to the
constructor.

Add some documentation to the beforeCommand method.

PiperOrigin-RevId: 158847128

11 files changed, 32 insertions, 31 deletions

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 1d340cde2d..53baececcf 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
@@ -69,7 +69,6 @@ import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue.Injected;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.SkyValueDirtinessChecker;
-import com.google.devtools.build.lib.util.AbruptExitException;
 import com.google.devtools.build.lib.util.io.TimestampGranularityMonitor;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Path;
@@ -215,7 +214,7 @@ public class BazelRepositoryModule extends BlazeModule {
   }
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment env) throws AbruptExitException {
+  public void beforeCommand(CommandEnvironment env) {
     delegator.setClientEnvironment(env.getActionClientEnv());
   }
 
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 812c60a305..3e079e7701 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
@@ -370,7 +370,7 @@ public class BazelWorkspaceStatusModule extends BlazeModule {
   private WorkspaceStatusAction.Options options;
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment env) {
+  public void beforeCommand(CommandEnvironment env) {
     this.env = env;
     env.getEventBus().register(this);
   }
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 7a0c9b9262..38a79ae224 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
@@ -24,12 +24,10 @@ 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.Command;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.runtime.WorkspaceBuilder;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.build.lib.util.ResourceFileLoader;
-
 import java.io.IOException;
 
 /**
@@ -40,7 +38,7 @@ public class BazelRulesModule extends BlazeModule {
   private CommandEnvironment env;
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment env) {
+  public void beforeCommand(CommandEnvironment env) {
     this.env = env;
     env.getEventBus().register(this);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java
index 76c446bd95..ad2b3fd9db 100644
--- a/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java
+++ b/src/main/java/com/google/devtools/build/lib/buildeventservice/BuildEventServiceModule.java
@@ -82,7 +82,7 @@ public abstract class BuildEventServiceModule<T extends BuildEventServiceOptions
   }
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment commandEnvironment)
+  public void beforeCommand(CommandEnvironment commandEnvironment)
       throws AbruptExitException {
     this.commandEnvironment = commandEnvironment;
     this.buildEventRecorder = new BuildEventRecorder();
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
index fedfee1485..0bd677f395 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java
@@ -86,7 +86,7 @@ public final class RemoteModule extends BlazeModule {
   }
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment env) {
+  public void beforeCommand(CommandEnvironment env) {
     this.env = env;
     env.getEventBus().register(this);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandDispatcher.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandDispatcher.java
index 55762d6fa6..7751ca60e0 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandDispatcher.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandDispatcher.java
@@ -348,14 +348,14 @@ public class BlazeCommandDispatcher {
     long execStartTimeNanos = runtime.getClock().nanoTime();
 
     // The initCommand call also records the start time for the timestamp granularity monitor.
-    CommandEnvironment env = runtime.getWorkspace().initCommand();
+    CommandEnvironment env = runtime.getWorkspace().initCommand(commandAnnotation);
     // Record the command's starting time for use by the commands themselves.
     env.recordCommandStartTime(firstContactTime);
 
     AbruptExitException exitCausingException = null;
     for (BlazeModule module : runtime.getBlazeModules()) {
       try {
-        module.beforeCommand(commandAnnotation, env);
+        module.beforeCommand(env);
       } catch (AbruptExitException e) {
         // Don't let one module's complaints prevent the other modules from doing necessary
         // setup. We promised to call beforeCommand exactly once per-module before each command
@@ -511,7 +511,6 @@ public class BlazeCommandDispatcher {
       try {
         // Notify the BlazeRuntime, so it can do some initial setup.
         env.beforeCommand(
-            commandAnnotation,
             optionsParser,
             commonOptions,
             execStartTimeNanos,
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 059160edb3..2d72ce307b 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
@@ -140,10 +140,13 @@ public abstract class BlazeModule {
   }
 
   /**
-   * Called before each command.
+   * Called to notify modules that the given command is about to be executed. This allows capturing
+   * the {@link com.google.common.eventbus.EventBus}, {@link Command}, or {@link OptionsProvider}.
+   *
+   * @param env the command
+   * @throws AbruptExitException modules can throw this exception to abort the command
    */
-  @SuppressWarnings("unused")
-  public void beforeCommand(Command command, CommandEnvironment env) throws AbruptExitException {
+  public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeWorkspace.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeWorkspace.java
index c0eda33b08..1f54096a38 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeWorkspace.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeWorkspace.java
@@ -181,9 +181,9 @@ public final class BlazeWorkspace {
    * <p>This method should be called from the "main" thread on which the command will execute;
    * that thread will receive interruptions if a module requests an early exit.
    */
-  public CommandEnvironment initCommand() {
+  public CommandEnvironment initCommand(Command command) {
     CommandEnvironment env = new CommandEnvironment(
-        runtime, this, new EventBus(eventBusExceptionHandler), Thread.currentThread(), null, null);
+        runtime, this, new EventBus(eventBusExceptionHandler), Thread.currentThread(), command);
     skyframeExecutor.setClientEnv(env.getClientEnv());
     return env;
   }
@@ -193,10 +193,10 @@ public final class BlazeWorkspace {
    * those values are set by {@code CommandEnvironment#beforeCommand()} which is not called for
    * testing. Use ONLY for testing purposes.
    */
-  public CommandEnvironment initCommandForTesting(String commandName, OptionsProvider options) {
+  public CommandEnvironment initCommandForTesting(Command command, OptionsProvider options) {
     CommandEnvironment env = new CommandEnvironment(
         runtime, this, new EventBus(eventBusExceptionHandler), Thread.currentThread(),
-        commandName, options);
+        command, options);
     skyframeExecutor.setClientEnv(env.getClientEnv());
     return env;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java b/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java
index 0061e3f937..ee2ca955e9 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java
@@ -44,7 +44,7 @@ public class BuildSummaryStatsModule extends BlazeModule {
   private boolean discardActions;
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment env) {
+  public void beforeCommand(CommandEnvironment env) {
     this.reporter = env.getReporter();
     this.eventBus = env.getEventBus();
     eventBus.register(this);
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java b/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java
index 922a097855..d1411c70bd 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java
@@ -85,6 +85,7 @@ public final class CommandEnvironment {
   private final Map<String, String> actionClientEnv = new TreeMap<>();
   private final TimestampGranularityMonitor timestampGranularityMonitor;
   private final Thread commandThread;
+  private final Command command;
 
   private String[] crashData;
 
@@ -94,7 +95,6 @@ public final class CommandEnvironment {
   private Path workingDirectory;
   private String workspaceName;
 
-  private String commandName;
   private OptionsProvider options;
 
   private AtomicReference<AbruptExitException> pendingException = new AtomicReference<>();
@@ -126,7 +126,8 @@ public final class CommandEnvironment {
    * commandThread passed is interrupted when a module requests an early exit.
    */
   CommandEnvironment(
-      BlazeRuntime runtime, BlazeWorkspace workspace, EventBus eventBus, Thread commandThread) {
+      BlazeRuntime runtime, BlazeWorkspace workspace, EventBus eventBus, Thread commandThread,
+      Command command) {
     this.runtime = runtime;
     this.workspace = workspace;
     this.directories = workspace.getDirectories();
@@ -134,6 +135,7 @@ public final class CommandEnvironment {
     this.reporter = new Reporter(eventBus);
     this.eventBus = eventBus;
     this.commandThread = commandThread;
+    this.command = command;
     this.blazeModuleEnvironment = new BlazeModuleEnvironment();
     this.timestampGranularityMonitor = new TimestampGranularityMonitor(runtime.getClock());
     // Record the command's starting time again, for use by
@@ -159,12 +161,10 @@ public final class CommandEnvironment {
   @VisibleForTesting
   CommandEnvironment(
       BlazeRuntime runtime, BlazeWorkspace workspace, EventBus eventBus, Thread commandThread,
-      String commandNameForTesting, OptionsProvider optionsForTesting) {
-    this(runtime, workspace, eventBus, commandThread);
-    // Both commandName and options are normally set by beforeCommand(); however this method is not
-    // called in tests (i.e. tests use BlazeRuntimeWrapper). These fields should only be set for
-    // testing.
-    this.commandName = commandNameForTesting;
+      Command command, OptionsProvider optionsForTesting) {
+    this(runtime, workspace, eventBus, commandThread, command);
+    // Options are normally set by beforeCommand(); however this method is not called in tests (i.e.
+    // tests use BlazeRuntimeWrapper). These fields should only be set for testing.
     this.options = optionsForTesting;
   }
 
@@ -203,8 +203,12 @@ public final class CommandEnvironment {
     return Collections.unmodifiableMap(clientEnv);
   }
 
+  public Command getCommand() {
+    return command;
+  }
+
   public String getCommandName() {
-    return commandName;
+    return command.name();
   }
 
   public OptionsProvider getOptions() {
@@ -553,7 +557,6 @@ public final class CommandEnvironment {
    * @throws AbruptExitException if this command is unsuitable to be run as specified
    */
   void beforeCommand(
-      Command command,
       OptionsParser optionsParser,
       CommonCommandOptions options,
       long execStartTimeNanos,
@@ -570,7 +573,6 @@ public final class CommandEnvironment {
         throw new IllegalStateException(e);
       }
     }
-    this.commandName = command.name();
     this.options = optionsParser;
 
     eventBus.post(
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
index ddd07ad36c..e0ef188849 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
@@ -49,7 +49,7 @@ public class WorkerModule extends BlazeModule {
   }
 
   @Override
-  public void beforeCommand(Command command, CommandEnvironment env) {
+  public void beforeCommand(CommandEnvironment env) {
     this.env = env;
     env.getEventBus().register(this);
   }