Add ActionExecutionMetadata as a parameter to `ActionInputPrefetche…

…r#prefetchFiles`.

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

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

@@ -34,6 +34,7 @@ public interface MetadataSupplier {
       new ActionInputPrefetcher() {
         @Override
         public ListenableFuture<Void> prefetchFiles(
+            ActionExecutionMetadata action,
             Iterable<? extends ActionInput> inputs,
             MetadataSupplier metadataSupplier,
             Priority priority) {
@@ -79,7 +80,10 @@ public enum Priority {
    * @return future success if prefetch is finished or {@link IOException}.
    */
   ListenableFuture<Void> prefetchFiles(
-      Iterable<? extends ActionInput> inputs, MetadataSupplier metadataSupplier, Priority priority);
+      ActionExecutionMetadata action,
+      Iterable<? extends ActionInput> inputs,
+      MetadataSupplier metadataSupplier,
+      Priority priority);
 
   /**
    * Whether the prefetcher requires the metadata for a tree artifact to be available whenever one

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

@@ -252,6 +252,7 @@ public ListenableFuture<Void> prefetchInputs()
             actionExecutionContext
                 .getActionInputPrefetcher()
                 .prefetchFiles(
+                    spawn.getResourceOwner(),
                     getInputMapping(PathFragment.EMPTY_FRAGMENT, /* willAccessRepeatedly= */ true)
                         .values(),
                     getInputMetadataProvider()::getInputMetadata,

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

@@ -29,6 +29,7 @@
 import com.google.common.flogger.GoogleLogger;
 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;
@@ -263,6 +264,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,
@@ -285,6 +287,7 @@ protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOExc
    */
   @Override
   public ListenableFuture<Void> prefetchFiles(
+      ActionExecutionMetadata action,
       Iterable<? extends ActionInput> inputs,
       MetadataSupplier metadataSupplier,
       Priority priority) {
@@ -308,7 +311,8 @@ public ListenableFuture<Void> prefetchFiles(
 
     Flowable<TransferResult> transfers =
         Flowable.fromIterable(files)
-            .flatMapSingle(input -> prefetchFile(dirCtx, metadataSupplier, input, priority));
+            .flatMapSingle(
+                input -> prefetchFile(action, dirCtx, metadataSupplier, input, priority));
 
     Completable prefetch =
         Completable.using(
@@ -318,6 +322,7 @@ public ListenableFuture<Void> prefetchFiles(
   }
 
   private Single<TransferResult> prefetchFile(
+      ActionExecutionMetadata action,
       DirectoryContext dirCtx,
       MetadataSupplier metadataSupplier,
       ActionInput input,
@@ -347,6 +352,7 @@ private Single<TransferResult> prefetchFile(
 
       Completable result =
           downloadFileNoCheckRx(
+              action,
               dirCtx,
               execRoot.getRelative(execPath),
               treeRootExecPath != null ? execRoot.getRelative(treeRootExecPath) : null,
@@ -421,6 +427,7 @@ private Symlink maybeGetSymlink(
   }
 
   private Completable downloadFileNoCheckRx(
+      ActionExecutionMetadata action,
       DirectoryContext dirCtx,
       Path path,
       @Nullable Path treeRoot,
@@ -445,6 +452,7 @@ private Completable downloadFileNoCheckRx(
                     toCompletable(
                             () ->
                                 doDownloadFile(
+                                    action,
                                     reporter,
                                     tempPath,
                                     finalPath.relativeTo(execRoot),
@@ -603,7 +611,7 @@ public void finalizeAction(Action action, OutputMetadataStore outputMetadataStor
       try (var s = Profiler.instance().profile(ProfilerTask.REMOTE_DOWNLOAD, "Download outputs")) {
         getFromFuture(
             prefetchFiles(
-                outputsToDownload, outputMetadataStore::getOutputMetadata, Priority.HIGH));
+                action, outputsToDownload, outputMetadataStore::getOutputMetadata, Priority.HIGH));
       }
     }
   }

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

@@ -25,6 +25,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.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ActionInputMap;
@@ -81,6 +82,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(
@@ -121,7 +123,8 @@ private boolean isRemote(PathFragment path) {
     return getRemoteMetadata(path) != null;
   }
 
-  public void updateContext(MetadataInjector metadataInjector) {
+  public void updateContext(ActionExecutionMetadata action, MetadataInjector metadataInjector) {
+    this.action = action;
     this.metadataInjector = metadataInjector;
   }
 
@@ -625,7 +628,7 @@ private void downloadFileIfRemote(PathFragment path) throws IOException {
       }
       getFromFuture(
           inputFetcher.prefetchFiles(
-              ImmutableList.of(input), this::getInputMetadata, Priority.CRITICAL));
+              action, ImmutableList.of(input), this::getInputMetadata, Priority.CRITICAL));
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new IOException(String.format("Received interrupt while fetching file '%s'", path), e);

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

@@ -14,18 +14,15 @@
 package com.google.devtools.build.lib.remote;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.devtools.build.lib.remote.common.ProgressStatusListener.NO_ACTION;
 
 import build.bazel.remote.execution.v2.Digest;
 import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
-import com.google.devtools.build.lib.events.ExtendedEventHandler.FetchProgress;
 import com.google.devtools.build.lib.events.Reporter;
-import com.google.devtools.build.lib.exec.SpawnProgressEvent;
-import com.google.devtools.build.lib.remote.RemoteCache.DownloadProgressReporter;
 import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.remote.util.TempPathGenerator;
@@ -79,6 +76,7 @@ protected boolean canDownloadFile(Path path, FileArtifactValue metadata) {
 
   @Override
   protected ListenableFuture<Void> doDownloadFile(
+      ActionExecutionMetadata action,
       Reporter reporter,
       Path tempPath,
       PathFragment execPath,
@@ -87,47 +85,18 @@ 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());
 
-    DownloadProgressReporter downloadProgressReporter;
-    if (priority == Priority.LOW) {
-      // Only report download progress for toplevel outputs
-      downloadProgressReporter =
-          new DownloadProgressReporter(
-              /* includeFile= */ false,
-              progress -> reporter.post(new DownloadProgress(progress)),
-              execPath.toString(),
-              metadata.getSize());
-    } else {
-      downloadProgressReporter = new DownloadProgressReporter(NO_ACTION, "", 0);
-    }
-
-    return remoteCache.downloadFile(context, tempPath, digest, downloadProgressReporter);
-  }
-
-  public static class DownloadProgress implements FetchProgress {
-    private final SpawnProgressEvent progress;
-
-    public DownloadProgress(SpawnProgressEvent progress) {
-      this.progress = progress;
-    }
-
-    @Override
-    public String getResourceIdentifier() {
-      return progress.progressId();
-    }
-
-    @Override
-    public String getProgress() {
-      return progress.progress();
-    }
-
-    @Override
-    public boolean isFinished() {
-      return progress.finished();
-    }
+    return remoteCache.downloadFile(
+        context,
+        tempPath,
+        digest,
+        new RemoteCache.DownloadProgressReporter(
+            progress -> progress.postTo(reporter, action),
+            execPath.toString(),
+            digest.getSizeBytes()));
   }
 }

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

@@ -21,6 +21,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;
@@ -99,11 +100,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

src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java

@@ -382,11 +382,13 @@ FileSystem createActionFileSystem(
   }
 
   private void updateActionFileSystemContext(
+      Action action,
       FileSystem actionFileSystem,
       Environment env,
       MetadataInjector metadataInjector,
       ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
-    outputService.updateActionFileSystemContext(actionFileSystem, env, metadataInjector, filesets);
+    outputService.updateActionFileSystemContext(
+        action, actionFileSystem, env, metadataInjector, filesets);
   }
 
   void executionOver() {
@@ -489,7 +491,8 @@ ActionExecutionValue executeAction(
       boolean hasDiscoveredInputs)
       throws ActionExecutionException, InterruptedException {
     if (actionFileSystem != null) {
-      updateActionFileSystemContext(actionFileSystem, env, metadataHandler, expandedFilesets);
+      updateActionFileSystemContext(
+          action, actionFileSystem, env, metadataHandler, expandedFilesets);
     }
 
     ActionExecutionContext actionExecutionContext =
@@ -834,6 +837,7 @@ NestedSet<Artifact> discoverInputs(
             threadStateReceiverFactory.apply(actionLookupData));
     if (actionFileSystem != null) {
       updateActionFileSystemContext(
+          action,
           actionFileSystem,
           env,
           THROWING_METADATA_INJECTOR_FOR_ACTIONFS,

src/main/java/com/google/devtools/build/lib/vfs/OutputService.java

@@ -18,6 +18,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 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;
@@ -196,6 +197,7 @@ default FileSystem createActionFileSystem(
    * @param filesets The Fileset symlinks known for this action.
    */
   default void updateActionFileSystemContext(
+      ActionExecutionMetadata action,
       FileSystem actionFileSystem,
       Environment env,
       MetadataInjector injector,