Widen interface for action input prefetching.

RELNOTES: None PiperOrigin-RevId: 703221309 Change-Id: Ie7a4607b530a01e49f671bc6804a307740c81175
uber-common · Dec 5, 2024 · b024eea · b024eea
1 parent 5fc2b01
commit b024eea
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 49 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java
@@ -20,21 +20,10 @@
 
 /** Prefetches files to local disk. */
 public interface ActionInputPrefetcher {
-  /**
-   * Returns the metadata for an {@link ActionInput}.
-   *
-   * <p>This will generally call through to a {@link InputMetadataProvider} and ask for the metadata
-   * of either an input or an output artifact.
-   */
-  public interface MetadataSupplier {
-    FileArtifactValue getMetadata(ActionInput actionInput) throws IOException, InterruptedException;
-  }
-
   public static final ActionInputPrefetcher NONE =
-      (action, inputs, metadataSupplier, priority) -> {
-        // Do nothing.
-        return immediateVoidFuture();
-      };
+      (action, inputs, metadataProvider, priority) ->
+          // Do nothing.
+          immediateVoidFuture();
 
   /** Priority for the staging task. */
   public enum Priority {
@@ -70,6 +59,6 @@ public enum Priority {
   ListenableFuture<Void> prefetchFiles(
       ActionExecutionMetadata action,
       Iterable<? extends ActionInput> inputs,
-      MetadataSupplier metadataSupplier,
+      InputMetadataProvider metadataProvider,
       Priority priority);
 }
diff --git a/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java
@@ -289,7 +289,7 @@ public ListenableFuture<Void> prefetchInputs() throws ForbiddenActionInputExcept
                     spawn.getResourceOwner(),
                     getInputMapping(PathFragment.EMPTY_FRAGMENT, /* willAccessRepeatedly= */ true)
                         .values(),
-                    getInputMetadataProvider()::getInputMetadata,
+                    getInputMetadataProvider(),
                     Priority.MEDIUM),
             BulkTransferException.class,
             (BulkTransferException e) -> {

diff --git a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java
@@ -40,6 +40,7 @@
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
+import com.google.devtools.build.lib.actions.InputMetadataProvider;
 import com.google.devtools.build.lib.actions.cache.OutputMetadataStore;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.events.Reporter;
@@ -89,6 +90,17 @@ enum DirectoryState {
     OUTPUT_PERMISSIONS
   }
 
+  /**
+   * Returns the metadata for an {@link ActionInput}.
+   *
+   * <p>This will generally call through to a {@link InputMetadataProvider} and ask for the metadata
+   * of either an input or an output artifact.
+   */
+  @VisibleForTesting
+  public interface MetadataSupplier {
+    FileArtifactValue getMetadata(ActionInput actionInput) throws IOException, InterruptedException;
+  }
+
   /**
    * Tracks directory permissions to minimize filesystem operations.
    *
@@ -260,6 +272,30 @@ protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOExc
    */
   @Override
   public ListenableFuture<Void> prefetchFiles(
+      ActionExecutionMetadata action,
+      Iterable<? extends ActionInput> inputs,
+      InputMetadataProvider metadataProvider,
+      Priority priority) {
+    return prefetchFilesInterruptibly(action, inputs, metadataProvider::getInputMetadata, priority);
+  }
+
+  /**
+   * Fetches remotely stored action outputs, that are inputs to this spawn, and stores them under
+   * their path in the output base.
+   *
+   * <p>The {@code inputs} may not contain any unexpanded directories.
+   *
+   * <p>This method is safe to be called concurrently from spawn runners before running any local
+   * spawn.
+   *
+   * <p>This method is similar to #prefetchFiles() above, but note that {@code metadataSupplier} may
+   * throw {@link InterruptedException}. If it does, this method will propagate this exception in
+   * the returned future.
+   *
+   * @return a future that is completed once all downloads have finished.
+   */
+  @VisibleForTesting
+  public ListenableFuture<Void> prefetchFilesInterruptibly(
       ActionExecutionMetadata action,
       Iterable<? extends ActionInput> inputs,
       MetadataSupplier metadataSupplier,
@@ -686,7 +722,7 @@ public void finalizeAction(Action action, OutputMetadataStore outputMetadataStor
     if (!outputsToDownload.isEmpty()) {
       try (var s = Profiler.instance().profile(ProfilerTask.REMOTE_DOWNLOAD, "Download outputs")) {
         getFromFuture(
-            prefetchFiles(
+            prefetchFilesInterruptibly(
                 action, outputsToDownload, outputMetadataStore::getOutputMetadata, Priority.HIGH));
       }
     }

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java
@@ -735,8 +735,7 @@ ActionInput getInput(String execPath) {
   @Nullable
   @VisibleForTesting
   FileArtifactValue getInputMetadata(ActionInput input) {
-    PathFragment execPath = input.getExecPath();
-    return inputArtifactData.getMetadata(execPath);
+    return inputArtifactData.getInputMetadata(input);
   }
 
   private void downloadFileIfRemote(PathFragment path) throws IOException {
@@ -753,7 +752,7 @@ private void downloadFileIfRemote(PathFragment path) throws IOException {
       }
       getFromFuture(
           inputFetcher.prefetchFiles(
-              action, ImmutableList.of(input), this::getInputMetadata, Priority.CRITICAL));
+              action, ImmutableList.of(input), inputArtifactData, Priority.CRITICAL));
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new IOException(String.format("Received interrupt while fetching file '%s'", path), e);

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
@@ -468,8 +468,7 @@ private void downloadArtifact(
         var action =
             ActionUtils.getActionForLookupData(env, derivedArtifact.getGeneratingActionKey());
         var future =
-            actionInputPrefetcher.prefetchFiles(
-                action, filesToDownload, inputMap::getInputMetadata, Priority.LOW);
+            actionInputPrefetcher.prefetchFiles(action, filesToDownload, inputMap, Priority.LOW);
         futures.add(future);
       }
     } else {
@@ -485,7 +484,7 @@ private void downloadArtifact(
             ActionUtils.getActionForLookupData(env, derivedArtifact.getGeneratingActionKey());
         var future =
             actionInputPrefetcher.prefetchFiles(
-                action, ImmutableList.of(artifact), inputMap::getInputMetadata, Priority.LOW);
+                action, ImmutableList.of(artifact), inputMap, Priority.LOW);
         futures.add(future);
       }
     }

diff --git a/src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java b/src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java
@@ -45,7 +45,6 @@
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInput;
-import com.google.devtools.build.lib.actions.ActionInputPrefetcher.MetadataSupplier;
 import com.google.devtools.build.lib.actions.ActionInputPrefetcher.Priority;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
@@ -56,6 +55,7 @@
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
+import com.google.devtools.build.lib.remote.AbstractActionInputPrefetcher.MetadataSupplier;
 import com.google.devtools.build.lib.remote.common.LostInputsEvent;
 import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
 import com.google.devtools.build.lib.testing.vfs.SpiedFileSystem;
@@ -296,7 +296,9 @@ public void prefetchFiles_fileExists_doNotDownload()
     FileSystemUtils.writeContent(a.getPath(), "hello world".getBytes(UTF_8));
     AbstractActionInputPrefetcher prefetcher = spy(createPrefetcher(cas));
 
-    wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM));
+    wait(
+        prefetcher.prefetchFilesInterruptibly(
+            action, metadata.keySet(), metadata::get, Priority.MEDIUM));
 
     verify(prefetcher, never()).doDownloadFile(eq(action), any(), any(), any(), any(), any());
     assertThat(prefetcher.downloadedFiles()).containsExactly(a.getPath());
@@ -312,7 +314,9 @@ public void prefetchFiles_fileExistsButContentMismatches_download()
     FileSystemUtils.writeContent(a.getPath(), "hello world local".getBytes(UTF_8));
     AbstractActionInputPrefetcher prefetcher = spy(createPrefetcher(cas));
 
-    wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM));
+    wait(
+        prefetcher.prefetchFilesInterruptibly(
+            action, metadata.keySet(), metadata::get, Priority.MEDIUM));
 
     verify(prefetcher).doDownloadFile(eq(action), any(), any(), eq(a.getExecPath()), any(), any());
     assertThat(prefetcher.downloadedFiles()).containsExactly(a.getPath());
@@ -328,7 +332,9 @@ public void prefetchFiles_downloadRemoteFiles() throws Exception {
     Artifact a2 = createRemoteArtifact("file2", "fizz buzz", metadata, cas);
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
-    wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM));
+    wait(
+        prefetcher.prefetchFilesInterruptibly(
+            action, metadata.keySet(), metadata::get, Priority.MEDIUM));
 
     assertThat(FileSystemUtils.readContent(a1.getPath(), UTF_8)).isEqualTo("hello world");
     assertReadableNonWritableAndExecutable(a1.getPath());
@@ -346,7 +352,9 @@ public void prefetchFiles_downloadRemoteFiles_withMaterializationExecPath() thro
     Artifact a = createRemoteArtifact("file", "hello world", targetExecPath, metadata, cas);
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
-    wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM));
+    wait(
+        prefetcher.prefetchFilesInterruptibly(
+            action, metadata.keySet(), metadata::get, Priority.MEDIUM));
 
     assertThat(a.getPath().isSymbolicLink()).isTrue();
     assertThat(a.getPath().readSymbolicLink())
@@ -376,7 +384,7 @@ public void prefetchFiles_downloadRemoteTrees() throws Exception {
 
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
-    wait(prefetcher.prefetchFiles(action, children, metadata::get, Priority.MEDIUM));
+    wait(prefetcher.prefetchFilesInterruptibly(action, children, metadata::get, Priority.MEDIUM));
 
     assertThat(FileSystemUtils.readContent(firstChild.getPath(), UTF_8)).isEqualTo("content1");
     assertThat(FileSystemUtils.readContent(secondChild.getPath(), UTF_8)).isEqualTo("content2");
@@ -407,7 +415,7 @@ public void prefetchFiles_downloadRemoteTrees_partial() throws Exception {
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
     wait(
-        prefetcher.prefetchFiles(
+        prefetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(firstChild, secondChild), metadata::get, Priority.MEDIUM));
 
     assertThat(firstChild.getPath().exists()).isFalse();
@@ -437,7 +445,7 @@ public void prefetchFiles_downloadRemoteTrees_withMaterializationExecPath() thro
 
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
-    wait(prefetcher.prefetchFiles(action, children, metadata::get, Priority.MEDIUM));
+    wait(prefetcher.prefetchFilesInterruptibly(action, children, metadata::get, Priority.MEDIUM));
 
     assertThat(tree.getPath().isSymbolicLink()).isTrue();
     assertThat(tree.getPath().readSymbolicLink())
@@ -475,13 +483,13 @@ public void prefetchFiles_downloadRemoteTrees_forActionTemplateExpansion() throw
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
     wait(
-        prefetcher.prefetchFiles(
+        prefetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(firstChild), metadata::get, Priority.MEDIUM));
 
     assertTreeReadableWritableAndExecutable(tree.getPath());
 
     wait(
-        prefetcher.prefetchFiles(
+        prefetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(secondChild), metadata::get, Priority.MEDIUM));
 
     assertTreeReadableWritableAndExecutable(tree.getPath());
@@ -497,7 +505,7 @@ public void prefetchFiles_missingFiles_fails() throws Exception {
         Exception.class,
         () ->
             wait(
-                prefetcher.prefetchFiles(
+                prefetcher.prefetchFilesInterruptibly(
                     action, ImmutableList.of(a), metadata::get, Priority.MEDIUM)));
 
     assertThat(prefetcher.downloadedFiles()).isEmpty();
@@ -515,7 +523,9 @@ public void prefetchFiles_ignoreNonRemoteFiles() throws Exception {
     ImmutableMap<ActionInput, FileArtifactValue> metadata = ImmutableMap.of(a, f);
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(new HashMap<>());
 
-    wait(prefetcher.prefetchFiles(action, ImmutableList.of(a), metadata::get, Priority.MEDIUM));
+    wait(
+        prefetcher.prefetchFilesInterruptibly(
+            action, ImmutableList.of(a), metadata::get, Priority.MEDIUM));
 
     assertThat(prefetcher.downloadedFiles()).isEmpty();
     assertThat(prefetcher.downloadsInProgress()).isEmpty();
@@ -541,7 +551,7 @@ public void prefetchFiles_ignoreNonRemoteFiles_tree() throws Exception {
 
     AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);
 
-    wait(prefetcher.prefetchFiles(action, children, metadata::get, Priority.MEDIUM));
+    wait(prefetcher.prefetchFilesInterruptibly(action, children, metadata::get, Priority.MEDIUM));
 
     assertThat(firstChild.getPath().exists()).isFalse();
     assertThat(FileSystemUtils.readContent(secondChild.getPath(), UTF_8)).isEqualTo("content2");
@@ -573,7 +583,7 @@ public void prefetchFiles_treeFiles_minimizeFilesystemOperations() throws Except
     reset(fs);
 
     wait(
-        prefetcher.prefetchFiles(
+        prefetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(firstChild, secondChild), metadata::get, Priority.MEDIUM));
 
     verify(fs).createWritableDirectory(root);
@@ -584,7 +594,7 @@ public void prefetchFiles_treeFiles_minimizeFilesystemOperations() throws Except
     reset(fs);
 
     wait(
-        prefetcher.prefetchFiles(
+        prefetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(firstChild, secondChild), metadata::get, Priority.MEDIUM));
 
     verify(fs, never()).createWritableDirectory(root);
@@ -623,7 +633,7 @@ public void prefetchFiles_treeFiles_multipleThreads_waitForPermissionsToBeSet()
     Callable<Void> prefetch =
         () -> {
           wait(
-              prefetcher.prefetchFiles(
+              prefetcher.prefetchFilesInterruptibly(
                   action, ImmutableList.of(child), metadata::get, Priority.MEDIUM));
           assertTreeReadableNonWritableAndExecutable(tree.getPath());
           return null;
@@ -657,7 +667,7 @@ public void prefetchFiles_multipleThreads_downloadIsCancelled() throws Exception
             () -> {
               try {
                 wait(
-                    prefetcher.prefetchFiles(
+                    prefetcher.prefetchFilesInterruptibly(
                         action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM));
               } catch (IOException | ExecException | InterruptedException ignored) {
                 // do nothing
@@ -669,7 +679,7 @@ public void prefetchFiles_multipleThreads_downloadIsCancelled() throws Exception
             () -> {
               try {
                 wait(
-                    prefetcher.prefetchFiles(
+                    prefetcher.prefetchFilesInterruptibly(
                         action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM));
               } catch (IOException | ExecException | InterruptedException ignored) {
                 // do nothing
@@ -707,7 +717,7 @@ public void prefetchFiles_multipleThreads_downloadIsNotCancelledByOtherThreads()
             () -> {
               try {
                 wait(
-                    prefetcher.prefetchFiles(
+                    prefetcher.prefetchFilesInterruptibly(
                         action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM));
               } catch (IOException | ExecException | InterruptedException ignored) {
                 // do nothing
@@ -720,7 +730,7 @@ public void prefetchFiles_multipleThreads_downloadIsNotCancelledByOtherThreads()
             () -> {
               try {
                 wait(
-                    prefetcher.prefetchFiles(
+                    prefetcher.prefetchFilesInterruptibly(
                         action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM));
                 successful.set(true);
               } catch (IOException | ExecException | InterruptedException ignored) {
@@ -764,7 +774,7 @@ public void prefetchFile_interruptingMetadataSupplier_interruptsDownload() throw
         };
 
     ListenableFuture<Void> future =
-        prefetcher.prefetchFiles(
+        prefetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(a1), interruptedMetadataSupplier, Priority.MEDIUM);
 
     assertThrows(InterruptedException.class, () -> getFromFuture(future));
@@ -792,7 +802,7 @@ public void prefetchFiles_onInterrupt_deletePartialDownloadedFile() throws Excep
             () -> {
               try {
                 getFromFuture(
-                    prefetcher.prefetchFiles(
+                    prefetcher.prefetchFilesInterruptibly(
                         action, ImmutableList.of(a1), metadata::get, Priority.MEDIUM));
               } catch (IOException ignored) {
                 // Intentionally left empty
@@ -831,7 +841,7 @@ public void onLostInputsEvent(LostInputsEvent event) {
         Exception.class,
         () ->
             wait(
-                prefetcher.prefetchFiles(
+                prefetcher.prefetchFilesInterruptibly(
                     action, metadata.keySet(), metadata::get, Priority.MEDIUM)));
 
     assertThat(lostInputsEvents).hasSize(1);

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
@@ -102,7 +102,7 @@ public void testStagingVirtualActionInput() throws Exception {
 
     // act
     wait(
-        actionInputFetcher.prefetchFiles(
+        actionInputFetcher.prefetchFilesInterruptibly(
             action, ImmutableList.of(a), (ActionInput unused) -> null, Priority.MEDIUM));
 
     // assert
@@ -131,7 +131,7 @@ public void testStagingEmptyVirtualActionInput() throws Exception {
 
     // act
     wait(
-        actionInputFetcher.prefetchFiles(
+        actionInputFetcher.prefetchFilesInterruptibly(
             action,
             ImmutableList.of(VirtualActionInput.EMPTY_MARKER),
             (ActionInput unused) -> null,
@@ -153,7 +153,7 @@ public void prefetchFiles_missingFiles_failsWithSpecificMessage() throws Excepti
             BulkTransferException.class,
             () ->
                 wait(
-                    prefetcher.prefetchFiles(
+                    prefetcher.prefetchFilesInterruptibly(
                         action, ImmutableList.of(a), metadata::get, Priority.MEDIUM)));
 
     assertThat(prefetcher.downloadedFiles()).isEmpty();