[6.3.0] Add ActionExecutionMetadata as a parameter to `ActionInputPrefetche…

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