Anchor input fetches to source action id

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

@@ -111,6 +111,15 @@ public int getLocationIndex() {
     return 0;
   }
 
+  /**
+   * Remote action source identifier for the file.
+   *
+   * <p>"" indicates that a remote action output was not the source of this artifact.
+   */
+  public String getActionId() {
+    return "";
+  }
+
   /** Returns {@code true} if this is a special marker as opposed to a representing a real file. */
   public boolean isMarkerValue() {
     return this instanceof Singleton;
@@ -544,11 +553,13 @@ public static final class RemoteFileArtifactValue extends FileArtifactValue {
     private final byte[] digest;
     private final long size;
     private final int locationIndex;
+    private final String actionId;
 
-    public RemoteFileArtifactValue(byte[] digest, long size, int locationIndex) {
+    public RemoteFileArtifactValue(byte[] digest, long size, int locationIndex, String actionId) {
       this.digest = digest;
       this.size = size;
       this.locationIndex = locationIndex;
+      this.actionId = actionId;
     }
 
     @Override
@@ -588,6 +599,11 @@ public long getSize() {
       return size;
     }
 
+    @Override
+    public String getActionId() {
+      return actionId;
+    }
+
     @Override
     public long getModifiedTime() {
       throw new UnsupportedOperationException(

src/main/java/com/google/devtools/build/lib/actions/cache/MetadataInjector.java

@@ -31,8 +31,9 @@ public interface MetadataInjector {
    * @param digest the digest of the file.
    * @param size the size of the file in bytes.
    * @param locationIndex is only used in Blaze.
+   * @param actionId the id of the action that produced this file.
    */
-  void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex);
+  void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex, String actionId);
 
   /**
    * Inject the metadata of a tree artifact whose contents are stored remotely.

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

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.remote;
 
 import build.bazel.remote.execution.v2.Digest;
+import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
@@ -32,6 +33,7 @@
 import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
+import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import io.grpc.Context;
 import java.io.IOException;
@@ -65,12 +67,12 @@ class RemoteActionInputFetcher implements ActionInputPrefetcher {
 
   private final RemoteCache remoteCache;
   private final Path execRoot;
-  private final Context ctx;
+  private final RequestMetadata requestMetadata;
 
-  RemoteActionInputFetcher(RemoteCache remoteCache, Path execRoot, Context ctx) {
+  RemoteActionInputFetcher(RemoteCache remoteCache, Path execRoot, RequestMetadata requestMetadata) {
     this.remoteCache = Preconditions.checkNotNull(remoteCache);
     this.execRoot = Preconditions.checkNotNull(execRoot);
-    this.ctx = Preconditions.checkNotNull(ctx);
+    this.requestMetadata = Preconditions.checkNotNull(requestMetadata);
   }
 
   /**
@@ -164,6 +166,10 @@ private ListenableFuture<Void> downloadFileAsync(Path path, FileArtifactValue me
 
       ListenableFuture<Void> download = downloadsInProgress.get(path);
       if (download == null) {
+        Context ctx = TracingMetadataUtils.contextWithMetadata(
+            requestMetadata.toBuilder()
+                .setActionId(metadata.getActionId())
+                .build());
         Context prevCtx = ctx.attach();
         try {
           Digest digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());

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

@@ -528,6 +528,7 @@ private List<ListenableFuture<FileMetadata>> downloadOutErr(ActionResult result,
    */
   @Nullable
   public InMemoryOutput downloadMinimal(
+      String actionId,
       ActionResult result,
       Collection<? extends ActionInput> outputs,
       @Nullable PathFragment inMemoryOutputPath,
@@ -570,7 +571,7 @@ public InMemoryOutput downloadMinimal(
         inMemoryOutput = output;
       }
       if (output instanceof Artifact) {
-        injectRemoteArtifact((Artifact) output, metadata, execRoot, metadataInjector);
+        injectRemoteArtifact((Artifact) output, metadata, execRoot, metadataInjector, actionId);
       }
     }
 
@@ -594,7 +595,8 @@ private void injectRemoteArtifact(
       Artifact output,
       ActionResultMetadata metadata,
       Path execRoot,
-      MetadataInjector metadataInjector)
+      MetadataInjector metadataInjector,
+      String actionId)
       throws IOException {
     if (output.isTreeArtifact()) {
       DirectoryMetadata directory =
@@ -617,7 +619,8 @@ private void injectRemoteArtifact(
             new RemoteFileArtifactValue(
                 DigestUtil.toBinaryDigest(file.digest()),
                 file.digest().getSizeBytes(),
-                /* locationIndex= */ 1);
+                /* locationIndex= */ 1,
+                actionId);
         childMetadata.put(p, r);
       }
       metadataInjector.injectRemoteDirectory(
@@ -633,7 +636,8 @@ private void injectRemoteArtifact(
           output,
           DigestUtil.toBinaryDigest(outputMetadata.digest()),
           outputMetadata.digest().getSizeBytes(),
-          /* locationIndex= */ 1);
+          /* locationIndex= */ 1,
+          actionId);
     }
   }
 

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

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.remote;
 
 import build.bazel.remote.execution.v2.DigestFunction;
+import build.bazel.remote.execution.v2.RequestMetadata;
 import build.bazel.remote.execution.v2.ServerCapabilities;
 import com.google.auth.Credentials;
 import com.google.common.base.Preconditions;
@@ -662,12 +663,14 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB
             env.getOptions().getOptions(RemoteOptions.class), "RemoteOptions");
     RemoteOutputsMode remoteOutputsMode = remoteOptions.remoteOutputsMode;
     if (!remoteOutputsMode.downloadAllOutputs()) {
-      Context ctx =
-          TracingMetadataUtils.contextWithMetadata(
-              env.getBuildRequestId(), env.getCommandId().toString(), "fetch-remote-inputs");
+      RequestMetadata requestMetadata =
+          RequestMetadata.newBuilder()
+              .setCorrelatedInvocationsId(env.getBuildRequestId())
+              .setToolInvocationId(env.getCommandId().toString())
+              .build();
       actionInputFetcher =
           new RemoteActionInputFetcher(
-              actionContextProvider.getRemoteCache(), env.getExecRoot(), ctx);
+              actionContextProvider.getRemoteCache(), env.getExecRoot(), requestMetadata);
       builder.setActionInputPrefetcher(actionInputFetcher);
       remoteOutputService.setActionInputFetcher(actionInputFetcher);
     }

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

@@ -186,6 +186,7 @@ public CacheHandle lookup(Spawn spawn, SpawnExecutionContext context)
                 prof.profile(ProfilerTask.REMOTE_DOWNLOAD, "download outputs minimal")) {
               inMemoryOutput =
                   remoteCache.downloadMinimal(
+                      actionKey.getDigest().getHash(),
                       result,
                       spawn.getOutputFiles(),
                       inMemoryOutputPath,

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

@@ -262,6 +262,7 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
           } else {
             try {
               return downloadAndFinalizeSpawnResult(
+                  actionKey.getDigest().getHash(),
                   cachedResult,
                   /* cacheHit= */ true,
                   spawn,
@@ -341,6 +342,7 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
 
               try {
                 return downloadAndFinalizeSpawnResult(
+                    actionKey.getDigest().getHash(),
                     actionResult,
                     reply.getCachedResult(),
                     spawn,
@@ -413,6 +415,7 @@ static void spawnMetricsAccounting(
   }
 
   private SpawnResult downloadAndFinalizeSpawnResult(
+      String actionId,
       ActionResult actionResult,
       boolean cacheHit,
       Spawn spawn,
@@ -441,6 +444,7 @@ private SpawnResult downloadAndFinalizeSpawnResult(
           Profiler.instance().profile(REMOTE_DOWNLOAD, "download outputs minimal")) {
         inMemoryOutput =
             remoteCache.downloadMinimal(
+                actionId,
                 actionResult,
                 spawn.getOutputFiles(),
                 inMemoryOutputPath,

src/main/java/com/google/devtools/build/lib/remote/util/TracingMetadataUtils.java

@@ -70,18 +70,21 @@ public static Context contextWithMetadata(
     Preconditions.checkNotNull(buildRequestId);
     Preconditions.checkNotNull(commandId);
     Preconditions.checkNotNull(actionId);
-    RequestMetadata.Builder metadata =
+    RequestMetadata metadata =
         RequestMetadata.newBuilder()
             .setCorrelatedInvocationsId(buildRequestId)
-            .setToolInvocationId(commandId);
-    metadata.setActionId(actionId);
-    metadata
-        .setToolDetails(
-            ToolDetails.newBuilder()
-                .setToolName("bazel")
-                .setToolVersion(BlazeVersionInfo.instance().getVersion()))
-        .build();
-    return Context.current().withValue(CONTEXT_KEY, metadata.build());
+            .setToolInvocationId(commandId)
+            .setActionId(actionId)
+            .setToolDetails(
+                ToolDetails.newBuilder()
+                    .setToolName("bazel")
+                    .setToolVersion(BlazeVersionInfo.instance().getVersion()))
+            .build();
+    return contextWithMetadata(metadata);
+  }
+
+  public static Context contextWithMetadata(RequestMetadata metadata) {
+    return Context.current().withValue(CONTEXT_KEY, metadata);
   }
 
   /**

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

@@ -542,7 +542,7 @@ public void injectDigest(ActionInput output, FileStatus statNoFollow, byte[] dig
   }
 
   @Override
-  public void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex) {
+  public void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex, String actionId) {
     Preconditions.checkArgument(
         isKnownOutput(output), output + " is not a declared output of this action");
     Preconditions.checkArgument(
@@ -551,7 +551,7 @@ public void injectRemoteFile(Artifact output, byte[] digest, long size, int loca
         output);
     Preconditions.checkState(
         executionMode.get(), "Tried to inject %s outside of execution", output);
-    store.injectRemoteFile(output, digest, size, locationIndex);
+    store.injectRemoteFile(output, digest, size, locationIndex, actionId);
   }
 
   @Override

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

@@ -95,9 +95,9 @@ void addTreeArtifactContents(Artifact artifact, TreeFileArtifact contents) {
     treeArtifactContents.computeIfAbsent(artifact, a -> Sets.newConcurrentHashSet()).add(contents);
   }
 
-  void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex) {
+  void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex, String actionId) {
     injectOutputData(
-        output, new FileArtifactValue.RemoteFileArtifactValue(digest, size, locationIndex));
+        output, new FileArtifactValue.RemoteFileArtifactValue(digest, size, locationIndex, actionId));
   }
 
   final void injectOutputData(Artifact output, FileArtifactValue artifactValue) {

src/test/java/com/google/devtools/build/lib/actions/util/ActionsTestUtil.java

@@ -967,7 +967,7 @@ public void injectDigest(ActionInput output, FileStatus statNoFollow, byte[] dig
     }
 
     @Override
-    public void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex) {
+    public void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex, String actionId) {
       throw new UnsupportedOperationException();
     }
 

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

@@ -360,7 +360,7 @@ private Artifact createRemoteArtifact(
     byte[] b = contents.getBytes(StandardCharsets.UTF_8);
     HashCode h = HashCode.fromString(DIGEST_UTIL.compute(b).getHash());
     FileArtifactValue f =
-        new RemoteFileArtifactValue(h.asBytes(), b.length, /* locationIndex= */ 1);
+        new RemoteFileArtifactValue(h.asBytes(), b.length, /* locationIndex= */ 1, "action-id");
     inputs.putWithNoDepOwner(a, f);
     return a;
   }

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

@@ -172,7 +172,7 @@ private Artifact createRemoteArtifact(
     byte[] b = contents.getBytes(StandardCharsets.UTF_8);
     HashCode h = HASH_FUNCTION.getHashFunction().hashBytes(b);
     FileArtifactValue f =
-        new RemoteFileArtifactValue(h.asBytes(), b.length, /* locationIndex= */ 1);
+        new RemoteFileArtifactValue(h.asBytes(), b.length, /* locationIndex= */ 1, "action-id");
     inputs.putWithNoDepOwner(a, f);
     return a;
   }

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

@@ -17,6 +17,7 @@
 import static org.junit.Assert.assertThrows;
 
 import build.bazel.remote.execution.v2.Digest;
+import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
@@ -85,7 +86,7 @@ public void testFetching() throws Exception {
     MetadataProvider metadataProvider = new StaticMetadataProvider(metadata);
     RemoteCache remoteCache = newCache(options, digestUtil, cacheEntries);
     RemoteActionInputFetcher actionInputFetcher =
-        new RemoteActionInputFetcher(remoteCache, execRoot, Context.current());
+        new RemoteActionInputFetcher(remoteCache, execRoot, RequestMetadata.getDefaultInstance());
 
     // act
     actionInputFetcher.prefetchFiles(metadata.keySet(), metadataProvider);
@@ -108,7 +109,7 @@ public void testStagingVirtualActionInput() throws Exception {
     MetadataProvider metadataProvider = new StaticMetadataProvider(new HashMap<>());
     RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
     RemoteActionInputFetcher actionInputFetcher =
-        new RemoteActionInputFetcher(remoteCache, execRoot, Context.current());
+        new RemoteActionInputFetcher(remoteCache, execRoot, RequestMetadata.getDefaultInstance());
     VirtualActionInput a = new StringActionInput("hello world", PathFragment.create("file1"));
 
     // act
@@ -133,7 +134,7 @@ public void testFileNotFound() throws Exception {
     MetadataProvider metadataProvider = new StaticMetadataProvider(metadata);
     RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
     RemoteActionInputFetcher actionInputFetcher =
-        new RemoteActionInputFetcher(remoteCache, execRoot, Context.current());
+        new RemoteActionInputFetcher(remoteCache, execRoot, RequestMetadata.getDefaultInstance());
 
     // act
     assertThrows(
@@ -157,7 +158,7 @@ public void testIgnoreNoneRemoteFiles() throws Exception {
     MetadataProvider metadataProvider = new StaticMetadataProvider(ImmutableMap.of(a, f));
     RemoteCache remoteCache = newCache(options, digestUtil, new HashMap<>());
     RemoteActionInputFetcher actionInputFetcher =
-        new RemoteActionInputFetcher(remoteCache, execRoot, Context.current());
+        new RemoteActionInputFetcher(remoteCache, execRoot, RequestMetadata.getDefaultInstance());
 
     // act
     actionInputFetcher.prefetchFiles(ImmutableList.of(a), metadataProvider);
@@ -175,7 +176,7 @@ public void testDownloadFile() throws Exception {
     Artifact a1 = createRemoteArtifact("file1", "hello world", metadata, cacheEntries);
     RemoteCache remoteCache = newCache(options, digestUtil, cacheEntries);
     RemoteActionInputFetcher actionInputFetcher =
-        new RemoteActionInputFetcher(remoteCache, execRoot, Context.current());
+        new RemoteActionInputFetcher(remoteCache, execRoot, RequestMetadata.getDefaultInstance());
 
     // act
     actionInputFetcher.downloadFile(a1.getPath(), metadata.get(a1));
@@ -198,7 +199,7 @@ private Artifact createRemoteArtifact(
     byte[] b = contents.getBytes(StandardCharsets.UTF_8);
     HashCode h = HASH_FUNCTION.getHashFunction().hashBytes(b);
     FileArtifactValue f =
-        new RemoteFileArtifactValue(h.asBytes(), b.length, /* locationIndex= */ 1);
+        new RemoteFileArtifactValue(h.asBytes(), b.length, /* locationIndex= */ 1, "action-id");
     metadata.put(a, f);
     cacheEntries.put(DigestUtil.buildDigest(h.asBytes(), b.length), ByteString.copyFrom(b));
     return a;