Enable Bazel commands to exit at any time.

This is the first step on a journey toward allowing commands to AbruptExit
wherever they please, similar to how the user can press Ctrl+C at any time
and we (should) bail out as fast as we can.

By interrupting the command's main thread, we at least offer the command
the ability to see that an error requiring a bail has happened, and it should
trigger at potentially more locations, rather than just between phases.

--
MOS_MIGRATED_REVID=135152330

6 files changed, 531 insertions, 14 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
index 974a63228c..afdde3c365 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
@@ -342,9 +342,19 @@ public final class BuildTool {
       }
       exitCode = e.getExitCode() != null ? e.getExitCode() : ExitCode.BUILD_FAILURE;
     } catch (InterruptedException e) {
-      exitCode = ExitCode.INTERRUPTED;
-      env.getReporter().handle(Event.error("build interrupted"));
-      env.getEventBus().post(new BuildInterruptedEvent());
+      // We may have been interrupted by an error, or the user's interruption may have raced with
+      // an error, so check to see if we should report that error code instead.
+      exitCode = env.getPendingExitCode();
+      if (exitCode == null) {
+        exitCode = ExitCode.INTERRUPTED;
+        env.getReporter().handle(Event.error("build interrupted"));
+        env.getEventBus().post(new BuildInterruptedEvent());
+      } else {
+        // Report the exception from the environment - the exception we're handling here is just an
+        // interruption.
+        reportExceptionError(env.getPendingException());
+        result.setCatastrophe();
+      }
     } catch (TargetParsingException | LoadingFailedException | ViewCreationFailedException e) {
       exitCode = ExitCode.PARSING_FAILURE;
       reportExceptionError(e);
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 126c1274b1..a2e71cc9a4 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
@@ -310,9 +310,7 @@ public abstract class BlazeModule {
         throws NoSuchThingException, InterruptedException, IOException;
 
     /**
-     * Exits Blaze as early as possible. This is currently a hack and should only be called in
-     * event handlers for {@code BuildStartingEvent}, {@code GotOptionsEvent} and
-     * {@code LoadingPhaseCompleteEvent}.
+     * Exits Blaze as early as possible by sending an interrupt to the command's main thread.
      */
     void exit(AbruptExitException exception);
   }
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 2b64b6ed80..e4fad28e57 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
@@ -234,6 +234,12 @@ public final class BlazeRuntime {
     }
   }
 
+  /**
+   * Initializes a CommandEnvironment to execute a command in this server.
+   *
+   * <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() {
     return workspace.initCommand();
   }
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 f536e937a3..98c56c48c8 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
@@ -184,9 +184,15 @@ public final class BlazeWorkspace {
     return lastExecutionRange;
   }
 
+  /**
+   * Initializes a CommandEnvironment to execute a command in this workspace.
+   *
+   * <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() {
-    CommandEnvironment env
-        = new CommandEnvironment(runtime, this, new EventBus(eventBusExceptionHandler));
+    CommandEnvironment env = new CommandEnvironment(
+        runtime, this, new EventBus(eventBusExceptionHandler), Thread.currentThread());
     skyframeExecutor.setClientEnv(env.getClientEnv());
     return env;
   }
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 1208995695..a68226d446 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
@@ -64,6 +64,7 @@ import java.util.TreeMap;
 import java.util.TreeSet;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
 
 /**
  * Encapsulates the state needed for a single command. The environment is dropped when the current
@@ -81,6 +82,7 @@ public final class CommandEnvironment {
   private final Map<String, String> clientEnv = new TreeMap<>();
   private final Set<String> visibleClientEnv = new TreeSet<>();
   private final TimestampGranularityMonitor timestampGranularityMonitor;
+  private final Thread commandThread;
 
   private String[] crashData;
 
@@ -104,17 +106,29 @@ public final class CommandEnvironment {
 
     @Override
     public void exit(AbruptExitException exception) {
-      pendingException.compareAndSet(null, exception);
+      Preconditions.checkNotNull(exception);
+      Preconditions.checkNotNull(exception.getExitCode());
+      if (pendingException.compareAndSet(null, exception)) {
+        // There was no exception, so we're the first one to ask for an exit. Interrupt the command.
+        commandThread.interrupt();
+      }
     }
   }
 
-  CommandEnvironment(BlazeRuntime runtime, BlazeWorkspace workspace, EventBus eventBus) {
+  /**
+   * Creates a new command environment which can be used for executing commands for the given
+   * runtime in the given workspace, which will publish events on the given eventBus. The
+   * commandThread passed is interrupted when a module requests an early exit.
+   */
+  CommandEnvironment(
+      BlazeRuntime runtime, BlazeWorkspace workspace, EventBus eventBus, Thread commandThread) {
     this.runtime = runtime;
     this.workspace = workspace;
     this.directories = workspace.getDirectories();
     this.commandId = null; // Will be set once we get the client environment
     this.reporter = new Reporter();
     this.eventBus = eventBus;
+    this.commandThread = commandThread;
     this.blazeModuleEnvironment = new BlazeModuleEnvironment();
     this.timestampGranularityMonitor = new TimestampGranularityMonitor(runtime.getClock());
     // Record the command's starting time again, for use by
@@ -377,6 +391,29 @@ public final class CommandEnvironment {
   }
 
   /**
+   * Prevents any further interruption of this command by modules, and returns the final exit code
+   * from modules, or null if no modules requested an abrupt exit.
+   *
+   * <p>Always returns the same value on subsequent calls.
+   */
+  @Nullable
+  private ExitCode finalizeExitCode() {
+    // Set the pending exception so that further calls to exit(AbruptExitException) don't lead to
+    // unwanted thread interrupts.
+    if (pendingException.compareAndSet(null, new AbruptExitException(null))) {
+      return null;
+    }
+    if (Thread.currentThread() == commandThread) {
+      // We may have interrupted the thread in the process, so clear the interrupted bit.
+      // Whether the command was interrupted or not, it's about to be over, so don't interrupt later
+      // things happening on this thread.
+      Thread.interrupted();
+    }
+    // Extract the exit code (it can be null if someone has already called finalizeExitCode()).
+    return getPendingExitCode();
+  }
+
+  /**
    * Hook method called by the BlazeCommandDispatcher right before the dispatch
    * of each command ends (while its outcome can still be modified).
    */
@@ -384,14 +421,32 @@ public final class CommandEnvironment {
     eventBus.post(new CommandPrecompleteEvent(originalExit));
     // If Blaze did not suffer an infrastructure failure, check for errors in modules.
     ExitCode exitCode = originalExit;
-    AbruptExitException exception = pendingException.get();
-    if (!originalExit.isInfrastructureFailure() && exception != null) {
-      exitCode = exception.getExitCode();
+    ExitCode newExitCode = finalizeExitCode();
+    if (!originalExit.isInfrastructureFailure() && newExitCode != null) {
+      exitCode = newExitCode;
     }
     return exitCode;
   }
 
   /**
+   * Returns the current exit code requested by modules, or null if no exit has been requested.
+   */
+  @Nullable
+  public ExitCode getPendingExitCode() {
+    AbruptExitException exception = getPendingException();
+    return exception == null ? null : exception.getExitCode();
+  }
+
+  /**
+   * Retrieves the exception currently queued by a Blaze module.
+   *
+   * <p>Prefer getPendingExitCode or throwPendingException where appropriate.
+   */
+  public AbruptExitException getPendingException() {
+    return pendingException.get();
+  }
+
+  /**
    * Throws the exception currently queued by a Blaze module.
    *
    * <p>This should be called as often as is practical so that errors are reported as soon as
@@ -399,8 +454,12 @@ public final class CommandEnvironment {
    * the exception this way.
    */
   public void throwPendingException() throws AbruptExitException {
-    AbruptExitException exception = pendingException.get();
+    AbruptExitException exception = getPendingException();
     if (exception != null) {
+      if (Thread.currentThread() == commandThread) {
+        // Throwing this exception counts as the requested interruption. Clear the interrupted bit.
+        Thread.interrupted();
+      }
       throw exception;
     }
   }
diff --git a/src/test/java/com/google/devtools/build/lib/runtime/CommandInterruptionTest.java b/src/test/java/com/google/devtools/build/lib/runtime/CommandInterruptionTest.java
new file mode 100644
index 0000000000..da71cbfd0e
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/runtime/CommandInterruptionTest.java
@@ -0,0 +1,438 @@
+// Copyright 2016 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.runtime;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.SettableFuture;
+import com.google.devtools.build.lib.analysis.BlazeDirectories;
+import com.google.devtools.build.lib.analysis.ConfigurationCollectionFactory;
+import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider;
+import com.google.devtools.build.lib.analysis.ServerDirectories;
+import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
+import com.google.devtools.build.lib.runtime.BlazeCommandDispatcher.ShutdownBlazeServerException;
+import com.google.devtools.build.lib.testutil.Scratch;
+import com.google.devtools.build.lib.testutil.TestConstants;
+import com.google.devtools.build.lib.testutil.TestUtils;
+import com.google.devtools.build.lib.util.AbruptExitException;
+import com.google.devtools.build.lib.util.ExitCode;
+import com.google.devtools.build.lib.util.io.OutErr;
+import com.google.devtools.common.options.Option;
+import com.google.devtools.common.options.OptionsBase;
+import com.google.devtools.common.options.OptionsParser;
+import com.google.devtools.common.options.OptionsProvider;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.BrokenBarrierException;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+/** Tests of CommandEnvironment's command-interrupting exit functionality. */
+@RunWith(JUnit4.class)
+public final class CommandInterruptionTest {
+
+  /** Options class to pass configuration to our dummy wait command. */
+  public static class WaitOptions extends OptionsBase {
+    public WaitOptions() {}
+
+    @Option(name = "expect_interruption", defaultValue = "false")
+    public boolean expectInterruption;
+  }
+
+  /**
+   * Command which retrieves an exit code off the queue and returns it, or INTERRUPTED if
+   * interrupted more than --expect_interruptions times while waiting.
+   */
+  @Command(
+    name = "snooze",
+    shortDescription = "",
+    help = "",
+    options = {WaitOptions.class}
+  )
+  private static final class WaitForCompletionCommand implements BlazeCommand {
+    private final AtomicBoolean isTestShuttingDown;
+    private final AtomicReference<SettableFuture<CommandState>> commandStateHandoff;
+
+    public WaitForCompletionCommand(AtomicBoolean isTestShuttingDown) {
+      this.isTestShuttingDown = isTestShuttingDown;
+      this.commandStateHandoff = new AtomicReference<>();
+    }
+
+    @Override
+    public ExitCode exec(CommandEnvironment env, OptionsProvider options) {
+      CommandState commandState = new CommandState(
+          env, options.getOptions(WaitOptions.class).expectInterruption, isTestShuttingDown);
+      commandStateHandoff.getAndSet(null).set(commandState);
+      return commandState.waitForExitCodeFromTest();
+    }
+
+    @Override
+    public void editOptions(CommandEnvironment env, OptionsParser optionsParser) {}
+
+    /**
+     * Runs an instance of this command on the given executor, waits for it to start and returns a
+     * CommandState which can be used to control and assert on the status of that command.
+     */
+    public CommandState runIn(
+        ExecutorService executor, BlazeCommandDispatcher dispatcher, boolean expectInterruption)
+        throws InterruptedException, ExecutionException {
+      SettableFuture<CommandState> newHandoff = SettableFuture.create();
+      if (!commandStateHandoff.compareAndSet(null, newHandoff)) {
+        throw new AssertionError("Another command is already starting at this time?!");
+      }
+      executor.submit(
+          new RunCommandThroughDispatcher(dispatcher, newHandoff, expectInterruption));
+      return newHandoff.get();
+    }
+  }
+
+  /** Callable to run the above command on a different thread. */
+  private static final class RunCommandThroughDispatcher implements Callable<Integer> {
+    private final BlazeCommandDispatcher dispatcher;
+    private final Future<CommandState> commandStateHandoff;
+    private final boolean expectInterruption;
+
+    public RunCommandThroughDispatcher(
+        BlazeCommandDispatcher dispatcher, Future<CommandState> commandStateHandoff,
+        boolean expectInterruption) {
+      this.dispatcher = dispatcher;
+      this.commandStateHandoff = commandStateHandoff;
+      this.expectInterruption = expectInterruption;
+    }
+
+    @Override
+    public Integer call()
+        throws ShutdownBlazeServerException, InterruptedException, ExecutionException {
+      int result = dispatcher.exec(
+          ImmutableList.of(
+              "snooze",
+              expectInterruption ? "--expect_interruption" : "--noexpect_interruption"),
+          BlazeCommandDispatcher.LockingMode.ERROR_OUT,
+          "CommandInterruptionTest",
+          OutErr.SYSTEM_OUT_ERR);
+      // TODO(mstaib): replace with Futures.getDone when Bazel uses Guava 20.0
+      commandStateHandoff.get().completeWithExitCode(result);
+      return result;
+    }
+  }
+
+  /**
+   * A remote control allowing the test to control and assert on the WaitForCompletionCommand.
+   */
+  private static final class CommandState {
+    private final SettableFuture<Integer> result;
+    private final CommandEnvironment commandEnvironment;
+    private final Thread thread;
+    private final BlockingQueue<ExitCode> exitCodeQueue;
+    private final AtomicBoolean isTestShuttingDown;
+    private boolean expectInterruption;
+    private final CyclicBarrier barrier;
+
+    private static final ExitCode SENTINEL =
+        ExitCode.createInfrastructureFailure(-1, "GO TO THE BARRIER");
+
+    public CommandState(
+        CommandEnvironment commandEnvironment, boolean expectInterruption,
+        AtomicBoolean isTestShuttingDown) {
+      this.result = SettableFuture.create();
+      this.commandEnvironment = commandEnvironment;
+      this.thread = Thread.currentThread();
+      this.exitCodeQueue = new ArrayBlockingQueue<ExitCode>(1);
+      this.isTestShuttingDown = isTestShuttingDown;
+      this.expectInterruption = expectInterruption;
+      this.barrier = new CyclicBarrier(2);
+    }
+
+    // command side
+
+    /**
+     * Marks the Future associated with this CommandState completed with the given exit code, then
+     * waits at the barrier for the test thread to catch up.
+     */
+    private void completeWithExitCode(int exitCode) {
+      result.set(exitCode);
+      if (!isTestShuttingDown.get()) {
+        // Wait at the barrier for the test to assert on status, unless the test is shutting down.
+        try {
+          barrier.await();
+        } catch (InterruptedException | BrokenBarrierException ex) {
+          // this is fine, we're only doing this for the test thread's benefit anyway
+        }
+      }
+    }
+
+    /**
+     * Waits for an exit code to come from the test, either INTERRUPTED via thread interruption, or
+     * a test-specified exit code via requestExitWith(). If expectInterruption was set,
+     * a single interruption will be ignored.
+     */
+    private ExitCode waitForExitCodeFromTest() {
+      while (true) {
+        ExitCode exitCode = null;
+        try {
+          exitCode = exitCodeQueue.take();
+          if (Thread.interrupted()) {
+            // the interruption and the exit code delivery may have come in simultaneously, which
+            // may result in a successful return from the queue with interrupted() set.
+            throw new InterruptedException();
+          }
+        } catch (InterruptedException ex) {
+          if (!expectInterruption || isTestShuttingDown.get()) {
+            // This is not an expected interruption (possibly because the test is shutting down and
+            // it's the executor's please stop interruption) so give up.
+            return ExitCode.INTERRUPTED;
+          }
+          // Otherwise, that was an expected interruption, so return to looking for exit codes.
+          // But we only expect one, so the next one will be fatal.
+          expectInterruption = false;
+          // We fall through the catch here in case we received an interruption and an exit code at
+          // the same time.
+        }
+
+        if (exitCode == SENTINEL) {
+          // The test just wants us to go wait at the barrier for an assertion.
+          try {
+            barrier.await();
+          } catch (InterruptedException | BrokenBarrierException impossible) {
+            // This should not happen in normal use, but if it does, exit gracefully so
+            // BlazeCommandDispatcher has a chance to clean up. Use the SENTINEL value to avoid
+            // accidentally passing any tests that might have been looking for INTERRUPTED.
+            return SENTINEL;
+          }
+          continue;
+        } else if (exitCode != null) {
+          return exitCode;
+        }
+      }
+    }
+
+    // test side
+
+    /** Gets the ModuleEnvironment modules will see when executing this command. */
+    public BlazeModule.ModuleEnvironment getModuleEnvironment() {
+      return commandEnvironment.getBlazeModuleEnvironment();
+    }
+
+    /** Sends an exit code to the command, which will then return with it if it is still running. */
+    public void requestExitWith(ExitCode exitCode) {
+      exitCodeQueue.offer(exitCode);
+    }
+
+    /** Sends an interrupt directly to the command's thread. */
+    public void interrupt() {
+      thread.interrupt();
+    }
+
+    /** Waits for the command to reach a stopping point to check if it has finished or not. */
+    private void synchronizeWithCommand() throws InterruptedException, BrokenBarrierException {
+      // If the future is already done, no need to wait at the barrier - we already know the state.
+      if (result.isDone()) {
+        // But if the command thread is waiting on the barrier, tell it to stop doing so.
+        barrier.reset();
+        return;
+      }
+      // Offer the sentinel to the queue - if the command is still waiting and it sees the sentinel,
+      // it will go to the barrier.
+      exitCodeQueue.offer(SENTINEL);
+      // Then wait for the command to finish processing.
+      barrier.await();
+    }
+
+    /** Asserts that the command finished and returned the given ExitCode. */
+    public void assertFinishedWith(ExitCode exitCode)
+        throws InterruptedException, ExecutionException, BrokenBarrierException {
+      synchronizeWithCommand();
+      assertWithMessage("The command should have been finished, but it was not.")
+          .that(result.isDone()).isTrue();
+      // TODO(mstaib): replace with Futures.getDone when Bazel uses Guava 20.0
+      assertThat(result.get()).isEqualTo(exitCode.getNumericExitCode());
+    }
+
+    /** Asserts that the command has not finished yet. */
+    public void assertNotFinishedYet()
+        throws InterruptedException, ExecutionException, BrokenBarrierException {
+      synchronizeWithCommand();
+      assertWithMessage("The command should not have been finished, but it was.")
+          .that(result.isDone()).isFalse();
+    }
+
+    /** Asserts that both commands were executed on the same thread. */
+    public void assertOnSameThreadAs(CommandState other) {
+      assertThat(thread).isSameAs(other.thread);
+    }
+  }
+
+  private ExecutorService executor;
+  private AtomicBoolean isTestShuttingDown;
+  private BlazeCommandDispatcher dispatcher;
+  private WaitForCompletionCommand snooze;
+
+  @Before
+  public void setUp() throws Exception {
+    executor = Executors.newSingleThreadExecutor();
+    Scratch scratch = new Scratch();
+    isTestShuttingDown = new AtomicBoolean(false);
+    String productName = TestConstants.PRODUCT_NAME;
+    ServerDirectories serverDirectories =
+        new ServerDirectories(scratch.dir("install"), scratch.dir("output"));
+    BlazeRuntime runtime =
+        new BlazeRuntime.Builder()
+            .setProductName(productName)
+            .setServerDirectories(serverDirectories)
+            .setStartupOptionsProvider(
+                OptionsParser.newOptionsParser(BlazeServerStartupOptions.class))
+            .addBlazeModule(
+                new BlazeModule() {
+                  @Override
+                  public void initializeRuleClasses(ConfiguredRuleClassProvider.Builder builder) {
+                    // Can't create a Skylark environment without a tools repository!
+                    builder.setToolsRepository(TestConstants.TOOLS_REPOSITORY);
+                    // Can't create a runtime without a configuration collection factory!
+                    builder.setConfigurationCollectionFactory(
+                        Mockito.mock(ConfigurationCollectionFactory.class));
+                    // Can't create a defaults package without the base options in there!
+                    builder.addConfigurationOptions(BuildConfiguration.Options.class);
+                  }
+                })
+            .build();
+    snooze = new WaitForCompletionCommand(isTestShuttingDown);
+    dispatcher = new BlazeCommandDispatcher(runtime, snooze);
+    BlazeDirectories blazeDirectories =
+        new BlazeDirectories(serverDirectories, scratch.dir("workspace"), productName);
+    runtime.initWorkspace(blazeDirectories, /*bintools=*/ null);
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    isTestShuttingDown.set(true);
+    executor.shutdownNow();
+    executor.awaitTermination(TestUtils.WAIT_TIMEOUT_MILLISECONDS, TimeUnit.MILLISECONDS);
+  }
+
+  // These tests are basically testing the functionality of the dummy command.
+  @Test
+  public void sendingExitCodeToTestCommandResultsInExitWithThatStatus() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    command.requestExitWith(ExitCode.SUCCESS);
+    command.assertFinishedWith(ExitCode.SUCCESS);
+  }
+
+  @Test
+  public void interruptingTestCommandMakesItExitWithInterruptedStatus() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    command.interrupt();
+    command.assertFinishedWith(ExitCode.INTERRUPTED);
+  }
+
+  @Test
+  public void commandIgnoresFirstInterruptionWhenExpectingInterruption() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ true);
+    command.interrupt();
+    command.assertNotFinishedYet();
+    command.requestExitWith(ExitCode.SUCCESS);
+    command.assertFinishedWith(ExitCode.SUCCESS);
+  }
+
+  @Test
+  public void commandExitsWithInterruptedAfterInterruptionCountExceeded() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ true);
+    command.interrupt();
+    command.assertNotFinishedYet();
+    command.interrupt();
+    command.assertFinishedWith(ExitCode.INTERRUPTED);
+  }
+
+  // These tests get into the meat of actual abrupt exits.
+  @Test
+  public void exitForbidsNullException() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    try {
+      command.getModuleEnvironment().exit(null);
+      throw new AssertionError("It shouldn't be allowed to pass null to exit()!");
+    } catch (NullPointerException expected) {
+      // Good!
+    }
+    command.assertNotFinishedYet();
+    command.requestExitWith(ExitCode.SUCCESS);
+  }
+
+  @Test
+  public void exitForbidsNullExitCode() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    try {
+      command.getModuleEnvironment().exit(new AbruptExitException(null));
+      throw new AssertionError(
+          "It shouldn't be allowed to pass an AbruptExitException with null ExitCode to exit()!");
+    } catch (NullPointerException expected) {
+      // Good!
+    }
+    command.assertNotFinishedYet();
+    command.requestExitWith(ExitCode.SUCCESS);
+  }
+
+  @Test
+  public void callingExitOnceInterruptsAndOverridesExitCode() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    command.getModuleEnvironment().exit(new AbruptExitException(ExitCode.NO_TESTS_FOUND));
+    command.assertFinishedWith(ExitCode.NO_TESTS_FOUND);
+  }
+
+  @Test
+  public void callingExitSecondTimeNeitherInterruptsNorReOverridesExitCode() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ true);
+    command.getModuleEnvironment().exit(new AbruptExitException(ExitCode.NO_TESTS_FOUND));
+    command.assertNotFinishedYet();
+    command.getModuleEnvironment().exit(new AbruptExitException(ExitCode.ANALYSIS_FAILURE));
+    command.assertNotFinishedYet();
+    command.requestExitWith(ExitCode.SUCCESS);
+    command.assertFinishedWith(ExitCode.NO_TESTS_FOUND);
+  }
+
+  @Test
+  public void abruptExitCodesDontOverrideInfrastructureFailures() throws Exception {
+    CommandState command = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ true);
+    command.getModuleEnvironment().exit(new AbruptExitException(ExitCode.NO_TESTS_FOUND));
+    command.assertNotFinishedYet();
+    command.requestExitWith(ExitCode.BLAZE_INTERNAL_ERROR);
+    command.assertFinishedWith(ExitCode.BLAZE_INTERNAL_ERROR);
+  }
+
+  @Test
+  public void callingExitAfterCommandCompletesDoesNothing() throws Exception {
+    CommandState firstCommand = snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    firstCommand.requestExitWith(ExitCode.SUCCESS);
+    firstCommand.assertFinishedWith(ExitCode.SUCCESS);
+    CommandState newCommandOnSameThread =
+        snooze.runIn(executor, dispatcher, /*expectInterruption=*/ false);
+    firstCommand.assertOnSameThreadAs(newCommandOnSameThread);
+    firstCommand.getModuleEnvironment().exit(new AbruptExitException(ExitCode.RUN_FAILURE));
+    newCommandOnSameThread.assertNotFinishedYet();
+    newCommandOnSameThread.requestExitWith(ExitCode.SUCCESS);
+  }
+}