getFromFuture will cancel the future by default on InterruptedException

Fixes #11339.

Closes #12453.

PiperOrigin-RevId: 341993215

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

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
index 7f3f83d..93ec205 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
@@ -34,6 +34,7 @@
 import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
+import com.google.devtools.build.lib.remote.util.Utils;
 import com.google.devtools.build.lib.sandbox.SandboxHelpers;
 import com.google.devtools.build.lib.vfs.Path;
 import io.grpc.Context;
@@ -42,7 +43,6 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
 import javax.annotation.concurrent.GuardedBy;
 
 /**
@@ -145,14 +145,7 @@
 
   void downloadFile(Path path, FileArtifactValue metadata)
       throws IOException, InterruptedException {
-    try {
-      downloadFileAsync(path, metadata).get();
-    } catch (ExecutionException e) {
-      if (e.getCause() instanceof IOException) {
-        throw (IOException) e.getCause();
-      }
-      throw new IOException(e.getCause());
-    }
+    Utils.getFromFuture(downloadFileAsync(path, metadata));
   }
 
   private ListenableFuture<Void> downloadFileAsync(Path path, FileArtifactValue metadata)

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

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
index 511c5bd..5fc3b3a 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
@@ -218,7 +218,7 @@
       try {
         if (interruptedException == null) {
           // Wait for all transfers to finish.
-          getFromFuture(transfer);
+          getFromFuture(transfer, cancelRemainingOnInterrupt);
         } else {
           transfer.cancel(true);
         }

src/main/java/com/google/devtools/build/lib/remote/util/Utils.java

diff --git a/src/main/java/com/google/devtools/build/lib/remote/util/Utils.java b/src/main/java/com/google/devtools/build/lib/remote/util/Utils.java
index 9d07d7a..86f21a5 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/util/Utils.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/util/Utils.java
@@ -57,9 +57,22 @@
   /**
    * Returns the result of a {@link ListenableFuture} if successful, or throws any checked {@link
    * Exception} directly if it's an {@link IOException} or else wraps it in an {@link IOException}.
+   *
+   * <p>Cancel the future on {@link InterruptedException}
    */
   public static <T> T getFromFuture(ListenableFuture<T> f)
       throws IOException, InterruptedException {
+    return getFromFuture(f, /* cancelOnInterrupt */ true);
+  }
+
+  /**
+   * Returns the result of a {@link ListenableFuture} if successful, or throws any checked {@link
+   * Exception} directly if it's an {@link IOException} or else wraps it in an {@link IOException}.
+   *
+   * @param cancelOnInterrupt cancel the future on {@link InterruptedException} if {@code true}.
+   */
+  public static <T> T getFromFuture(ListenableFuture<T> f, boolean cancelOnInterrupt)
+      throws IOException, InterruptedException {
     try {
       return f.get();
     } catch (ExecutionException e) {
@@ -74,6 +87,11 @@
         throw (RuntimeException) cause;
       }
       throw new IOException(cause);
+    } catch (InterruptedException e) {
+      if (cancelOnInterrupt) {
+        f.cancel(true);
+      }
+      throw e;
     }
   }
 

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

diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java
index 5f991de..4dded52 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java
@@ -15,6 +15,9 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import build.bazel.remote.execution.v2.Digest;
 import build.bazel.remote.execution.v2.RequestMetadata;
@@ -22,6 +25,8 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.hash.HashCode;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.SettableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
@@ -46,6 +51,8 @@
 import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -186,6 +193,54 @@
     assertThat(a1.getPath().isWritable()).isTrue();
   }
 
+  @Test
+  public void testDownloadFile_onInterrupt_deletePartialDownloadedFile() throws Exception {
+    Semaphore startSemaphore = new Semaphore(0);
+    Semaphore endSemaphore = new Semaphore(0);
+    Map<ActionInput, FileArtifactValue> metadata = new HashMap<>();
+    Map<Digest, ByteString> cacheEntries = new HashMap<>();
+    Artifact a1 = createRemoteArtifact("file1", "hello world", metadata, cacheEntries);
+    RemoteCache remoteCache = mock(RemoteCache.class);
+    when(remoteCache.downloadFile(any(), any()))
+        .thenAnswer(
+            invocation -> {
+              Path path = invocation.getArgument(0);
+              Digest digest = invocation.getArgument(1);
+              ByteString content = cacheEntries.get(digest);
+              if (content == null) {
+                return Futures.immediateFailedFuture(new IOException("Not found"));
+              }
+              content.writeTo(path.getOutputStream());
+
+              startSemaphore.release();
+              return SettableFuture
+                  .create(); // A future that never complete so we can interrupt later
+            });
+    RemoteActionInputFetcher actionInputFetcher =
+        new RemoteActionInputFetcher(remoteCache, execRoot, RequestMetadata.getDefaultInstance());
+
+    AtomicBoolean interrupted = new AtomicBoolean(false);
+    Thread t =
+        new Thread(
+            () -> {
+              try {
+                actionInputFetcher.downloadFile(a1.getPath(), metadata.get(a1));
+              } catch (IOException ignored) {
+                interrupted.set(false);
+              } catch (InterruptedException e) {
+                interrupted.set(true);
+              }
+              endSemaphore.release();
+            });
+    t.start();
+    startSemaphore.acquire();
+    t.interrupt();
+    endSemaphore.acquire();
+
+    assertThat(interrupted.get()).isTrue();
+    assertThat(a1.getPath().exists()).isFalse();
+  }
+
   private Artifact createRemoteArtifact(
       String pathFragment,
       String contents,