Strictify exception types in remote package to reduce wrapping

This makes it easier to understand what kinds of exceptions can be thrown by a particular function and avoids wrapping in `RuntimeException` that can turn regular failures into crashes.

In particular, this resolves an issue seen on a PR in which `RemoteSpawnRunner` turned an `ExecException` raised by `RES.uploadInputsIfNotPresent` due to a lost input into a `RuntimeException`, thus resulting in a crash. It's not clear whether this behavior can also happen at HEAD or whether all callers unwrap the `RuntimeException`, so this change doesn't necessarily fix a bug.

Work towards #21378

Closes #26650.

PiperOrigin-RevId: 792329428
Change-Id: Iced921ce8633926e6488d38502c9b93c9fffa393

Author: fmeum

src/main/java/com/google/devtools/build/lib/remote/BazelOutputService.java

@@ -360,7 +360,7 @@ private FinalizeBuildResponse finalizeBuild(FinalizeBuildRequest request)
 
   @Override
   public void finalizeAction(Action action, OutputMetadataStore outputMetadataStore)
-      throws IOException, EnvironmentalExecException, InterruptedException {
+      throws IOException, InterruptedException {
     var execRoot = execRootSupplier.get();
     var outputPath = outputPathSupplier.get();
 

src/main/java/com/google/devtools/build/lib/remote/ReferenceCountedChannel.java

@@ -78,8 +78,6 @@ public ServerCapabilities getServerCapabilities() throws IOException {
     try (var s = Profiler.instance().profile("getServerCapabilities")) {
       return blockingGet(
           withChannelConnection(ChannelConnectionWithServerCapabilities::getServerCapabilities));
-    } catch (ExecutionException e) {
-      throw new IOException(e);
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new IOException(e);
@@ -90,21 +88,30 @@ public boolean isShutdown() {
     return dynamicConnectionPool.isClosed();
   }
 
+  /**
+   * A specialized {@link Function} that can only throw {@link IOException} and {@link
+   * InterruptedException}.
+   */
+  @FunctionalInterface
+  public interface IOFunction<T, R> extends Function<T, R> {
+    @Override
+    R apply(T t) throws IOException, InterruptedException;
+  }
+
   @CheckReturnValue
   public <T> ListenableFuture<T> withChannelFuture(
-      Function<Channel, ? extends ListenableFuture<T>> source) {
+      IOFunction<Channel, ? extends ListenableFuture<T>> source) {
     return RxFutures.toListenableFuture(
         withChannel(channel -> RxFutures.toSingle(() -> source.apply(channel), directExecutor())));
   }
 
-  public <T> T withChannelBlocking(Function<Channel, T> source)
-      throws ExecutionException, IOException, InterruptedException {
+  public <T> T withChannelBlocking(IOFunction<Channel, T> source)
+      throws IOException, InterruptedException {
     return blockingGet(withChannel(channel -> Single.just(source.apply(channel))));
   }
 
   // prevents rxjava silent possible wrap of RuntimeException and misinterpretation
-  private <T> T blockingGet(Single<T> single)
-      throws ExecutionException, IOException, InterruptedException {
+  private <T> T blockingGet(Single<T> single) throws IOException, InterruptedException {
     SettableFuture<T> future = SettableFuture.create();
     single.subscribe(
         new SingleObserver<T>() {
@@ -135,8 +142,9 @@ public void onSubscribe(Disposable d) {
     } catch (ExecutionException e) {
       Throwable cause = e.getCause();
       Throwables.throwIfInstanceOf(cause, IOException.class);
+      Throwables.throwIfInstanceOf(cause, InterruptedException.class);
       Throwables.throwIfUnchecked(cause);
-      throw e;
+      throw new IllegalStateException("Unexpected exception type", cause);
     }
   }
 

src/main/java/com/google/devtools/build/lib/remote/RemoteRetrier.java

@@ -18,15 +18,13 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Throwables;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
 import com.google.devtools.build.lib.remote.Retrier.ResultClassifier.Result;
 import com.google.devtools.build.lib.remote.options.RemoteOptions;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import java.io.IOException;
 import java.time.Duration;
-import java.util.concurrent.Callable;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
@@ -106,33 +104,13 @@ public RemoteRetrier(
     super(backoff, resultClassifier, retryScheduler, circuitBreaker, sleeper);
   }
 
-  /**
-   * Execute a callable with retries. {@link IOException} and {@link InterruptedException} are
-   * propagated directly to the caller. All other exceptions are wrapped in {@link
-   * RuntimeException}.
-   */
+  /** Execute a callable with retries. */
   @Override
-  public <T> T execute(Callable<T> call) throws IOException, InterruptedException {
+  public <T, E extends Exception> T execute(RetryableCallable<T, E> call)
+      throws E, IOException, InterruptedException {
     return execute(call, newBackoff());
   }
 
-  /**
-   * Execute a callable with retries and given {@link Backoff}. {@link IOException} and {@link
-   * InterruptedException} are propagated directly to the caller. All other exceptions are wrapped
-   * in {@link RuntimeException}.
-   */
-  @Override
-  public <T> T execute(Callable<T> call, Backoff backoff) throws IOException, InterruptedException {
-    try {
-      return super.execute(call, backoff);
-    } catch (Exception e) {
-      Throwables.throwIfInstanceOf(e, IOException.class);
-      Throwables.throwIfInstanceOf(e, InterruptedException.class);
-      Throwables.throwIfUnchecked(e);
-      throw new RuntimeException(e);
-    }
-  }
-
   /** Backoff strategy that backs off exponentially. */
   public static class ExponentialBackoff implements Backoff {
 
@@ -154,8 +132,8 @@ public static class ExponentialBackoff implements Backoff {
      *     jitter, and 1 providing a duration that is 0-200% of the non-jittered duration.
      * @param maxAttempts Maximal times to attempt a retry 0 means no retries.
      */
-    ExponentialBackoff(Duration initial, Duration max, double multiplier, double jitter,
-        int maxAttempts) {
+    ExponentialBackoff(
+        Duration initial, Duration max, double multiplier, double jitter, int maxAttempts) {
       Preconditions.checkArgument(multiplier > 1, "multipler must be > 1");
       Preconditions.checkArgument(jitter >= 0 && jitter <= 1, "jitter must be in the range (0, 1)");
       Preconditions.checkArgument(maxAttempts >= 0, "maxAttempts must be >= 0");

src/main/java/com/google/devtools/build/lib/remote/Retrier.java

@@ -224,33 +224,38 @@ public Retrier(
     this.sleeper = sleeper;
   }
 
-  ListeningScheduledExecutorService getRetryService() {
-    return retryService;
+  /** A {@link Callable} that can be retried in case of transient failure. */
+  @FunctionalInterface
+  public interface RetryableCallable<T, E extends Exception> extends Callable<T> {
+    @Override
+    T call() throws IOException, InterruptedException, E;
   }
 
   /**
-   * Execute a {@link Callable}, retrying execution in case of transient failure and returning the
-   * result in case of success.
+   * Execute a {@link RetryableCallable}, retrying execution in case of transient failure and
+   * returning the result in case of success.
    */
-  public <T> T execute(Callable<T> call) throws Exception {
+  public <T, E extends Exception> T execute(RetryableCallable<T, E> call)
+      throws E, IOException, InterruptedException {
     return execute(call, newBackoff());
   }
 
   /**
-   * Execute a {@link Callable}, retrying execution in case of transient failure and returning the
-   * result in case of success with give {@link Backoff}.
+   * Execute a {@link RetryableCallable}, retrying execution in case of transient failure and
+   * returning the result in case of success with give {@link Backoff}.
    *
    * <p>{@link InterruptedException} is not retried.
    *
    * @param call the {@link Callable} to execute.
-   * @throws Exception if the {@code call} didn't succeed within the framework specified by {@code
-   *     backoffSupplier} and {@code resultClassifier}.
+   * @throws E or {@link IOException} if the {@code call} didn't succeed within the framework
+   *     specified by {@code backoffSupplier} and {@code resultClassifier}.
    * @throws CircuitBreakerException in case a call was rejected because the circuit breaker
    *     tripped.
    * @throws InterruptedException if the {@code call} throws an {@link InterruptedException} or the
    *     current thread's interrupted flag is set.
    */
-  public <T> T execute(Callable<T> call, Backoff backoff) throws Exception {
+  public <T, E extends Exception> T execute(RetryableCallable<T, E> call, Backoff backoff)
+      throws E, IOException, InterruptedException {
     while (true) {
       State circuitState = circuitBreaker.state();
       if (State.REJECT_CALLS.equals(circuitState)) {

src/main/java/com/google/devtools/build/lib/remote/downloader/GrpcRemoteDownloader.java

@@ -183,17 +183,18 @@ public void download(
       }
       final Digest blobDigest = response.getBlobDigest();
 
-      retrier.execute(
-          () -> {
-            try (OutputStream out = newOutputStream(destination, checksum)) {
-              Utils.getFromFuture(
-                  cacheClient.downloadBlob(remoteActionExecutionContext, blobDigest, out));
-            } catch (OutputDigestMismatchException e) {
-              e.setOutputPath(destination.getPathString());
-              throw e;
-            }
-            return null;
-          });
+      var unused =
+          retrier.execute(
+              () -> {
+                try (OutputStream out = newOutputStream(destination, checksum)) {
+                  Utils.getFromFuture(
+                      cacheClient.downloadBlob(remoteActionExecutionContext, blobDigest, out));
+                } catch (OutputDigestMismatchException e) {
+                  e.setOutputPath(destination.getPathString());
+                  throw e;
+                }
+                return null;
+              });
 
     } catch (StatusRuntimeException | IOException e) {
       eventHandler.post(new FetchEvent(eventUri, FetchId.Downloader.GRPC, /* success= */ false));

src/test/java/com/google/devtools/build/lib/remote/RetrierTest.java

@@ -158,19 +158,20 @@ public void nestedRetriesShouldWork() throws Exception {
     AtomicInteger attemptsLvl1 = new AtomicInteger();
     AtomicInteger attemptsLvl2 = new AtomicInteger();
     try {
-      r.execute(
-          () -> {
-            attemptsLvl0.incrementAndGet();
-            return r.execute(
-                () -> {
-                  attemptsLvl1.incrementAndGet();
-                  return r.execute(
-                      () -> {
-                        attemptsLvl2.incrementAndGet();
-                        throw new Exception("failure message");
-                      });
-                });
-          });
+      var unused =
+          r.execute(
+              () -> {
+                attemptsLvl0.incrementAndGet();
+                return r.execute(
+                    () -> {
+                      attemptsLvl1.incrementAndGet();
+                      return r.execute(
+                          () -> {
+                            attemptsLvl2.incrementAndGet();
+                            throw new Exception("failure message");
+                          });
+                    });
+              });
     } catch (Exception e) {
       assertThat(e).hasMessageThat().isEqualTo("failure message");
       assertThat(attemptsLvl0.get()).isEqualTo(2);
@@ -229,10 +230,11 @@ public void circuitBreakerHalfOpenIsNotRetried() throws Exception {
     cb.trialCall();
 
     try {
-      r.execute(
-          () -> {
-            throw new Exception("call failed");
-          });
+      var unused =
+          r.execute(
+              () -> {
+                throw new Exception("call failed");
+              });
     } catch (Exception expected) {
       // Intentionally left empty.
     }