Make the module environment invoked exit code path thread-safe.

The module environment can be called from any number of threads, not
necessarily from the main thread. I don't know if it's a problem right now -
we don't have any problem reports that could be caused by this - but better be
safe than sorry.

--
MOS_MIGRATED_REVID=103277567

1 files changed, 9 insertions, 7 deletions

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 5d402a1563..381225c0a0 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
@@ -43,6 +43,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Encapsulates the state needed for a single command. The environment is dropped when the current
@@ -55,7 +56,7 @@ public final class CommandEnvironment {
   private final BlazeModule.ModuleEnvironment blazeModuleEnvironment;
   private final Map<String, String> clientEnv = new HashMap<>();
 
-  private AbruptExitException pendingException;
+  private AtomicReference<AbruptExitException> pendingException = new AtomicReference<>();
 
   private class BlazeModuleEnvironment implements BlazeModule.ModuleEnvironment {
     @Override
@@ -69,8 +70,7 @@ public final class CommandEnvironment {
 
     @Override
     public void exit(AbruptExitException exception) {
-      Preconditions.checkState(pendingException == null);
-      pendingException = exception;
+      pendingException.compareAndSet(null, exception);
     }
   }
 
@@ -175,8 +175,9 @@ 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;
-    if (!originalExit.isInfrastructureFailure() && pendingException != null) {
-      exitCode = pendingException.getExitCode();
+    AbruptExitException exception = pendingException.get();
+    if (!originalExit.isInfrastructureFailure() && exception != null) {
+      exitCode = exception.getExitCode();
     }
     return exitCode;
   }
@@ -189,8 +190,9 @@ public final class CommandEnvironment {
    * the exception this way.
    */
   public void throwPendingException() throws AbruptExitException {
-    if (pendingException != null) {
-      throw pendingException;
+    AbruptExitException exception = pendingException.get();
+    if (exception != null) {
+      throw exception;
     }
   }
 }