Add ActionExecutionMetadata as a parameter to ActionInputPrefetcher#prefetchFiles. (#18656)

Due to recent overhual of `prefetchFiles`, the call sites now always have the reference to `ActionExecutionMetadata`, which makes it possible to pass it to `prefetchFiles`.

The reason why we want to access `ActionExecutionMetadata` inside `prefetchFiles` is to associate downloads with actions. For example:

1. When downloading from remote cache, we want to provide action metadata to the RPCs so that remote server can associate them with actions. `actionId` was removed from remote file metadata by f62a8b9 because it is not the right action to associate with for downloads inside this context. We left the action metadata empty for RPCs until this change. #17724 (comment)
2. When building with the bytes, we are able to report download progress to UI. However, `prefetchFiles` is the major way to download when building without the bytes and we can't report download progress on action basis. #17604 added the download report for background downloads, but now, we no longer download toplevel artifacts in background (a5dde12). With this change, we are able to report download progress for each action, either when they download outputs, or prefetch inputs for building without the bytes.

PiperOrigin-RevId: 539635790
Change-Id: Ic9cc1024f5d3560ac46bc462bd549dd81712d92f

Co-authored-by: Ian (Hee) Cha <[email protected]>

Author: coeuvre

src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java

@@ -24,7 +24,9 @@ public interface ActionInputPrefetcher {
       new ActionInputPrefetcher() {
         @Override
         public ListenableFuture<Void> prefetchFiles(
-            Iterable<? extends ActionInput> inputs, MetadataProvider metadataProvider) {
+            ActionExecutionMetadata action,
+            Iterable<? extends ActionInput> inputs,
+            MetadataProvider metadataProvider) {
           // Do nothing.
           return immediateVoidFuture();
         }
@@ -43,7 +45,9 @@ public boolean supportsPartialTreeArtifactInputs() {
    * @return future success if prefetch is finished or {@link IOException}.
    */
   ListenableFuture<Void> prefetchFiles(
-      Iterable<? extends ActionInput> inputs, MetadataProvider metadataProvider);
+      ActionExecutionMetadata action,
+      Iterable<? extends ActionInput> inputs,
+      MetadataProvider metadataProvider);
 
   /**
    * Whether the prefetcher is able to fetch individual files in a tree artifact without fetching

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

@@ -246,6 +246,7 @@ public ListenableFuture<Void> prefetchInputs()
         return actionExecutionContext
             .getActionInputPrefetcher()
             .prefetchFiles(
+                spawn.getResourceOwner(),
                 getInputMapping(PathFragment.EMPTY_FRAGMENT, /* willAccessRepeatedly= */ true)
                     .values(),
                 getMetadataProvider());

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

@@ -33,6 +33,7 @@
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.Action;
+import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
 import com.google.devtools.build.lib.actions.Artifact;
@@ -292,6 +293,7 @@ private boolean shouldDownloadFile(Path path, FileArtifactValue metadata) {
    * @param tempPath the temporary path which the input should be written to.
    */
   protected abstract ListenableFuture<Void> doDownloadFile(
+      ActionExecutionMetadata action,
       Reporter reporter,
       Path tempPath,
       PathFragment execPath,
@@ -317,11 +319,13 @@ protected Completable onErrorResumeNext(Throwable error) {
    */
   @Override
   public ListenableFuture<Void> prefetchFiles(
+      ActionExecutionMetadata action,
       Iterable<? extends ActionInput> inputs, MetadataProvider metadataProvider) {
-    return prefetchFiles(inputs, metadataProvider, Priority.MEDIUM);
+    return prefetchFiles(action, inputs, metadataProvider, Priority.MEDIUM);
   }
 
   protected ListenableFuture<Void> prefetchFiles(
+      ActionExecutionMetadata action,
       Iterable<? extends ActionInput> inputs,
       MetadataProvider metadataProvider,
       Priority priority) {
@@ -346,7 +350,7 @@ protected ListenableFuture<Void> prefetchFiles(
     Flowable<TransferResult> transfers =
         Flowable.fromIterable(files)
             .flatMapSingle(
-                input -> toTransferResult(prefetchFile(dirCtx, metadataProvider, input, priority)));
+                input -> toTransferResult(prefetchFile(action, dirCtx, metadataProvider, input, priority)));
 
     Completable prefetch =
         Completable.using(
@@ -357,6 +361,7 @@ protected ListenableFuture<Void> prefetchFiles(
   }
 
   private Completable prefetchFile(
+      ActionExecutionMetadata action,
       DirectoryContext dirCtx,
       MetadataProvider metadataProvider,
       ActionInput input,
@@ -386,6 +391,7 @@ private Completable prefetchFile(
 
     Completable result =
         downloadFileNoCheckRx(
+            action,
             dirCtx,
             execRoot.getRelative(execPath),
             treeRootExecPath != null ? execRoot.getRelative(treeRootExecPath) : null,
@@ -461,6 +467,7 @@ private Symlink maybeGetSymlink(
    * download finished.
    */
   private Completable downloadFileRx(
+      ActionExecutionMetadata action,
       DirectoryContext dirCtx,
       Path path,
       @Nullable Path treeRoot,
@@ -470,10 +477,11 @@ private Completable downloadFileRx(
     if (!canDownloadFile(path, metadata)) {
       return Completable.complete();
     }
-    return downloadFileNoCheckRx(dirCtx, path, treeRoot, actionInput, metadata, priority);
+    return downloadFileNoCheckRx(action, dirCtx, path, treeRoot, actionInput, metadata, priority);
   }
 
   private Completable downloadFileNoCheckRx(
+      ActionExecutionMetadata action,
       DirectoryContext dirCtx,
       Path path,
       @Nullable Path treeRoot,
@@ -498,6 +506,7 @@ private Completable downloadFileNoCheckRx(
                     toCompletable(
                             () ->
                                 doDownloadFile(
+                                    action,
                                     reporter,
                                     tempPath,
                                     finalPath.relativeTo(execRoot),
@@ -542,12 +551,15 @@ private Completable downloadFileNoCheckRx(
    * <p>The file will be written into a temporary file and moved to the final destination after the
    * download finished.
    */
-  public void downloadFile(Path path, @Nullable ActionInput actionInput, FileArtifactValue metadata)
+  public void downloadFile(
+      ActionExecutionMetadata action,
+      Path path, @Nullable ActionInput actionInput, FileArtifactValue metadata)
       throws IOException, InterruptedException {
-    getFromFuture(downloadFileAsync(path.asFragment(), actionInput, metadata, Priority.CRITICAL));
+    getFromFuture(downloadFileAsync(action, path.asFragment(), actionInput, metadata, Priority.CRITICAL));
   }
 
   protected ListenableFuture<Void> downloadFileAsync(
+      ActionExecutionMetadata action,
       PathFragment path,
       @Nullable ActionInput actionInput,
       FileArtifactValue metadata,
@@ -557,6 +569,7 @@ protected ListenableFuture<Void> downloadFileAsync(
             DirectoryContext::new,
             dirCtx ->
                 downloadFileRx(
+                    action,
                     dirCtx,
                     execRoot.getFileSystem().getPath(path),
                     /* treeRoot= */ null,
@@ -695,7 +708,7 @@ public void finalizeAction(Action action, MetadataHandler metadataHandler) {
     }
 
     if (!inputsToDownload.isEmpty()) {
-      var future = prefetchFiles(inputsToDownload, metadataHandler, Priority.HIGH);
+      var future = prefetchFiles(action, inputsToDownload, metadataHandler, Priority.HIGH);
       addCallback(
           future,
           new FutureCallback<Void>() {
@@ -716,7 +729,7 @@ public void onFailure(Throwable throwable) {
     }
 
     if (!outputsToDownload.isEmpty()) {
-      var future = prefetchFiles(outputsToDownload, metadataHandler, Priority.LOW);
+      var future = prefetchFiles(action, outputsToDownload, metadataHandler, Priority.LOW);
       addCallback(
           future,
           new FutureCallback<Void>() {

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

@@ -9,8 +9,8 @@ filegroup(
     srcs = glob(["*"]) + [
         "//src/main/java/com/google/devtools/build/lib/remote/circuitbreaker:srcs",
         "//src/main/java/com/google/devtools/build/lib/remote/common:srcs",
-        "//src/main/java/com/google/devtools/build/lib/remote/downloader:srcs",
         "//src/main/java/com/google/devtools/build/lib/remote/disk:srcs",
+        "//src/main/java/com/google/devtools/build/lib/remote/downloader:srcs",
         "//src/main/java/com/google/devtools/build/lib/remote/grpc:srcs",
         "//src/main/java/com/google/devtools/build/lib/remote/http:srcs",
         "//src/main/java/com/google/devtools/build/lib/remote/logging:srcs",
@@ -206,6 +206,7 @@ java_library(
     srcs = ["ToplevelArtifactsDownloader.java"],
     deps = [
         ":abstract_action_input_prefetcher",
+        "//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/actions:file_metadata",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",

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

@@ -24,6 +24,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputMap;
 import com.google.devtools.build.lib.actions.Artifact;
@@ -77,6 +78,7 @@ public class RemoteActionFileSystem extends DelegateFileSystem {
   private final RemoteActionInputFetcher inputFetcher;
   private final RemoteInMemoryFileSystem remoteOutputTree;
 
+  @Nullable private ActionExecutionMetadata action = null;
   @Nullable private MetadataInjector metadataInjector = null;
 
   RemoteActionFileSystem(
@@ -111,7 +113,8 @@ boolean isRemote(Path path) {
     return getRemoteMetadata(path.asFragment()) != null;
   }
 
-  public void updateContext(MetadataInjector metadataInjector) {
+  public void updateContext(ActionExecutionMetadata action, MetadataInjector metadataInjector) {
+    this.action = action;
     this.metadataInjector = metadataInjector;
   }
 
@@ -579,7 +582,7 @@ private void downloadFileIfRemote(PathFragment path) throws IOException {
     FileArtifactValue m = getRemoteMetadata(path);
     if (m != null) {
       try {
-        inputFetcher.downloadFile(delegateFs.getPath(path), getActionInput(path), m);
+        inputFetcher.downloadFile(action, delegateFs.getPath(path), getActionInput(path), m);
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
         throw new IOException(

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

@@ -21,6 +21,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
@@ -96,6 +97,7 @@ protected boolean canDownloadFile(Path path, FileArtifactValue metadata) {
 
   @Override
   protected ListenableFuture<Void> doDownloadFile(
+      ActionExecutionMetadata action,
       Reporter reporter,
       Path tempPath,
       PathFragment execPath,
@@ -104,7 +106,7 @@ protected ListenableFuture<Void> doDownloadFile(
       throws IOException {
     checkArgument(metadata.isRemote(), "Cannot download file that is not a remote file.");
     RequestMetadata requestMetadata =
-        TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "prefetcher", null);
+        TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "prefetcher", action);
     RemoteActionExecutionContext context = RemoteActionExecutionContext.create(requestMetadata);
 
     Digest digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.eventbus.Subscribe;
 import com.google.devtools.build.lib.actions.Action;
+import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInputMap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
@@ -84,11 +85,12 @@ public FileSystem createActionFileSystem(
 
   @Override
   public void updateActionFileSystemContext(
+      ActionExecutionMetadata action,
       FileSystem actionFileSystem,
       Environment env,
       MetadataInjector injector,
       ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
-    ((RemoteActionFileSystem) actionFileSystem).updateContext(injector);
+    ((RemoteActionFileSystem) actionFileSystem).updateContext(action, injector);
   }
 
   @Override