Wire up PathMapper in sandbox code

`PathMapper`s rewrite paths in command lines to make them more cache friendly, which requires executor support to stage files at the rewritten paths. This commit wires up the `PathMapper` used by a given `Spawn` with sandbox outputs logic for sandboxed and worker sandboxed execution.

An end-to-end test will be added in bazelbuild#18155, but requires bazelbuild#19718, bazelbuild#19719, and bazelbuild#19721.

Work towards bazelbuild#6526

Closes bazelbuild#19719.

PiperOrigin-RevId: 571950903
Change-Id: I4a9d15da578da11f8677279b09e0864b8c33e9fc

src/main/java/com/google/devtools/build/lib/sandbox/AbstractContainerizingSandboxedSpawn.java

@@ -124,8 +124,7 @@ public void createFileSystem() throws IOException, InterruptedException {
         dirsToCreate,
         Iterables.concat(
             ImmutableSet.of(), inputs.getFiles().keySet(), inputs.getSymlinks().keySet()),
-        outputs.files(),
-        outputs.dirs());
+        outputs);
 
     // Allow subclasses to filter out inputs and dirs that don't need to be created.
     filterInputsAndDirsToCreate(inputsToCreate, dirsToCreate);

src/main/java/com/google/devtools/build/lib/sandbox/SandboxHelpers.java

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.sandbox;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.devtools.build.lib.vfs.Dirent.Type.DIRECTORY;
 import static com.google.devtools.build.lib.vfs.Dirent.Type.SYMLINK;
 
@@ -85,9 +86,10 @@ public final class SandboxHelpers {
    */
   public static void moveOutputs(SandboxOutputs outputs, Path sourceRoot, Path targetRoot)
       throws IOException {
-    for (PathFragment output : Iterables.concat(outputs.files(), outputs.dirs())) {
-      Path source = sourceRoot.getRelative(output);
-      Path target = targetRoot.getRelative(output);
+    for (Entry<PathFragment, PathFragment> output :
+        Iterables.concat(outputs.files().entrySet(), outputs.dirs().entrySet())) {
+      Path source = sourceRoot.getRelative(output.getValue());
+      Path target = targetRoot.getRelative(output.getKey());
       if (source.isFile() || source.isSymbolicLink()) {
         // Ensure the target directory exists in the target. The directories for the action outputs
         // have already been created, but the spawn outputs may be different from the overall action
@@ -221,8 +223,7 @@ public static void populateInputsAndDirsToCreate(
       Set<PathFragment> inputsToCreate,
       LinkedHashSet<PathFragment> dirsToCreate,
       Iterable<PathFragment> inputFiles,
-      ImmutableSet<PathFragment> outputFiles,
-      ImmutableSet<PathFragment> outputDirs) {
+      SandboxOutputs outputs) {
     // Add all worker files, input files, and the parent directories.
     for (PathFragment input : inputFiles) {
       inputsToCreate.add(input);
@@ -231,12 +232,12 @@ public static void populateInputsAndDirsToCreate(
 
     // And all parent directories of output files. Note that we don't add the files themselves --
     // any pre-existing files that have the same path as an output should get deleted.
-    for (PathFragment file : outputFiles) {
+    for (PathFragment file : outputs.files().values()) {
       dirsToCreate.add(file.getParentDirectory());
     }
 
     // Add all output directories.
-    dirsToCreate.addAll(outputDirs);
+    dirsToCreate.addAll(outputs.dirs().values());
 
     // Add some directories that should be writable, and thus exist.
     dirsToCreate.addAll(writableDirs);
@@ -547,32 +548,43 @@ public SandboxInputs processInputFiles(
   /** The file and directory outputs of a sandboxed spawn. */
   @AutoValue
   public abstract static class SandboxOutputs {
-    public abstract ImmutableSet<PathFragment> files();
 
-    public abstract ImmutableSet<PathFragment> dirs();
+    /** A map from output file exec paths to paths in the sandbox. */
+    public abstract ImmutableMap<PathFragment, PathFragment> files();
+
+    /** A map from output directory exec paths to paths in the sandbox. */
+    public abstract ImmutableMap<PathFragment, PathFragment> dirs();
 
     private static final SandboxOutputs EMPTY_OUTPUTS =
-        SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of());
+        SandboxOutputs.create(ImmutableMap.of(), ImmutableMap.of());
 
     public static SandboxOutputs create(
-        ImmutableSet<PathFragment> files, ImmutableSet<PathFragment> dirs) {
+        ImmutableMap<PathFragment, PathFragment> files,
+        ImmutableMap<PathFragment, PathFragment> dirs) {
       return new AutoValue_SandboxHelpers_SandboxOutputs(files, dirs);
     }
 
+    public static SandboxOutputs create(
+        ImmutableSet<PathFragment> files, ImmutableSet<PathFragment> dirs) {
+      return new AutoValue_SandboxHelpers_SandboxOutputs(
+          files.stream().collect(toImmutableMap(f -> f, f -> f)),
+          dirs.stream().collect(toImmutableMap(d -> d, d -> d)));
+    }
+
     public static SandboxOutputs getEmptyInstance() {
       return EMPTY_OUTPUTS;
     }
   }
 
   public SandboxOutputs getOutputs(Spawn spawn) {
-    ImmutableSet.Builder<PathFragment> files = ImmutableSet.builder();
-    ImmutableSet.Builder<PathFragment> dirs = ImmutableSet.builder();
+    ImmutableMap.Builder<PathFragment, PathFragment> files = ImmutableMap.builder();
+    ImmutableMap.Builder<PathFragment, PathFragment> dirs = ImmutableMap.builder();
     for (ActionInput output : spawn.getOutputFiles()) {
-      PathFragment path = PathFragment.create(output.getExecPathString());
+      PathFragment mappedPath = spawn.getPathMapper().map(output.getExecPath());
       if (output instanceof Artifact && ((Artifact) output).isTreeArtifact()) {
-        dirs.add(path);
+        dirs.put(output.getExecPath(), mappedPath);
       } else {
-        files.add(path);
+        files.put(output.getExecPath(), mappedPath);
       }
     }
     return SandboxOutputs.create(files.build(), dirs.build());

src/main/java/com/google/devtools/build/lib/sandbox/SandboxfsSandboxedSpawn.java

@@ -136,10 +136,10 @@ class SandboxfsSandboxedSpawn implements SandboxedSpawn {
     this.arguments = arguments;
     this.environment = environment;
     this.inputs = inputs;
-    for (PathFragment path : outputs.files()) {
+    for (PathFragment path : outputs.files().values()) {
       checkArgument(!path.isAbsolute(), "outputs %s must be relative", path);
     }
-    for (PathFragment path : outputs.dirs()) {
+    for (PathFragment path : outputs.dirs().values()) {
       checkArgument(!path.isAbsolute(), "outputs %s must be relative", path);
     }
     this.outputs = outputs;
@@ -203,10 +203,10 @@ public void createFileSystem() throws IOException {
     sandboxScratchDir.createDirectory();
 
     Set<PathFragment> dirsToCreate = new HashSet<>(writableDirs);
-    for (PathFragment output : outputs.files()) {
+    for (PathFragment output : outputs.files().values()) {
       dirsToCreate.add(output.getParentDirectory());
     }
-    dirsToCreate.addAll(outputs.dirs());
+    dirsToCreate.addAll(outputs.dirs().values());
     for (PathFragment dir : dirsToCreate) {
       sandboxScratchDir.getRelative(dir).createDirectoryAndParents();
     }

src/main/java/com/google/devtools/build/lib/worker/SandboxedWorkerProxy.java

@@ -79,8 +79,7 @@ public void prepareExecution(
         inputsToCreate,
         dirsToCreate,
         Iterables.concat(inputFiles.getFiles().keySet(), inputFiles.getSymlinks().keySet()),
-        outputs.files(),
-        outputs.dirs());
+        outputs);
     SandboxHelpers.cleanExisting(
         sandboxDir.getParentDirectory(), inputFiles, inputsToCreate, dirsToCreate, sandboxDir);
     // Finally, create anything that is still missing. This is non-strict only for historical

src/main/java/com/google/devtools/build/lib/worker/WorkerExecRoot.java

@@ -58,8 +58,7 @@ public void createFileSystem(
         inputsToCreate,
         dirsToCreate,
         Iterables.concat(workerFiles, inputs.getFiles().keySet(), inputs.getSymlinks().keySet()),
-        outputs.files(),
-        outputs.dirs());
+        outputs);
 
     // Then do a full traversal of the parent directory of `workDir`. This will use what we computed
     // above, delete anything unnecessary and update `inputsToCreate`/`dirsToCreate` if something is

src/main/java/com/google/devtools/build/lib/worker/WorkerMultiplexer.java

@@ -154,8 +154,7 @@ public synchronized void createSandboxedProcess(
           inputsToCreate,
           dirsToCreate,
           workerFiles,
-          SandboxOutputs.getEmptyInstance().files(),
-          SandboxOutputs.getEmptyInstance().dirs());
+          SandboxOutputs.getEmptyInstance());
       SandboxHelpers.cleanExisting(
           workDir.getParentDirectory(), inputFiles, inputsToCreate, dirsToCreate, workDir);
       SandboxHelpers.createDirectories(dirsToCreate, workDir, /* strict=*/ false);

src/test/java/com/google/devtools/build/lib/sandbox/BUILD

@@ -85,6 +85,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs",
         "//src/main/java/com/google/devtools/common/options",
         "//src/test/java/com/google/devtools/build/lib/actions/util",
+        "//src/test/java/com/google/devtools/build/lib/exec/util",
         "//src/test/java/com/google/devtools/build/lib/testutil",
         "//src/test/java/com/google/devtools/build/lib/testutil:TestUtils",
         "//third_party:auto_value",

src/test/java/com/google/devtools/build/lib/sandbox/SandboxHelpersTest.java

@@ -24,11 +24,15 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.ArtifactRoot;
 import com.google.devtools.build.lib.actions.CommandLines.ParamFileActionInput;
 import com.google.devtools.build.lib.actions.ParameterFile.ParameterFileType;
+import com.google.devtools.build.lib.actions.PathMapper;
+import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
 import com.google.devtools.build.lib.exec.BinTools;
+import com.google.devtools.build.lib.exec.util.SpawnBuilder;
 import com.google.devtools.build.lib.sandbox.SandboxHelpers.SandboxInputs;
 import com.google.devtools.build.lib.sandbox.SandboxHelpers.SandboxOutputs;
 import com.google.devtools.build.lib.testutil.Scratch;
@@ -268,11 +272,7 @@ public void cleanExisting_updatesDirs() throws IOException, InterruptedException
         Iterables.concat(
             ImmutableSet.of(), inputs.getFiles().keySet(), inputs.getSymlinks().keySet()),
         SandboxOutputs.create(
-                ImmutableSet.of(PathFragment.create("out/dir/output.txt")), ImmutableSet.of())
-            .files(),
-        SandboxOutputs.create(
-                ImmutableSet.of(PathFragment.create("out/dir/output.txt")), ImmutableSet.of())
-            .dirs());
+            ImmutableSet.of(PathFragment.create("out/dir/output.txt")), ImmutableSet.of()));
 
     PathFragment inputDir1 = input1.getParentDirectory();
     PathFragment inputDir2 = input2.getParentDirectory();
@@ -303,4 +303,64 @@ public void cleanExisting_updatesDirs() throws IOException, InterruptedException
     assertThat(execRootPath.getRelative("out").exists()).isTrue();
     assertThat(execRootPath.getRelative("out/dir").exists()).isFalse();
   }
+
+  @Test
+  public void populateInputsAndDirsToCreate_createsMappedDirectories() {
+    ArtifactRoot outputRoot =
+        ArtifactRoot.asDerivedRoot(execRootPath, ArtifactRoot.RootType.Output, "outputs");
+    ActionInput outputFile = ActionsTestUtil.createArtifact(outputRoot, "bin/config/dir/file");
+    ActionInput outputDir =
+        ActionsTestUtil.createTreeArtifactWithGeneratingAction(
+            outputRoot, "bin/config/other_dir/subdir");
+    PathMapper pathMapper =
+        execPath -> PathFragment.create(execPath.getPathString().replace("config/", ""));
+    Spawn spawn =
+        new SpawnBuilder().withOutputs(outputFile, outputDir).setPathMapper(pathMapper).build();
+    var sandboxHelpers = new SandboxHelpers();
+    LinkedHashSet<PathFragment> writableDirs = new LinkedHashSet<>();
+    LinkedHashSet<PathFragment> inputsToCreate = new LinkedHashSet<>();
+    LinkedHashSet<PathFragment> dirsToCreate = new LinkedHashSet<>();
+
+    SandboxHelpers.populateInputsAndDirsToCreate(
+        writableDirs,
+        inputsToCreate,
+        dirsToCreate,
+        ImmutableList.of(),
+        sandboxHelpers.getOutputs(spawn));
+
+    assertThat(writableDirs).isEmpty();
+    assertThat(inputsToCreate).isEmpty();
+    assertThat(dirsToCreate)
+        .containsExactly(
+            PathFragment.create("outputs/bin/dir"),
+            PathFragment.create("outputs/bin/other_dir/subdir"));
+  }
+
+  @Test
+  public void moveOutputs_mappedPathMovedToUnmappedPath() throws Exception {
+    PathFragment unmappedOutputPath = PathFragment.create("bin/config/output");
+    PathMapper pathMapper =
+        execPath -> PathFragment.create(execPath.getPathString().replace("config/", ""));
+    Spawn spawn =
+        new SpawnBuilder()
+            .withOutputs(unmappedOutputPath.getPathString())
+            .setPathMapper(pathMapper)
+            .build();
+    var sandboxHelpers = new SandboxHelpers();
+    Path sandboxBase = execRootPath.getRelative("sandbox");
+    PathFragment mappedOutputPath = PathFragment.create("bin/output");
+    sandboxBase.getRelative(mappedOutputPath).getParentDirectory().createDirectoryAndParents();
+    FileSystemUtils.writeLinesAs(
+        sandboxBase.getRelative(mappedOutputPath), UTF_8, "hello", "pathmapper");
+
+    Path realBase = execRootPath.getRelative("real");
+    SandboxHelpers.moveOutputs(sandboxHelpers.getOutputs(spawn), sandboxBase, realBase);
+
+    assertThat(
+            FileSystemUtils.readLines(
+                realBase.getRelative(unmappedOutputPath.getPathString()), UTF_8))
+        .containsExactly("hello", "pathmapper")
+        .inOrder();
+    assertThat(sandboxBase.getRelative(mappedOutputPath).exists()).isFalse();
+  }
 }