Make it possible for a task submitted to the BlockWaitingModule to th…

…row an AbruptExitException.

BlockWaitingModule#submit is meant to be a replacement for BlazeModule#afterCommand when one needs to ensure that the code runs after every other module's afterCommand. Therefore, it should accept tasks with the same signature.

PiperOrigin-RevId: 475816559
Change-Id: Id19526ee9c8f274e1745571bdeb447a435f19bda

src/main/java/com/google/devtools/build/lib/runtime/BlockWaitingModule.java

@@ -19,28 +19,62 @@
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.devtools.build.lib.concurrent.ExecutorUtil;
 import com.google.devtools.build.lib.util.AbruptExitException;
+import java.util.ArrayList;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import javax.annotation.Nullable;
 
 /** A {@link BlazeModule} that waits for submitted tasks to terminate after every command. */
 public class BlockWaitingModule extends BlazeModule {
+
+  /** A task to be submitted. */
+  public interface Task {
+    void call() throws AbruptExitException;
+  }
+
+  /**
+   * Wraps an AbruptExitException thrown by a task.
+   *
+   * <p>This is needed because a task that can throw a checked exception cannot be submitted to
+   * {@link ExecutorService}.
+   */
+  private static class TaskException extends RuntimeException {
+    TaskException(AbruptExitException cause) {
+      super(cause);
+    }
+  }
+
   @Nullable private ExecutorService executorService;
+  @Nullable private ArrayList<Future<?>> submittedTasks;
 
   @Override
   public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
     checkState(executorService == null, "executorService must be null");
+    checkState(submittedTasks == null, "submittedTasks must be null");
 
     executorService =
         Executors.newCachedThreadPool(
             new ThreadFactoryBuilder().setNameFormat("block-waiting-%d").build());
+
+    submittedTasks = new ArrayList<>();
   }
 
   @SuppressWarnings("FutureReturnValueIgnored")
-  public void submit(Runnable task) {
+  public void submit(Task task) {
     checkNotNull(executorService, "executorService must not be null");
+    checkNotNull(submittedTasks, "submittedTasks must be null");
 
-    executorService.submit(task);
+    submittedTasks.add(
+        executorService.submit(
+            () -> {
+              try {
+                task.call();
+              } catch (AbruptExitException e) {
+                throw new TaskException(e);
+              }
+            }));
   }
 
   @Override
@@ -51,6 +85,22 @@ public void afterCommand() throws AbruptExitException {
       Thread.currentThread().interrupt();
     }
 
+    for (Future<?> f : submittedTasks) {
+      try {
+        f.get(); // guaranteed to have completed.
+      } catch (InterruptedException e) {
+        throw new AssertionError("task should not have been interrupted");
+      } catch (ExecutionException e) {
+        Throwable cause = e.getCause();
+        if (cause instanceof TaskException) {
+          checkState(cause.getCause() instanceof AbruptExitException);
+          throw (AbruptExitException) cause.getCause();
+        }
+        throw new RuntimeException(e);
+      }
+    }
+
     executorService = null;
+    submittedTasks = null;
   }
 }

src/test/java/com/google/devtools/build/lib/runtime/BlockWaitingModuleTest.java

@@ -0,0 +1,128 @@
+// Copyright 2022 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 org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import com.google.devtools.build.lib.runtime.BlockWaitingModule.Task;
+import com.google.devtools.build.lib.server.FailureDetails.Crash;
+import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
+import com.google.devtools.build.lib.util.AbruptExitException;
+import com.google.devtools.build.lib.util.DetailedExitCode;
+import java.util.concurrent.ExecutionException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+
+/** Tests for {@link BlockWaitingModule}. */
+@RunWith(JUnit4.class)
+public final class BlockWaitingModuleTest {
+
+  private static final DetailedExitCode CRASH =
+      DetailedExitCode.of(
+          FailureDetail.newBuilder()
+              .setMessage("crash")
+              .setCrash(Crash.newBuilder().setCode(Crash.Code.CRASH_UNKNOWN))
+              .build());
+
+  BlockWaitingModule module;
+  @Mock CommandEnvironment env;
+
+  @Test
+  public void testSubmitZeroTasks() throws Exception {
+    // arrange
+    BlockWaitingModule m = new BlockWaitingModule();
+
+    // act
+    m.beforeCommand(env);
+    m.afterCommand();
+
+    // nothing to assert
+  }
+
+  @Test
+  public void testSubmitOneTask() throws Exception {
+    // arrange
+    BlockWaitingModule m = new BlockWaitingModule();
+    Task t = mock(Task.class);
+
+    // act
+    m.beforeCommand(env);
+    m.submit(t);
+    m.afterCommand();
+
+    // assert
+    verify(t).call();
+  }
+
+  @Test
+  public void testSubmitMultipleTasks() throws Exception {
+    // arrange
+    BlockWaitingModule m = new BlockWaitingModule();
+    Task t1 = mock(Task.class);
+    Task t2 = mock(Task.class);
+    Task t3 = mock(Task.class);
+
+    // act
+    m.beforeCommand(env);
+    m.submit(t1);
+    m.submit(t2);
+    m.submit(t3);
+    m.afterCommand();
+
+    // assert
+    verify(t1).call();
+    verify(t2).call();
+    verify(t3).call();
+  }
+
+  @Test
+  public void testTaskThrowsAbruptExitException() throws Exception {
+    // arrange
+    BlockWaitingModule m = new BlockWaitingModule();
+    Task t = mock(Task.class);
+    doThrow(new AbruptExitException(CRASH)).when(t).call();
+
+    // act
+    m.beforeCommand(env);
+    m.submit(t);
+
+    // assert
+    Throwable e = assertThrows(AbruptExitException.class, m::afterCommand);
+    assertThat(((AbruptExitException) e).getDetailedExitCode()).isEqualTo(CRASH);
+  }
+
+  @Test
+  public void testTaskThrowsUnrecognizedException() throws Exception {
+    // arrange
+    BlockWaitingModule m = new BlockWaitingModule();
+    Task t = mock(Task.class);
+    doThrow(new IllegalStateException("illegal state")).when(t).call();
+
+    // act
+    m.beforeCommand(env);
+    m.submit(t);
+
+    // assert
+    Throwable e = assertThrows(RuntimeException.class, m::afterCommand);
+    assertThat(e).hasCauseThat().isInstanceOf(ExecutionException.class);
+    assertThat(e).hasCauseThat().hasCauseThat().isInstanceOf(IllegalStateException.class);
+    assertThat(e).hasCauseThat().hasCauseThat().hasMessageThat().contains("illegal state");
+  }
+}