Wire up PathMapper in RemoteExecutionService

`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` in `RemoteExecutionService`, which ensures that paths of inputs
and outputs are correctly mapped before being sent off to the remote
executor and mapped back to the correct local paths when downloading the
results.

Work towards #6526

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

@@ -117,6 +117,10 @@ public Command getCommand() {
     return command;
   }
 
+  public RemotePathResolver getRemotePathResolver() {
+    return remotePathResolver;
+  }
+
   @Nullable
   public MerkleTree getMerkleTree() {
     return merkleTree;

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

@@ -160,7 +160,13 @@ public class RemoteExecutionService {
   private final Reporter reporter;
   private final boolean verboseFailures;
   private final Path execRoot;
-  private final RemotePathResolver remotePathResolver;
+
+  /**
+   * Do not use directly, instead use the per-spawn resolver created in {@link
+   * #buildRemoteAction(Spawn, SpawnExecutionContext)}.
+   */
+  private final RemotePathResolver baseRemotePathResolver;
+
   private final String buildRequestId;
   private final String commandId;
   private final DigestUtil digestUtil;
@@ -200,7 +206,7 @@ public RemoteExecutionService(
     this.reporter = reporter;
     this.verboseFailures = verboseFailures;
     this.execRoot = execRoot;
-    this.remotePathResolver = remotePathResolver;
+    this.baseRemotePathResolver = remotePathResolver;
     this.buildRequestId = buildRequestId;
     this.commandId = commandId;
     this.digestUtil = digestUtil;
@@ -343,7 +349,8 @@ Cache<Object, CompletableFuture<MerkleTree>> getMerkleTreeCache() {
     return merkleTreeCache;
   }
 
-  private SortedMap<PathFragment, ActionInput> buildOutputDirMap(Spawn spawn) {
+  private SortedMap<PathFragment, ActionInput> buildOutputDirMap(
+      Spawn spawn, RemotePathResolver remotePathResolver) {
     TreeMap<PathFragment, ActionInput> outputDirMap = new TreeMap<>();
     for (ActionInput output : spawn.getOutputFiles()) {
       if (output instanceof Artifact && ((Artifact) output).isTreeArtifact()) {
@@ -360,12 +367,14 @@ private MerkleTree buildInputMerkleTree(
       Spawn spawn,
       SpawnExecutionContext context,
       ToolSignature toolSignature,
-      @Nullable SpawnScrubber spawnScrubber)
+      @Nullable SpawnScrubber spawnScrubber,
+      RemotePathResolver remotePathResolver)
       throws IOException, ForbiddenActionInputException {
     // Add output directories to inputs so that they are created as empty directories by the
     // executor. The spec only requires the executor to create the parent directory of an output
     // directory, which differs from the behavior of both local and sandboxed execution.
-    SortedMap<PathFragment, ActionInput> outputDirMap = buildOutputDirMap(spawn);
+    SortedMap<PathFragment, ActionInput> outputDirMap =
+        buildOutputDirMap(spawn, remotePathResolver);
     boolean useMerkleTreeCache = remoteOptions.remoteMerkleTreeCache;
     if (toolSignature != null || spawnScrubber != null) {
       // The Merkle tree cache is not yet compatible with scrubbing or marking tool files.
@@ -531,10 +540,16 @@ public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context
       throws IOException, ExecException, ForbiddenActionInputException, InterruptedException {
     remoteActionBuildingSemaphore.acquire();
     try {
+      // Create a remote path resolver that is aware of the spawn's path mapper, which rewrites
+      // the paths of the inputs and outputs as well as paths appearing in the command line for
+      // execution. This is necessary to ensure that artifacts are correctly emitted into and staged
+      // from the unmapped location locally.
+      RemotePathResolver remotePathResolver =
+          RemotePathResolver.createMapped(baseRemotePathResolver, execRoot, spawn.getPathMapper());
       ToolSignature toolSignature = getToolSignature(spawn, context);
       SpawnScrubber spawnScrubber = scrubber != null ? scrubber.forSpawn(spawn) : null;
       final MerkleTree merkleTree =
-          buildInputMerkleTree(spawn, context, toolSignature, spawnScrubber);
+          buildInputMerkleTree(spawn, context, toolSignature, spawnScrubber, remotePathResolver);
 
       // Get the remote platform properties.
       Platform platform;
@@ -751,7 +766,8 @@ private ListenableFuture<FileMetadata> downloadFile(
       RemoteActionExecutionContext context,
       ProgressStatusListener progressStatusListener,
       FileMetadata file,
-      Path tmpPath) {
+      Path tmpPath,
+      RemotePathResolver remotePathResolver) {
     checkNotNull(remoteCache, "remoteCache can't be null");
 
     try {
@@ -992,7 +1008,9 @@ private DirectoryMetadata parseDirectory(
   }
 
   ActionResultMetadata parseActionResultMetadata(
-      RemoteActionExecutionContext context, RemoteActionResult result)
+      RemoteActionExecutionContext context,
+      RemoteActionResult result,
+      RemotePathResolver remotePathResolver)
       throws IOException, InterruptedException {
     checkNotNull(remoteCache, "remoteCache can't be null");
 
@@ -1097,7 +1115,7 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
 
     ActionResultMetadata metadata;
     try (SilentCloseable c = Profiler.instance().profile("Remote.parseActionResultMetadata")) {
-      metadata = parseActionResultMetadata(context, result);
+      metadata = parseActionResultMetadata(context, result, action.getRemotePathResolver());
     }
 
     // The expiration time for remote cache entries.
@@ -1134,7 +1152,9 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
       if (!isInMemoryOutputFile && shouldDownload(result, execPath)) {
         Path tmpPath = tempPathGenerator.generateTempPath();
         realToTmpPath.put(file.path, tmpPath);
-        downloadsBuilder.add(downloadFile(context, progressStatusListener, file, tmpPath));
+        downloadsBuilder.add(
+            downloadFile(
+                context, progressStatusListener, file, tmpPath, action.getRemotePathResolver()));
       } else {
         remoteActionFileSystem.injectRemoteFile(
             file.path().asFragment(),
@@ -1164,7 +1184,9 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
         if (shouldDownload(result, file.path.relativeTo(execRoot))) {
           Path tmpPath = tempPathGenerator.generateTempPath();
           realToTmpPath.put(file.path, tmpPath);
-          downloadsBuilder.add(downloadFile(context, progressStatusListener, file, tmpPath));
+          downloadsBuilder.add(
+              downloadFile(
+                  context, progressStatusListener, file, tmpPath, action.getRemotePathResolver()));
         } else {
           remoteActionFileSystem.injectRemoteFile(
               file.path().asFragment(),
@@ -1322,7 +1344,7 @@ private Single<UploadManifest> buildUploadManifestAsync(
               remoteOptions,
               remoteCache.getCacheCapabilities(),
               digestUtil,
-              remotePathResolver,
+              action.getRemotePathResolver(),
               action.getActionKey(),
               action.getAction(),
               action.getCommand(),
@@ -1442,7 +1464,9 @@ public void uploadInputsIfNotPresent(RemoteAction action, boolean force)
         SpawnExecutionContext context = action.getSpawnExecutionContext();
         ToolSignature toolSignature = getToolSignature(spawn, context);
         SpawnScrubber spawnScrubber = scrubber != null ? scrubber.forSpawn(spawn) : null;
-        merkleTree = buildInputMerkleTree(spawn, context, toolSignature, spawnScrubber);
+        merkleTree =
+            buildInputMerkleTree(
+                spawn, context, toolSignature, spawnScrubber, action.getRemotePathResolver());
       }
 
       remoteExecutionCache.ensureInputsPresent(

src/main/java/com/google/devtools/build/lib/remote/common/RemotePathResolver.java

@@ -15,16 +15,20 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ForbiddenActionInputException;
+import com.google.devtools.build.lib.actions.PathMapper;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.exec.SpawnInputExpander;
+import com.google.devtools.build.lib.exec.SpawnInputExpander.InputVisitor;
 import com.google.devtools.build.lib.exec.SpawnRunner.SpawnExecutionContext;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
 import java.util.SortedMap;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * A {@link RemotePathResolver} is used to resolve input/output paths for remote execution from
@@ -225,4 +229,70 @@ public Path outputPathToLocalPath(ActionInput actionInput) {
       return ActionInputHelper.toInputPath(actionInput, execRoot);
     }
   }
+
+  /**
+   * Adapts a given base {@link RemotePathResolver} to also apply a {@link PathMapper} to map (and
+   * inverse map) paths.
+   */
+  static RemotePathResolver createMapped(
+      RemotePathResolver base, Path execRoot, PathMapper pathMapper) {
+    if (pathMapper.isNoop()) {
+      return base;
+    }
+    return new RemotePathResolver() {
+      private final ConcurrentHashMap<PathFragment, PathFragment> inverse =
+          new ConcurrentHashMap<>();
+
+      @Override
+      public String getWorkingDirectory() {
+        return base.getWorkingDirectory();
+      }
+
+      @Override
+      public SortedMap<PathFragment, ActionInput> getInputMapping(
+          SpawnExecutionContext context, boolean willAccessRepeatedly)
+          throws IOException, ForbiddenActionInputException {
+        return base.getInputMapping(context, willAccessRepeatedly);
+      }
+
+      @Override
+      public void walkInputs(Spawn spawn, SpawnExecutionContext context, InputVisitor visitor)
+          throws IOException, ForbiddenActionInputException {
+        base.walkInputs(spawn, context, visitor);
+      }
+
+      @Override
+      public String localPathToOutputPath(Path path) {
+        return localPathToOutputPath(path.relativeTo(execRoot));
+      }
+
+      @Override
+      public String localPathToOutputPath(PathFragment execPath) {
+        return base.localPathToOutputPath(map(execPath));
+      }
+
+      @Override
+      public Path outputPathToLocalPath(String outputPath) {
+        return execRoot.getRelative(
+            inverseMap(base.outputPathToLocalPath(outputPath).relativeTo(execRoot)));
+      }
+
+      private PathFragment map(PathFragment path) {
+        PathFragment mappedPath = pathMapper.map(path);
+        PathFragment previousPath = inverse.put(mappedPath, path);
+        Preconditions.checkState(
+            previousPath == null || previousPath.equals(path),
+            "Two different paths %s and %s map to the same path %s",
+            previousPath,
+            path,
+            mappedPath);
+        return mappedPath;
+      }
+
+      private PathFragment inverseMap(PathFragment path) {
+        return Preconditions.checkNotNull(
+            inverse.get(path), "Failed to find original path for mapped path %s", path);
+      }
+    };
+  }
 }

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

@@ -97,6 +97,7 @@
 import com.google.devtools.build.lib.remote.common.RemotePathResolver.SiblingRepositoryLayoutResolver;
 import com.google.devtools.build.lib.remote.merkletree.MerkleTree;
 import com.google.devtools.build.lib.remote.options.RemoteOptions;
+import com.google.devtools.build.lib.remote.options.RemoteOutputsMode;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.remote.util.FakeSpawnExecutionContext;
 import com.google.devtools.build.lib.remote.util.InMemoryCacheClient;
@@ -115,6 +116,8 @@
 import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
 import com.google.devtools.common.options.Options;
 import com.google.protobuf.ByteString;
+import com.google.testing.junit.testparameterinjector.TestParameter;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import java.io.IOException;
 import java.util.Random;
 import java.util.concurrent.CountDownLatch;
@@ -126,14 +129,13 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
 /** Tests for {@link RemoteExecutionService}. */
-@RunWith(JUnit4.class)
+@RunWith(TestParameterInjector.class)
 public class RemoteExecutionServiceTest {
   @Rule public final MockitoRule mockito = MockitoJUnit.rule();
   @Rule public final RxNoGlobalErrorsRule rxNoGlobalErrorsRule = new RxNoGlobalErrorsRule();
@@ -1540,6 +1542,52 @@ public void downloadOutputs_missingMandatoryOutputs_reportError() throws Excepti
     assertThat(error).hasMessageThat().containsMatch("expected output .+ does not exist.");
   }
 
+  @Test
+  public void downloadOutputs_pathUnmapped() throws Exception {
+    // Test that the output of a remote action with path mapping applied is downloaded into the
+    // correct unmapped local path.
+    Digest d1 = cache.addContents(remoteActionExecutionContext, "content1");
+    Digest d2 = cache.addContents(remoteActionExecutionContext, "content2");
+    Artifact output1 = ActionsTestUtil.createArtifact(artifactRoot, "bin/config/dir/output1");
+    Artifact output2 = ActionsTestUtil.createArtifact(artifactRoot, "bin/other_dir/output2");
+    ActionResult r =
+        ActionResult.newBuilder()
+            .setExitCode(0)
+            // The action result includes the mapped paths.
+            .addOutputFiles(
+                OutputFile.newBuilder().setPath("outputs/bin/dir/output1").setDigest(d1))
+            .addOutputFiles(
+                OutputFile.newBuilder().setPath("outputs/bin/other_dir/output2").setDigest(d2))
+            .build();
+    PathMapper pathMapper =
+        execPath -> PathFragment.create(execPath.getPathString().replaceAll("config/", ""));
+    Spawn spawn =
+        new SpawnBuilder("unused")
+            .withOutput(output1)
+            .withOutput(output2)
+            .setPathMapper(pathMapper)
+            .build();
+    RemoteActionResult result = RemoteActionResult.createFromCache(CachedActionResult.remote(r));
+    FakeSpawnExecutionContext context = newSpawnExecutionContext(spawn);
+    remoteOutputChecker =
+        new RemoteOutputChecker(
+            new JavaClock(), "build", RemoteOutputsMode.TOPLEVEL, ImmutableList.of());
+    remoteOutputChecker.addOutputToDownload(output1);
+    remoteOutputChecker.addOutputToDownload(output2);
+    RemoteExecutionService service = newRemoteExecutionService();
+    RemoteAction action = service.buildRemoteAction(spawn, context);
+
+    InMemoryOutput inMemoryOutput = service.downloadOutputs(action, result);
+
+    assertThat(inMemoryOutput).isNull();
+    RemoteActionFileSystem actionFs = context.getActionFileSystem();
+    assertThat(actionFs.getDigest(output1.getPath().asFragment())).isEqualTo(toBinaryDigest(d1));
+    assertThat(readContent(output1.getPath(), UTF_8)).isEqualTo("content1");
+    assertThat(actionFs.getDigest(output2.getPath().asFragment())).isEqualTo(toBinaryDigest(d2));
+    assertThat(readContent(output2.getPath(), UTF_8)).isEqualTo("content2");
+    assertThat(context.isLockOutputFilesCalled()).isTrue();
+  }
+
   @Test
   public void uploadOutputs_uploadDirectory_works() throws Exception {
     // Test that uploading a directory works.
@@ -2266,6 +2314,61 @@ public void buildRemoteActionWithScrubbing() throws Exception {
                 .toByteString());
   }
 
+  @Test
+  public void buildRemoteActionWithPathMapping(@TestParameter boolean remoteMerkleTreeCache)
+      throws Exception {
+    remoteOptions.remoteMerkleTreeCache = remoteMerkleTreeCache;
+
+    var mappedInput = ActionsTestUtil.createArtifact(artifactRoot, "bin/config/input1");
+    fakeFileCache.createScratchInput(mappedInput, "value1");
+    var unmappedInput = ActionsTestUtil.createArtifact(artifactRoot, "bin/input2");
+    fakeFileCache.createScratchInput(unmappedInput, "value2");
+    var outputDir =
+        ActionsTestUtil.createTreeArtifactWithGeneratingAction(
+            artifactRoot, "bin/config/output_dir");
+    PathMapper pathMapper =
+        execPath -> PathFragment.create(execPath.getPathString().replaceAll("config/", ""));
+    Spawn spawn =
+        new SpawnBuilder("unused")
+            .withInputs(mappedInput, unmappedInput)
+            .withOutputs("outputs/bin/config/dir/output1", "outputs/bin/other_dir/output2")
+            .withOutputs(outputDir)
+            .setPathMapper(pathMapper)
+            .build();
+    FakeSpawnExecutionContext context = newSpawnExecutionContext(spawn);
+    RemoteExecutionService service = newRemoteExecutionService(remoteOptions);
+
+    // Check that inputs and outputs of the remote action are mapped correctly.
+    var remoteAction = service.buildRemoteAction(spawn, context);
+    assertThat(remoteAction.getInputMap(false))
+        .containsExactly(
+            PathFragment.create("outputs/bin/input1"), mappedInput,
+            PathFragment.create("outputs/bin/input2"), unmappedInput);
+    assertThat(remoteAction.getCommand().getOutputFilesList())
+        .containsExactly("outputs/bin/dir/output1", "outputs/bin/other_dir/output2");
+    assertThat(remoteAction.getCommand().getOutputDirectoriesList())
+        .containsExactly("outputs/bin/output_dir");
+    assertThat(remoteAction.getCommand().getOutputPathsList())
+        .containsExactly(
+            "outputs/bin/dir/output1", "outputs/bin/other_dir/output2", "outputs/bin/output_dir");
+
+    // Check that the Merkle tree nodes are mapped correctly, including the output directory.
+    var merkleTree = remoteAction.getMerkleTree();
+    var outputsDirectory =
+        merkleTree.getDirectoryByDigest(merkleTree.getRootProto().getDirectories(0).getDigest());
+    assertThat(outputsDirectory.getDirectoriesCount()).isEqualTo(1);
+    var binDirectory =
+        merkleTree.getDirectoryByDigest(outputsDirectory.getDirectories(0).getDigest());
+    assertThat(
+            binDirectory.getFilesList().stream().map(FileNode::getName).collect(toImmutableList()))
+        .containsExactly("input1", "input2");
+    assertThat(
+            binDirectory.getDirectoriesList().stream()
+                .map(DirectoryNode::getName)
+                .collect(toImmutableList()))
+        .containsExactly("output_dir");
+  }
+
   private Spawn newSpawnFromResult(RemoteActionResult result) {
     return newSpawnFromResult(ImmutableMap.of(), result);
   }