Move the clientEnv to the CommandEnvironment.

--
MOS_MIGRATED_REVID=103269584

9 files changed, 99 insertions, 87 deletions

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 87a06cbd3a..d492e8434f 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
@@ -39,7 +39,6 @@ import com.google.devtools.build.lib.analysis.WorkspaceStatusAction.Key;
 import com.google.devtools.build.lib.analysis.WorkspaceStatusAction.KeyType;
 import com.google.devtools.build.lib.events.Event;
 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;
@@ -48,6 +47,7 @@ import com.google.devtools.build.lib.shell.CommandResult;
 import com.google.devtools.build.lib.util.CommandBuilder;
 import com.google.devtools.build.lib.util.NetUtil;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
+import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.common.options.OptionsBase;
 
@@ -76,7 +76,8 @@ public class BazelWorkspaceStatusModule extends BlazeModule {
 
     private BazelWorkspaceStatusAction(
         WorkspaceStatusAction.Options options,
-        BlazeRuntime runtime,
+        Map<String, String> clientEnv,
+        Path workspace,
         Artifact stableStatus,
         Artifact volatileStatus) {
       super(BuildInfoHelper.BUILD_INFO_ACTION_OWNER, Artifact.NO_ARTIFACTS,
@@ -95,8 +96,8 @@ public class BazelWorkspaceStatusModule extends BlazeModule {
                   // Pass client env, because certain SCM client(like
                   // perforce, git) relies on environment variables to work
                   // correctly.
-                  .setEnv(runtime.getClientEnv())
-                  .setWorkingDir(runtime.getWorkspace())
+                  .setEnv(clientEnv)
+                  .setWorkingDir(workspace)
                   .useShell(true)
                   .build();
     }
@@ -219,14 +220,15 @@ public class BazelWorkspaceStatusModule extends BlazeModule {
     @Override
     public WorkspaceStatusAction createWorkspaceStatusAction(
         ArtifactFactory factory, ArtifactOwner artifactOwner, Supplier<UUID> buildId) {
-      Root root = runtime.getDirectories().getBuildDataDirectory();
+      Root root = env.getDirectories().getBuildDataDirectory();
 
       Artifact stableArtifact = factory.getDerivedArtifact(
           new PathFragment("stable-status.txt"), root, artifactOwner);
       Artifact volatileArtifact = factory.getConstantMetadataArtifact(
           new PathFragment("volatile-status.txt"), root, artifactOwner);
 
-      return new BazelWorkspaceStatusAction(options, runtime, stableArtifact, volatileArtifact);
+      return new BazelWorkspaceStatusAction(options, env.getClientEnv(),
+          env.getDirectories().getWorkspace(), stableArtifact, volatileArtifact);
     }
   }
 
@@ -255,16 +257,22 @@ public class BazelWorkspaceStatusModule extends BlazeModule {
     }
   }
 
-  private BlazeRuntime runtime;
+  private CommandEnvironment env;
   private WorkspaceStatusAction.Options options;
 
   @Override
   public void beforeCommand(Command command, CommandEnvironment env) {
-    this.runtime = env.getRuntime();
+    this.env = env;
     env.getEventBus().register(this);
   }
 
   @Override
+  public void afterCommand() {
+    this.env = null;
+    this.options = null;
+  }
+
+  @Override
   public Iterable<Class<? extends OptionsBase>> getCommandOptions(Command command) {
     return command.builds()
         ? ImmutableList.<Class<? extends OptionsBase>>of(WorkspaceStatusAction.Options.class)
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 5a63d34c50..b1bdaaee2a 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
@@ -31,7 +31,6 @@ import com.google.devtools.build.lib.rules.cpp.CppLinkActionContext;
 import com.google.devtools.build.lib.rules.cpp.IncludeScanningContext;
 import com.google.devtools.build.lib.rules.genquery.GenQuery;
 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;
@@ -39,7 +38,6 @@ import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.common.options.Converters.AssignmentConverter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionsBase;
-import com.google.devtools.common.options.OptionsProvider;
 
 import java.util.List;
 import java.util.Map;
@@ -128,16 +126,22 @@ public class BazelRulesModule extends BlazeModule {
     }
   }
 
-  private BlazeRuntime runtime;
-  private OptionsProvider optionsProvider;
+  private CommandEnvironment env;
+  private BazelExecutionOptions options;
 
   @Override
   public void beforeCommand(Command command, CommandEnvironment env) {
-    this.runtime = env.getRuntime();
+    this.env = env;
     env.getEventBus().register(this);
   }
 
   @Override
+  public void afterCommand() {
+    this.env = null;
+    this.options = null;
+  }
+
+  @Override
   public Iterable<Class<? extends OptionsBase>> getCommandOptions(Command command) {
     return command.builds()
         ? ImmutableList.<Class<? extends OptionsBase>>of(BazelExecutionOptions.class)
@@ -147,18 +151,18 @@ public class BazelRulesModule extends BlazeModule {
   @Override
   public Iterable<ActionContextProvider> getActionContextProviders() {
     return ImmutableList.<ActionContextProvider>of(new SimpleActionContextProvider(
-        new WriteAdbArgsActionContext(runtime.getClientEnv().get("HOME"))));
+        new WriteAdbArgsActionContext(env.getClientEnv().get("HOME"))));
   }
 
   @Override
   public Iterable<ActionContextConsumer> getActionContextConsumers() {
-    return ImmutableList.<ActionContextConsumer>of(new BazelActionContextConsumer(
-        optionsProvider.getOptions(BazelExecutionOptions.class)));
+    return ImmutableList.<ActionContextConsumer>of(
+        new BazelActionContextConsumer(options));
   }
 
   @Subscribe
   public void gotOptions(GotOptionsEvent event) {
-    optionsProvider = event.getOptions();
+    options = event.getOptions().getOptions(BazelExecutionOptions.class);
   }
 
   @Override
@@ -173,7 +177,7 @@ public class BazelRulesModule extends BlazeModule {
         new Supplier<ImmutableList<OutputFormatter>>() {
           @Override
           public ImmutableList<OutputFormatter> get() {
-            return runtime.getQueryOutputFormatters();
+            return env.getRuntime().getQueryOutputFormatters();
           }
         }));
   }
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 9d77e80fd8..aedbecadc0 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
@@ -125,10 +125,7 @@ import java.io.StringWriter;
 import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.Date;
-import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -190,8 +187,6 @@ public final class BlazeRuntime {
 
   private final AtomicInteger storedExitCode = new AtomicInteger();
 
-  private final Map<String, String> clientEnv;
-
   // We pass this through here to make it available to the MasterLogWriter.
   private final OptionsProvider startupOptionsProvider;
 
@@ -212,7 +207,6 @@ public final class BlazeRuntime {
       PackageFactory pkgFactory, ConfiguredRuleClassProvider ruleClassProvider,
       ConfigurationFactory configurationFactory, Clock clock,
       OptionsProvider startupOptionsProvider, Iterable<BlazeModule> blazeModules,
-      Map<String, String> clientEnv,
       TimestampGranularityMonitor timestampGranularityMonitor,
       SubscriberExceptionHandler eventBusExceptionHandler,
       BinTools binTools, ProjectFile.Provider projectFileProvider) {
@@ -230,8 +224,6 @@ public final class BlazeRuntime {
         skyframeExecutor.getPackageManager(),
         pkgFactory.getRuleClassNames());
 
-    this.clientEnv = clientEnv;
-
     this.blazeModules = blazeModules;
     this.ruleClassProvider = ruleClassProvider;
     this.configurationFactory = configurationFactory;
@@ -683,7 +675,7 @@ public final class BlazeRuntime {
     }
     loadingPhaseRunner.updatePatternEvaluator(workingDirectory.relativeTo(workspace));
 
-    updateClientEnv(options.clientEnv, options.ignoreClientEnv);
+    env.updateClientEnv(options.clientEnv, options.ignoreClientEnv);
 
     // Fail fast in the case where a Blaze command forgets to install the package path correctly.
     skyframeExecutor.setActive(false);
@@ -727,12 +719,13 @@ public final class BlazeRuntime {
       try {
         for (Map.Entry<String, String> entry : testEnv.entrySet()) {
           if (entry.getValue() == null) {
-            String clientValue = clientEnv.get(entry.getKey());
+            String clientValue = env.getClientEnv().get(entry.getKey());
             if (clientValue != null) {
               optionsParser.parse(OptionPriority.SOFTWARE_REQUIREMENT,
                   "test environment variable from client environment",
                   ImmutableList.of(
-                      "--test_env=" + entry.getKey() + "=" + clientEnv.get(entry.getKey())));
+                      "--test_env=" + entry.getKey() + "="
+                          + env.getClientEnv().get(entry.getKey())));
             }
           }
         }
@@ -745,7 +738,7 @@ public final class BlazeRuntime {
     }
 
     env.getEventBus().post(
-        new CommandStartEvent(command.name(), commandId, clientEnv, workingDirectory));
+        new CommandStartEvent(command.name(), commandId, env.getClientEnv(), workingDirectory));
     // Initialize exit code to dummy value for afterCommand.
     storedExitCode.set(ExitCode.RESERVED.getNumericExitCode());
   }
@@ -803,26 +796,6 @@ public final class BlazeRuntime {
   }
 
   /**
-   * Return an unmodifiable view of the blaze client's environment when it
-   * invoked the most recent command. Updates from future requests will be
-   * accessible from this view.
-   */
-  public Map<String, String> getClientEnv() {
-    return Collections.unmodifiableMap(clientEnv);
-  }
-
-  @VisibleForTesting
-  void updateClientEnv(List<Map.Entry<String, String>> clientEnvList, boolean ignoreClientEnv) {
-    clientEnv.clear();
-
-    Collection<Map.Entry<String, String>> env =
-        ignoreClientEnv ? System.getenv().entrySet() : clientEnvList;
-    for (Map.Entry<String, String> entry : env) {
-      clientEnv.put(entry.getKey(), entry.getValue());
-    }
-  }
-
-  /**
    * Returns the Clock-instance used for the entire build. Before,
    * individual classes (such as Profiler) used to specify the type
    * of clock (e.g. EpochClock) they wanted to use. This made it
@@ -1500,7 +1473,6 @@ public final class BlazeRuntime {
       UUID instanceId =  (this.instanceId == null) ? UUID.randomUUID() : this.instanceId;
 
       Preconditions.checkNotNull(clock);
-      Map<String, String> clientEnv = new HashMap<>();
       TimestampGranularityMonitor timestampMonitor = new TimestampGranularityMonitor(clock);
 
       Preprocessor.Factory.Supplier preprocessorFactorySupplier = null;
@@ -1653,8 +1625,7 @@ public final class BlazeRuntime {
       return new BlazeRuntime(directories, reporter, workspaceStatusActionFactory, skyframeExecutor,
           pkgFactory, ruleClassProvider, configurationFactory,
           clock, startupOptionsProvider, ImmutableList.copyOf(blazeModules),
-          clientEnv, timestampMonitor,
-          eventBusExceptionHandler, binTools, projectFileProvider);
+          timestampMonitor, eventBusExceptionHandler, binTools, projectFileProvider);
     }
 
     public Builder setBinTools(BinTools binTools) {
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 1edd606240..5d402a1563 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
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.runtime;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.eventbus.EventBus;
@@ -36,6 +37,10 @@ import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.common.options.OptionsProvider;
 
 import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
@@ -48,6 +53,7 @@ public final class CommandEnvironment {
   private final Reporter reporter;
   private final EventBus eventBus;
   private final BlazeModule.ModuleEnvironment blazeModuleEnvironment;
+  private final Map<String, String> clientEnv = new HashMap<>();
 
   private AbruptExitException pendingException;
 
@@ -98,8 +104,23 @@ public final class CommandEnvironment {
     return blazeModuleEnvironment;
   }
 
+  /**
+   * Return an unmodifiable view of the blaze client's environment when it invoked the current
+   * command.
+   */
   public Map<String, String> getClientEnv() {
-    return runtime.getClientEnv();
+    return Collections.unmodifiableMap(clientEnv);
+  }
+
+  @VisibleForTesting
+  void updateClientEnv(List<Map.Entry<String, String>> clientEnvList, boolean ignoreClientEnv) {
+    Preconditions.checkState(clientEnv.isEmpty());
+
+    Collection<Map.Entry<String, String>> env =
+        ignoreClientEnv ? System.getenv().entrySet() : clientEnvList;
+    for (Map.Entry<String, String> entry : env) {
+      clientEnv.put(entry.getKey(), entry.getValue());
+    }
   }
 
   public PackageManager getPackageManager() {
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java
index 263758c6b0..67b769cf43 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/RunCommand.java
@@ -305,7 +305,7 @@ public class RunCommand implements BlazeCommand  {
         null, "Running command line: " + ShellEscaper.escapeJoinAll(prettyCmdLine)));
 
     com.google.devtools.build.lib.shell.Command command = new CommandBuilder()
-        .addArgs(cmdLine).setEnv(runtime.getClientEnv()).setWorkingDir(workingDir).build();
+        .addArgs(cmdLine).setEnv(env.getClientEnv()).setWorkingDir(workingDir).build();
 
     try {
       // Restore a raw EventHandler if it is registered. This allows for blaze run to produce the
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/SandboxActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/sandbox/SandboxActionContextProvider.java
index 61cd2cdbd1..65b35c9cdc 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/SandboxActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/SandboxActionContextProvider.java
@@ -19,7 +19,7 @@ import com.google.devtools.build.lib.actions.ActionContextProvider;
 import com.google.devtools.build.lib.actions.Executor.ActionContext;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.exec.ExecutionOptions;
-import com.google.devtools.build.lib.runtime.BlazeRuntime;
+import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.util.OS;
 
 import java.util.concurrent.ExecutorService;
@@ -33,7 +33,7 @@ public class SandboxActionContextProvider extends ActionContextProvider {
   private final ImmutableList<ActionContext> strategies;
 
   public SandboxActionContextProvider(
-      BlazeRuntime runtime, BuildRequest buildRequest, ExecutorService backgroundWorkers) {
+      CommandEnvironment env, BuildRequest buildRequest, ExecutorService backgroundWorkers) {
     boolean verboseFailures = buildRequest.getOptions(ExecutionOptions.class).verboseFailures;
     boolean sandboxDebug = buildRequest.getOptions(SandboxOptions.class).sandboxDebug;
     Builder<ActionContext> strategies = ImmutableList.builder();
@@ -41,8 +41,8 @@ public class SandboxActionContextProvider extends ActionContextProvider {
     if (OS.getCurrent() == OS.LINUX) {
       strategies.add(
           new LinuxSandboxedStrategy(
-              runtime.getClientEnv(),
-              runtime.getDirectories(),
+              env.getClientEnv(),
+              env.getDirectories(),
               backgroundWorkers,
               verboseFailures,
               sandboxDebug));
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/SandboxModule.java b/src/main/java/com/google/devtools/build/lib/sandbox/SandboxModule.java
index d9e0dc988b..594d79d38b 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/SandboxModule.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/SandboxModule.java
@@ -36,30 +36,33 @@ import java.util.concurrent.TimeUnit;
  * This module provides the Sandbox spawn strategy.
  */
 public class SandboxModule extends BlazeModule {
-  private final ExecutorService backgroundWorkers = Executors.newCachedThreadPool();
-  private BuildRequest buildRequest;
-  private CommandEnvironment env;
-  private BlazeRuntime runtime;
-  private Boolean sandboxingSupported = null;
-
   public static final String SANDBOX_NOT_SUPPORTED_MESSAGE =
       "Sandboxed execution is not supported on your system and thus hermeticity of actions cannot "
           + "be guaranteed. See http://bazel.io/docs/bazel-user-manual.html#sandboxing for more "
           + "information. You can turn off this warning via --ignore_unsupported_sandboxing";
 
-  @Override
-  public Iterable<ActionContextProvider> getActionContextProviders() {
-    Preconditions.checkNotNull(buildRequest);
-    Preconditions.checkNotNull(env);
+  // Per-server state
+  private final ExecutorService backgroundWorkers = Executors.newCachedThreadPool();
+  private Boolean sandboxingSupported = null;
+
+  // Per-command state
+  private CommandEnvironment env;
+  private BuildRequest buildRequest;
 
-    // Cache
+  private synchronized boolean isSandboxingSupported(BlazeRuntime runtime) {
     if (sandboxingSupported == null) {
       sandboxingSupported = NamespaceSandboxRunner.isSupported(runtime);
     }
+    return sandboxingSupported.booleanValue();
+  }
 
-    if (sandboxingSupported) {
+  @Override
+  public Iterable<ActionContextProvider> getActionContextProviders() {
+    Preconditions.checkNotNull(buildRequest);
+    Preconditions.checkNotNull(env);
+    if (isSandboxingSupported(env.getRuntime())) {
       return ImmutableList.<ActionContextProvider>of(
-          new SandboxActionContextProvider(runtime, buildRequest, backgroundWorkers));
+          new SandboxActionContextProvider(env, buildRequest, backgroundWorkers));
     }
 
     // For now, sandboxing is only supported on Linux and there's not much point in showing a scary
@@ -75,15 +78,9 @@ public class SandboxModule extends BlazeModule {
   @Override
   public Iterable<ActionContextConsumer> getActionContextConsumers() {
     Preconditions.checkNotNull(env);
-
-    if (sandboxingSupported == null) {
-      sandboxingSupported = NamespaceSandboxRunner.isSupported(runtime);
-    }
-
-    if (sandboxingSupported) {
+    if (isSandboxingSupported(env.getRuntime())) {
       return ImmutableList.<ActionContextConsumer>of(new SandboxActionContextConsumer());
     }
-
     return ImmutableList.of();
   }
 
@@ -97,10 +94,15 @@ public class SandboxModule extends BlazeModule {
   @Override
   public void beforeCommand(Command command, CommandEnvironment env) {
     this.env = env;
-    this.runtime = env.getRuntime();
     env.getEventBus().register(this);
   }
 
+  @Override
+  public void afterCommand() {
+    this.env = null;
+    this.buildRequest = null;
+  }
+
   @Subscribe
   public void buildStarting(BuildStartingEvent event) {
     buildRequest = event.getRequest();
diff --git a/src/main/java/com/google/devtools/build/lib/standalone/StandaloneActionContextProvider.java b/src/main/java/com/google/devtools/build/lib/standalone/StandaloneActionContextProvider.java
index 7fe35f69c7..5639363e4a 100644
--- a/src/main/java/com/google/devtools/build/lib/standalone/StandaloneActionContextProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/standalone/StandaloneActionContextProvider.java
@@ -32,6 +32,7 @@ import com.google.devtools.build.lib.rules.test.ExclusiveTestStrategy;
 import com.google.devtools.build.lib.rules.test.StandaloneTestStrategy;
 import com.google.devtools.build.lib.rules.test.TestActionContext;
 import com.google.devtools.build.lib.runtime.BlazeRuntime;
+import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 
 import java.io.IOException;
@@ -63,13 +64,13 @@ public class StandaloneActionContextProvider extends ActionContextProvider {
   private final ImmutableList<ActionContext> strategies;
   private final BlazeRuntime runtime;
 
-  public StandaloneActionContextProvider(BlazeRuntime runtime, BuildRequest buildRequest) {
+  public StandaloneActionContextProvider(CommandEnvironment env, BuildRequest buildRequest) {
     boolean verboseFailures = buildRequest.getOptions(ExecutionOptions.class).verboseFailures;
 
-    this.runtime = runtime;
+    this.runtime = env.getRuntime();
 
     TestActionContext testStrategy = new StandaloneTestStrategy(buildRequest,
-        runtime.getStartupOptionsProvider(), runtime.getBinTools(), runtime.getClientEnv(),
+        runtime.getStartupOptionsProvider(), runtime.getBinTools(), env.getClientEnv(),
         runtime.getWorkspace());
 
     Builder<ActionContext> strategiesBuilder = ImmutableList.builder();
diff --git a/src/main/java/com/google/devtools/build/lib/standalone/StandaloneModule.java b/src/main/java/com/google/devtools/build/lib/standalone/StandaloneModule.java
index e880043ea0..b3a0a3e04c 100644
--- a/src/main/java/com/google/devtools/build/lib/standalone/StandaloneModule.java
+++ b/src/main/java/com/google/devtools/build/lib/standalone/StandaloneModule.java
@@ -19,7 +19,6 @@ import com.google.devtools.build.lib.actions.ActionContextProvider;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.buildtool.buildevent.BuildStartingEvent;
 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;
 
@@ -27,21 +26,27 @@ import com.google.devtools.build.lib.runtime.CommandEnvironment;
  * StandaloneModule provides pluggable functionality for blaze.
  */
 public class StandaloneModule extends BlazeModule {
+  private CommandEnvironment env;
   private BuildRequest buildRequest;
-  private BlazeRuntime runtime;
 
   @Override
   public Iterable<ActionContextProvider> getActionContextProviders() {
     return ImmutableList.<ActionContextProvider>of(
-        new StandaloneActionContextProvider(runtime, buildRequest));
+        new StandaloneActionContextProvider(env, buildRequest));
   }
 
   @Override
   public void beforeCommand(Command command, CommandEnvironment env) {
-    this.runtime = env.getRuntime();
+    this.env = env;
     env.getEventBus().register(this);
   }
 
+  @Override
+  public void afterCommand() {
+    this.env = null;
+    this.buildRequest = null;
+  }
+
   @Subscribe
   public void buildStarting(BuildStartingEvent event) {
     buildRequest = event.getRequest();