Inject metadata when creating a filesystem symlink for a non-symlink …

…artifact.

A ctx.actions.symlink whose output is a declare_file or declare_directory
(as opposed to a declare_symlink) has "copy" semantics, i.e., the output
artifact is indistinguishable from the referent except for its name; the
symlink is just a filesystem-level optimization to avoid an actual copy,
and is transparently resolved when collecting the action output metadata.

When the symlink points to an artifact that was built remotely and without the
bytes, we currently must download it before executing the symlink action, so
that the output metadata can be constructed from the local filesystem. This
change short-circuits the filesystem traversal by injecting output metadata,
which is identical to the input plus a pointer to the original path. This is
used by the prefetcher to avoid downloading the same files multiple times when
they're symlinked more than once.

(An alternative would have been to teach all of the RemoteActionFileSystem
methods to resolve symlinks by patching together the local and remote metadata,
but that would have resulted in an awful lot of complexity.)

Fixes bazelbuild#15678.

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

@@ -84,6 +84,21 @@ TrieArtifact getOrPutSubFolder(String name) {
       return trieArtifact;
     }
 
+    @Nullable
+    TreeArtifactValue findTreeArtifactNode(PathFragment execPath) {
+      TrieArtifact current = this;
+      for (String segment : execPath.segments()) {
+        current = current.subFolders.get(segment);
+        if (current == null) {
+          break;
+        }
+      }
+      if (current != null && current.treeArtifactValue != null) {
+        return current.treeArtifactValue;
+      }
+      return null;
+    }
+
     @Nullable
     TreeArtifactValue findTreeArtifactNodeAtPrefix(PathFragment execPath) {
       TrieArtifact current = this;
@@ -233,6 +248,11 @@ public FileArtifactValue getMetadata(PathFragment execPath) {
     return getMetadataFromTreeArtifacts(execPath);
   }
 
+  @Nullable
+  public TreeArtifactValue getTreeMetadata(PathFragment execPath) {
+    return treeArtifactsRoot.findTreeArtifactNode(execPath);
+  }
+
   @Nullable
   private FileArtifactValue getMetadataFromTreeArtifacts(PathFragment execPath) {
     TreeArtifactValue tree = treeArtifactsRoot.findTreeArtifactNodeAtPrefix(execPath);

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

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.actions;
 
 import static java.nio.charset.StandardCharsets.ISO_8859_1;
+import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
@@ -539,16 +540,44 @@ public static class RemoteFileArtifactValue extends FileArtifactValue {
     private final long size;
     private final int locationIndex;
     private final String actionId;
-
-    public RemoteFileArtifactValue(byte[] digest, long size, int locationIndex, String actionId) {
+    @Nullable private final PathFragment stagingExecPath;
+
+    /**
+     * Creates a {@link RemoteFileArtifactValue} with a hint that the contents should be downloaded
+     * to some other location.
+     *
+     * <p>This is used by ctx.actions.symlink to implement efficient copies.
+     */
+    public static RemoteFileArtifactValue createWithStagingExecPath(
+        RemoteFileArtifactValue original, PathFragment stagingExecPath) {
+      return new RemoteFileArtifactValue(
+          original.digest,
+          original.size,
+          original.locationIndex,
+          original.actionId,
+          // If this is a double indirection, download to the original location.
+          original.stagingExecPath != null ? original.stagingExecPath : stagingExecPath);
+    }
+
+    private RemoteFileArtifactValue(
+        byte[] digest,
+        long size,
+        int locationIndex,
+        String actionId,
+        PathFragment stagingExecPath) {
       this.digest = Preconditions.checkNotNull(digest, actionId);
       this.size = size;
       this.locationIndex = locationIndex;
       this.actionId = actionId;
+      this.stagingExecPath = stagingExecPath;
+    }
+
+    public RemoteFileArtifactValue(byte[] digest, long size, int locationIndex, String actionId) {
+      this(digest, size, locationIndex, actionId, /* symlinkTargetPath= */ null);
     }
 
     public RemoteFileArtifactValue(byte[] digest, long size, int locationIndex) {
-      this(digest, size, locationIndex, /* actionId= */ "");
+      this(digest, size, locationIndex, /* actionId= */ "", /* stagingExecPath= */ null);
     }
 
     @Override
@@ -561,7 +590,8 @@ public boolean equals(Object o) {
       return Arrays.equals(digest, that.digest)
           && size == that.size
           && locationIndex == that.locationIndex
-          && Objects.equals(actionId, that.actionId);
+          && Objects.equals(actionId, that.actionId)
+          && Objects.equals(stagingExecPath, that.stagingExecPath);
     }
 
     @Override
@@ -615,6 +645,11 @@ public boolean isRemote() {
       return true;
     }
 
+    @Nullable
+    public PathFragment getStagingExecPath() {
+      return stagingExecPath;
+    }
+
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(this)

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

@@ -23,6 +23,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.flogger.GoogleLogger;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
@@ -31,17 +32,20 @@
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
+import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 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;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Single;
 import java.io.IOException;
+import java.rmi.Remote;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -153,8 +157,27 @@ private Completable prefetchInputTree(
       MetadataProvider provider,
       SpecialArtifact tree,
       List<TreeFileArtifact> treeFiles,
-      Priority priority) {
-    Path treeRoot = execRoot.getRelative(tree.getExecPath());
+      Priority priority)
+      throws IOException {
+    PathFragment execPath = tree.getExecPath();
+
+    if (!treeFiles.isEmpty()) {
+      // Grab the tree artifact root from the first file.
+      // TODO(tjgq): Refactor this class so that this isn't needed.
+      TreeFileArtifact treeFile = Iterables.getFirst(treeFiles, null);
+      FileArtifactValue metadata = provider.getMetadata(treeFile);
+      if (metadata instanceof RemoteFileArtifactValue) {
+        PathFragment stagingExecPath = ((RemoteFileArtifactValue) metadata).getStagingExecPath();
+        if (stagingExecPath != null) {
+          execPath =
+              stagingExecPath.subFragment(
+                  0,
+                  stagingExecPath.segmentCount() - treeFile.getParentRelativePath().segmentCount());
+        }
+      }
+    }
+
+    Path treeRoot = execRoot.getRelative(execPath);
     HashMap<TreeFileArtifact, Path> treeFileTmpPathMap = new HashMap<>();
 
     Flowable<TransferResult> transfers =
@@ -243,6 +266,14 @@ private Completable prefetchInputFile(
     }
 
     Path path = execRoot.getRelative(input.getExecPath());
+
+    if (metadata instanceof RemoteFileArtifactValue) {
+      PathFragment stagingExecPath = ((RemoteFileArtifactValue) metadata).getStagingExecPath();
+      if (stagingExecPath != null) {
+        path = execRoot.getRelative(stagingExecPath);
+      }
+    }
+
     return downloadFileRx(path, metadata, priority);
   }
 

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

@@ -178,6 +178,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/remote/util",
         "//src/main/java/com/google/devtools/build/lib/vfs",
+        "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:rxjava3",

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

@@ -16,7 +16,12 @@
 package com.google.devtools.build.lib.remote;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.Streams.stream;
 
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputMap;
 import com.google.devtools.build.lib.actions.Artifact;
@@ -43,7 +48,7 @@
 
 /**
  * This is a basic implementation and incomplete implementation of an action file system that's been
- * tuned to what native (non-spawn) actions in Bazel currently use.
+ * tuned to what internal (non-spawn) actions in Bazel currently use.
  *
  * <p>The implementation mostly delegates to the local file system except for the case where an
  * action input is a remotely stored action output. Most notably {@link
@@ -52,12 +57,13 @@
  * <p>This implementation only supports creating local action outputs.
  */
 public class RemoteActionFileSystem extends DelegateFileSystem {
-
   private final PathFragment execRoot;
   private final PathFragment outputBase;
   private final ActionInputMap inputArtifactData;
+  private final ImmutableMap<PathFragment, Artifact> outputArtifactByExecPath;
   private final RemoteActionInputFetcher inputFetcher;
-  private final Map<PathFragment, RemoteFileArtifactValue> injectedMetadata = new HashMap<>();
+  private final HashMap<PathFragment, RemoteFileArtifactValue> injectedMetadata = new HashMap<>();
+  private final HashMap<PathFragment, TreeArtifactValue> injectedTreeMetadata = new HashMap<>();
 
   @Nullable private MetadataInjector metadataInjector = null;
 
@@ -66,11 +72,15 @@ public class RemoteActionFileSystem extends DelegateFileSystem {
       PathFragment execRootFragment,
       String relativeOutputPath,
       ActionInputMap inputArtifactData,
+      Iterable<Artifact> outputArtifacts,
       RemoteActionInputFetcher inputFetcher) {
     super(localDelegate);
     this.execRoot = checkNotNull(execRootFragment, "execRootFragment");
     this.outputBase = execRoot.getRelative(checkNotNull(relativeOutputPath, "relativeOutputPath"));
     this.inputArtifactData = checkNotNull(inputArtifactData, "inputArtifactData");
+    this.outputArtifactByExecPath =
+        stream(checkNotNull(outputArtifacts, "outputArtifacts"))
+            .collect(toImmutableMap(Artifact::getExecPath, a -> a));
     this.inputFetcher = checkNotNull(inputFetcher, "inputFetcher");
   }
 
@@ -86,12 +96,13 @@ public void updateContext(MetadataInjector metadataInjector) {
   void injectTree(SpecialArtifact tree, TreeArtifactValue metadata) {
     checkNotNull(metadataInjector, "metadataInject is null");
 
+    injectedTreeMetadata.put(tree.getExecPath(), metadata);
+
     for (Map.Entry<TreeFileArtifact, FileArtifactValue> entry :
         metadata.getChildValues().entrySet()) {
       FileArtifactValue childMetadata = entry.getValue();
       if (childMetadata instanceof RemoteFileArtifactValue) {
-        TreeFileArtifact child = entry.getKey();
-        injectedMetadata.put(child.getExecPath(), (RemoteFileArtifactValue) childMetadata);
+        injectedMetadata.put(entry.getKey().getExecPath(), (RemoteFileArtifactValue) childMetadata);
       }
     }
 
@@ -227,14 +238,63 @@ protected PathFragment readSymbolicLink(PathFragment path) throws IOException {
   @Override
   protected void createSymbolicLink(PathFragment linkPath, PathFragment targetFragment)
       throws IOException {
-    /*
-     * TODO(buchgr): Optimize the case where we are creating a symlink to a remote output. This does
-     * add a non-trivial amount of complications though (as symlinks tend to do).
-     */
-    downloadFileIfRemote(execRoot.getRelative(targetFragment));
+    FileStatus stat = statNullable(linkPath, /* followSymlinks= */ false);
+    if (stat != null) {
+      throw new IOException(linkPath + " (File exists)");
+    }
+
+    // If the output is a regular file or tree artifact, it has "copy" semantics, i.e., it will be
+    // transparently resolved into the artifact it points to; the symlink is just a filesystem-level
+    // optimization to avoid making an actual copy. We can inject its metadata and avoid the need to
+    // compute it from the filesystem later.
+    Artifact output = outputArtifactByExecPath.get(linkPath.relativeTo(execRoot));
+    if (output != null) {
+      if (output.isTreeArtifact()) {
+        TreeArtifactValue metadata = getRemoteInputTreeMetadata(targetFragment);
+        if (metadata != null) {
+          // The archived representation is a Google-internal feature disabled on open source Bazel.
+          checkState(
+              metadata.getArchivedRepresentation().isEmpty(),
+              "symlink to archived tree artifact is not supported");
+          checkState(targetFragment.isAbsolute());
+          injectTree(
+              (SpecialArtifact) output,
+              createTreeArtifactValueWithStagingExecPath(
+                  (SpecialArtifact) output, metadata, targetFragment.relativeTo(execRoot)));
+        }
+      } else if (!output.isSymlink()) {
+        RemoteFileArtifactValue metadata = getRemoteInputMetadata(targetFragment);
+        if (metadata != null) {
+          checkState(targetFragment.isAbsolute());
+          injectFile(
+              output,
+              RemoteFileArtifactValue.createWithStagingExecPath(
+                  metadata, targetFragment.relativeTo(execRoot)));
+        }
+      }
+    }
+
     super.createSymbolicLink(linkPath, targetFragment);
   }
 
+  @VisibleForTesting
+  static TreeArtifactValue createTreeArtifactValueWithStagingExecPath(
+      SpecialArtifact output, TreeArtifactValue originalValue, PathFragment targetExecPath) {
+    TreeArtifactValue.Builder treeBuilder = TreeArtifactValue.newBuilder(output);
+
+    for (Map.Entry<TreeFileArtifact, FileArtifactValue> entry :
+        originalValue.getChildValues().entrySet()) {
+      PathFragment parentRelativePath = entry.getKey().getParentRelativePath();
+      treeBuilder.putChild(
+          TreeFileArtifact.createTreeOutput(output, parentRelativePath),
+          RemoteFileArtifactValue.createWithStagingExecPath(
+              (RemoteFileArtifactValue) entry.getValue(),
+              targetExecPath.getRelative(parentRelativePath)));
+    }
+
+    return treeBuilder.build();
+  }
+
   // -------------------- Implementations based on stat() --------------------
 
   @Override
@@ -370,10 +430,22 @@ private RemoteFileArtifactValue getRemoteInputMetadata(PathFragment path) {
     if (m != null && m.isRemote()) {
       return (RemoteFileArtifactValue) m;
     }
-
     return injectedMetadata.get(execPath);
   }
 
+  @Nullable
+  private TreeArtifactValue getRemoteInputTreeMetadata(PathFragment path) {
+    if (!path.startsWith(outputBase)) {
+      return null;
+    }
+    PathFragment execPath = path.relativeTo(execRoot);
+    TreeArtifactValue m = inputArtifactData.getTreeMetadata(execPath);
+    if (m != null && m.isEntirelyRemote()) {
+      return m;
+    }
+    return injectedTreeMetadata.get(execPath);
+  }
+
   private void downloadFileIfRemote(PathFragment path) throws IOException {
     FileArtifactValue m = getRemoteInputMetadata(path);
     if (m != null) {

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

@@ -70,6 +70,7 @@ public FileSystem createActionFileSystem(
         execRootFragment,
         relativeOutputPath,
         inputArtifactData,
+        outputArtifacts,
         actionInputFetcher);
   }
 
@@ -142,7 +143,12 @@ public ArtifactPathResolver createPathResolverForArtifactValues(
       Map<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
     FileSystem remoteFileSystem =
         new RemoteActionFileSystem(
-            fileSystem, execRoot, relativeOutputPath, actionInputMap, actionInputFetcher);
+            fileSystem,
+            execRoot,
+            relativeOutputPath,
+            actionInputMap,
+            ImmutableList.of(),
+            actionInputFetcher);
     return ArtifactPathResolver.createPathResolver(remoteFileSystem, fileSystem.getPath(execRoot));
   }
 }

src/test/java/com/google/devtools/build/lib/actions/ActionInputMapTest.java

@@ -134,6 +134,10 @@ public void putTreeArtifact_addsTreeArtifactAndAllChildren() {
     assertContainsEqualMetadata(tree, treeValue.getMetadata());
     assertContainsSameInstance(child1, child1Metadata);
     assertContainsSameInstance(child2, child2Metadata);
+
+    assertThat(map.getMetadata(tree)).isEqualTo(treeValue.getMetadata());
+    assertThat(map.getTreeMetadata(tree.getExecPath())).isEqualTo(treeValue);
+    assertThat(map.getTreeMetadata(child1.getExecPath())).isNull();
   }
 
   @Test