Assign priority to input prefech tasks.

High priority tasks are tasks that are critical to the build wall time e.g. staging inputs to local actions. Low priority tasks are tasks that don't have impact on wall time e.g. staging outputs of toplevel targets/aspects.

By assiging priority to these tasks, we can reduce the wall time when the network traffic is high by prioritizing high priority tasks.

PiperOrigin-RevId: 475296193
Change-Id: I961a9e55fb7b563a6813ceef7e58f0cb275a687b

Author: coeuvre

src/main/java/com/google/devtools/build/lib/exec/BUILD

@@ -282,6 +282,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/events",
+        "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/util/io",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",

src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java

@@ -30,6 +30,9 @@
 import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.actions.cache.MetadataInjector;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Path;
@@ -161,7 +164,8 @@ interface SpawnExecutionContext {
     default void prefetchInputsAndWait()
         throws IOException, InterruptedException, ForbiddenActionInputException {
       ListenableFuture<Void> future = prefetchInputs();
-      try {
+      try (SilentCloseable s =
+          Profiler.instance().profile(ProfilerTask.REMOTE_DOWNLOAD, "stage remote inputs")) {
         future.get();
       } catch (ExecutionException e) {
         Throwable cause = e.getCause();

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

@@ -33,9 +33,6 @@
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
-import com.google.devtools.build.lib.profiler.Profiler;
-import com.google.devtools.build.lib.profiler.ProfilerTask;
-import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.remote.util.AsyncTaskCache;
 import com.google.devtools.build.lib.remote.util.RxUtils.TransferResult;
 import com.google.devtools.build.lib.remote.util.TempPathGenerator;
@@ -64,6 +61,20 @@ public abstract class AbstractActionInputPrefetcher implements ActionInputPrefet
 
   protected final Path execRoot;
 
+  /** Priority for the staging task. */
+  protected enum Priority {
+    /**
+     * High priority tasks are tasks that are critical to the wall time e.g. staging inputs to local
+     * actions.
+     */
+    HIGH,
+    /**
+     * Low priority tasks are tasks that don't have impact on the wall time e.g. staging outputs of
+     * toplevel targets/aspects.
+     */
+    LOW,
+  }
+
   protected AbstractActionInputPrefetcher(Path execRoot, TempPathGenerator tempPathGenerator) {
     this.execRoot = execRoot;
     this.tempPathGenerator = tempPathGenerator;
@@ -77,7 +88,7 @@ protected AbstractActionInputPrefetcher(Path execRoot, TempPathGenerator tempPat
    * @param tempPath the temporary path which the input should be written to.
    */
   protected abstract ListenableFuture<Void> doDownloadFile(
-      Path tempPath, FileArtifactValue metadata) throws IOException;
+      Path tempPath, FileArtifactValue metadata, Priority priority) throws IOException;
 
   protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOException {}
 
@@ -98,6 +109,13 @@ protected Completable onErrorResumeNext(Throwable error) {
   @Override
   public ListenableFuture<Void> prefetchFiles(
       Iterable<? extends ActionInput> inputs, MetadataProvider metadataProvider) {
+    return prefetchFiles(inputs, metadataProvider, Priority.HIGH);
+  }
+
+  protected ListenableFuture<Void> prefetchFiles(
+      Iterable<? extends ActionInput> inputs,
+      MetadataProvider metadataProvider,
+      Priority priority) {
     Map<SpecialArtifact, List<TreeFileArtifact>> trees = new HashMap<>();
     List<ActionInput> files = new ArrayList<>();
     for (ActionInput input : inputs) {
@@ -120,22 +138,22 @@ public ListenableFuture<Void> prefetchFiles(
             .flatMapSingle(
                 entry ->
                     toTransferResult(
-                        prefetchInputTree(metadataProvider, entry.getKey(), entry.getValue())));
+                        prefetchInputTree(
+                            metadataProvider, entry.getKey(), entry.getValue(), priority)));
     Flowable<TransferResult> fileDownloads =
         Flowable.fromIterable(files)
-            .flatMapSingle(input -> toTransferResult(prefetchInputFile(metadataProvider, input)));
+            .flatMapSingle(
+                input -> toTransferResult(prefetchInputFile(metadataProvider, input, priority)));
     Flowable<TransferResult> transfers = Flowable.merge(treeDownloads, fileDownloads);
     Completable prefetch = mergeBulkTransfer(transfers).onErrorResumeNext(this::onErrorResumeNext);
-    Completable prefetchWithProfiler =
-        Completable.using(
-            () -> Profiler.instance().profile(ProfilerTask.REMOTE_DOWNLOAD, "stage remote inputs"),
-            profiler -> prefetch,
-            SilentCloseable::close);
-    return toListenableFuture(prefetchWithProfiler);
+    return toListenableFuture(prefetch);
   }
 
   private Completable prefetchInputTree(
-      MetadataProvider provider, SpecialArtifact tree, List<TreeFileArtifact> treeFiles) {
+      MetadataProvider provider,
+      SpecialArtifact tree,
+      List<TreeFileArtifact> treeFiles,
+      Priority priority) {
     Path treeRoot = execRoot.getRelative(tree.getExecPath());
     HashMap<TreeFileArtifact, Path> treeFileTmpPathMap = new HashMap<>();
 
@@ -153,7 +171,8 @@ private Completable prefetchInputTree(
                   treeFileTmpPathMap.put(treeFile, tempPath);
 
                   return toTransferResult(
-                      toCompletable(() -> doDownloadFile(tempPath, metadata), directExecutor()));
+                      toCompletable(
+                          () -> doDownloadFile(tempPath, metadata, priority), directExecutor()));
                 });
 
     AtomicBoolean completed = new AtomicBoolean();
@@ -211,8 +230,8 @@ private Completable prefetchInputTree(
     return downloadCache.executeIfNot(treeRoot, download);
   }
 
-  private Completable prefetchInputFile(MetadataProvider metadataProvider, ActionInput input)
-      throws IOException {
+  private Completable prefetchInputFile(
+      MetadataProvider metadataProvider, ActionInput input, Priority priority) throws IOException {
     if (input instanceof VirtualActionInput) {
       prefetchVirtualActionInput((VirtualActionInput) input);
       return Completable.complete();
@@ -224,7 +243,7 @@ private Completable prefetchInputFile(MetadataProvider metadataProvider, ActionI
     }
 
     Path path = execRoot.getRelative(input.getExecPath());
-    return downloadFileRx(path, metadata);
+    return downloadFileRx(path, metadata, priority);
   }
 
   /**
@@ -233,7 +252,7 @@ private Completable prefetchInputFile(MetadataProvider metadataProvider, ActionI
    * <p>The file will be written into a temporary file and moved to the final destination after the
    * download finished.
    */
-  public Completable downloadFileRx(Path path, FileArtifactValue metadata) {
+  public Completable downloadFileRx(Path path, FileArtifactValue metadata, Priority priority) {
     if (!shouldDownloadFile(path, metadata)) {
       return Completable.complete();
     }
@@ -253,7 +272,7 @@ public Completable downloadFileRx(Path path, FileArtifactValue metadata) {
         Completable.using(
             tempPathGenerator::generateTempPath,
             tempPath ->
-                toCompletable(() -> doDownloadFile(tempPath, metadata), directExecutor())
+                toCompletable(() -> doDownloadFile(tempPath, metadata, priority), directExecutor())
                     .doOnComplete(
                         () -> {
                           finalizeDownload(tempPath, finalPath);
@@ -277,7 +296,7 @@ public Completable downloadFileRx(Path path, FileArtifactValue metadata) {
    * download finished.
    */
   public ListenableFuture<Void> downloadFileAsync(Path path, FileArtifactValue metadata) {
-    return toListenableFuture(downloadFileRx(path, metadata));
+    return toListenableFuture(downloadFileRx(path, metadata, Priority.LOW));
   }
 
   /**

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

@@ -71,8 +71,8 @@ protected boolean shouldDownloadFile(Path path, FileArtifactValue metadata) {
   }
 
   @Override
-  protected ListenableFuture<Void> doDownloadFile(Path tempPath, FileArtifactValue metadata)
-      throws IOException {
+  protected ListenableFuture<Void> doDownloadFile(
+      Path tempPath, FileArtifactValue metadata, Priority priority) throws IOException {
     checkArgument(metadata.isRemote(), "Cannot download file that is not a remote file.");
     RequestMetadata requestMetadata =
         TracingMetadataUtils.buildMetadata(buildRequestId, commandId, metadata.getActionId(), null);

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

@@ -386,7 +386,7 @@ protected static void mockDownload(
               return resultSupplier.get();
             })
         .when(prefetcher)
-        .doDownloadFile(any(), any());
+        .doDownloadFile(any(), any(), any());
   }
 
   private void assertReadableNonWritableAndExecutable(Path path) throws IOException {