Path map command lines lazily during expansion

This removes a peak memory inefficiency of path mapping compared to regular execution.

Closes bazelbuild#25368.

PiperOrigin-RevId: 747933848
Change-Id: Ib2d8ffb626b8181dd7d831d9a5207bd13b581420

Author: fmeum

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

@@ -28,13 +28,13 @@ public interface ArgChunk {
    * <p>The returned {@link Iterable} may lazily materialize strings during iteration, so consumers
    * should attempt to avoid iterating more times than necessary.
    */
-  Iterable<String> arguments();
+  Iterable<String> arguments(PathMapper pathMapper);
 
   /**
    * Counts the total length of all arguments in this chunk.
    *
    * <p>Implementations that lazily materialize strings may be able to compute the total argument
    * length without actually materializing the arguments.
    */
-  int totalArgLength();
+  int totalArgLength(PathMapper pathMapper);
 }

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

@@ -138,7 +138,6 @@ java_library(
     deps = [
         ":action_lookup_data",
         ":action_lookup_key",
-        ":arg_chunk",
         ":artifact_owner",
         ":artifacts",
         ":commandline_item",
@@ -168,8 +167,6 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
-        "//src/main/java/com/google/devtools/build/lib/profiler",
-        "//src/main/java/com/google/devtools/build/lib/profiler:google-auto-profiler-utils",
         "//src/main/java/com/google/devtools/build/lib/remote:store",
         "//src/main/java/com/google/devtools/build/lib/shell",
         "//src/main/java/com/google/devtools/build/lib/skyframe:detailed_exceptions",
@@ -213,11 +210,6 @@ java_library(
     ],
 )
 
-java_library(
-    name = "arg_chunk",
-    srcs = ["ArgChunk.java"],
-)
-
 java_library(
     name = "action_environment",
     srcs = ["ActionEnvironment.java"],
@@ -299,6 +291,7 @@ java_library(
     srcs = [
         "ActionInput.java",
         "ActionKeyContext.java",
+        "ArgChunk.java",
         "Artifact.java",
         "ArtifactCodecs.java",
         "ArtifactFactory.java",
@@ -312,7 +305,6 @@ java_library(
     deps = [
         ":action_lookup_data",
         ":action_lookup_key",
-        ":arg_chunk",
         ":artifact_owner",
         ":commandline_item",
         ":package_roots",

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

@@ -57,12 +57,12 @@ public SimpleArgChunk(Iterable<String> args) {
     }
 
     @Override
-    public Iterable<String> arguments() {
+    public Iterable<String> arguments(PathMapper pathMapper) {
       return args;
     }
 
     @Override
-    public int totalArgLength() {
+    public int totalArgLength(PathMapper pathMapper) {
       int total = 0;
       for (String arg : args) {
         total += arg.length() + 1;

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

@@ -105,14 +105,14 @@ ExpandedCommandLines expand(
       ParamFileInfo paramFileInfo = pair.paramFileInfo;
       ArgChunk chunk = commandLine.expand(inputMetadataProvider, pathMapper);
       if (paramFileInfo == null) {
-        arguments.addAll(chunk.arguments());
-        cmdLineLength += chunk.totalArgLength();
+        arguments.addAll(chunk.arguments(pathMapper));
+        cmdLineLength += chunk.totalArgLength(pathMapper);
       } else {
         boolean useParamFile = true;
         if (!paramFileInfo.always()) {
-          int tentativeCmdLineLength = cmdLineLength + chunk.totalArgLength();
+          int tentativeCmdLineLength = cmdLineLength + chunk.totalArgLength(pathMapper);
           if (tentativeCmdLineLength <= conservativeMaxLength) {
-            arguments.addAll(chunk.arguments());
+            arguments.addAll(chunk.arguments(pathMapper));
             cmdLineLength = tentativeCmdLineLength;
             useParamFile = false;
           }
@@ -135,16 +135,16 @@ ExpandedCommandLines expand(
             paramFiles.add(
                 new ParamFileActionInput(
                     paramFileExecPath,
-                    ParameterFile.flagsOnly(chunk.arguments()),
+                    ParameterFile.flagsOnly(chunk.arguments(pathMapper)),
                     paramFileInfo.getFileType()));
-            for (String positionalArg : ParameterFile.nonFlags(chunk.arguments())) {
+            for (String positionalArg : ParameterFile.nonFlags(chunk.arguments(pathMapper))) {
               arguments.add(positionalArg);
               cmdLineLength += positionalArg.length() + 1;
             }
           } else {
             paramFiles.add(
                 new ParamFileActionInput(
-                    paramFileExecPath, chunk.arguments(), paramFileInfo.getFileType()));
+                    paramFileExecPath, chunk.arguments(pathMapper), paramFileInfo.getFileType()));
           }
         }
       }

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

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.actions;
 
 import com.google.devtools.build.docgen.annot.DocCategory;
+import com.google.devtools.build.lib.actions.Artifact.DerivedArtifact;
 import com.google.devtools.build.lib.actions.CommandLineItem.ExceptionlessMapFn;
 import com.google.devtools.build.lib.actions.CommandLineItem.MapFn;
 import com.google.devtools.build.lib.analysis.config.CoreOptions;
@@ -105,6 +106,17 @@ default String getMappedExecPathString(ActionInput artifact) {
     return map(artifact.getExecPath()).getPathString();
   }
 
+  /**
+   * Returns the difference {@code artifact.getExecPathString().length() -
+   * getMappedExecPathString(artifact).length()}, i.e., the unmapped path length minus the mapped
+   * path length.
+   *
+   * <p>Implementations should provide a more efficient implementation that avoids allocations.
+   */
+  default int computeExecPathLengthDiff(DerivedArtifact artifact) {
+    return artifact.getExecPathString().length() - getMappedExecPathString(artifact).length();
+  }
+
   /**
    * We don't yet have a Starlark API for mapping paths in command lines. Simple Starlark calls like
    * {@code args.add(arg_name, file_path} are automatically handled. But calls that involve custom
@@ -193,6 +205,16 @@ public PathFragment map(PathFragment execPath) {
           return execPath;
         }
 
+        @Override
+        public String getMappedExecPathString(ActionInput artifact) {
+          return artifact.getExecPathString();
+        }
+
+        @Override
+        public int computeExecPathLengthDiff(DerivedArtifact artifact) {
+          return 0;
+        }
+
         @Override
         public FileRootApi mapRoot(Artifact artifact) {
           return artifact.getRoot();

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

@@ -1329,7 +1329,6 @@ java_library(
     deps = [
         ":actions/abstract_file_write_action",
         "//src/main/java/com/google/devtools/build/lib/actions",
-        "//src/main/java/com/google/devtools/build/lib/actions:arg_chunk",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
         "//src/main/java/com/google/devtools/build/lib/analysis/config:core_options",
@@ -1353,7 +1352,6 @@ java_library(
     deps = [
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:action_input_helper",
-        "//src/main/java/com/google/devtools/build/lib/actions:arg_chunk",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
         "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
@@ -1759,7 +1757,6 @@ java_library(
     srcs = ["starlark/StarlarkCustomCommandLine.java"],
     deps = [
         "//src/main/java/com/google/devtools/build/lib/actions",
-        "//src/main/java/com/google/devtools/build/lib/actions:arg_chunk",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
         "//src/main/java/com/google/devtools/build/lib/actions:fileset_output_symlink",

src/main/java/com/google/devtools/build/lib/analysis/actions/ParameterFileWriteAction.java

@@ -167,7 +167,7 @@ private static class ParamFileWriter implements DeterministicWriter {
 
     @Override
     public void writeTo(OutputStream out) throws IOException {
-      ParameterFile.writeParameterFile(out, arguments.arguments(), type);
+      ParameterFile.writeParameterFile(out, arguments.arguments(PathMapper.NOOP), type);
     }
   }
 

src/main/java/com/google/devtools/build/lib/analysis/actions/StrippingPathMapper.java

@@ -27,9 +27,9 @@
 import com.google.devtools.build.lib.actions.CommandLineItem;
 import com.google.devtools.build.lib.actions.CommandLineItem.ExceptionlessMapFn;
 import com.google.devtools.build.lib.actions.CommandLineItem.MapFn;
-import com.google.devtools.build.lib.actions.CommandLines.ParamFileActionInput;
 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.starlarkbuildapi.FileRootApi;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.HashMap;
@@ -145,7 +145,7 @@ static Optional<PathMapper> tryCreate(AbstractAction action, boolean isStarlarkA
 
   @Override
   public String getMappedExecPathString(ActionInput artifact) {
-    if (isSupportedInputType(artifact) && isOutputPath(artifact, outputRoot)) {
+    if (isSupported(artifact)) {
       return strip(artifact.getExecPath()).getPathString();
     } else {
       return artifact.getExecPathString();
@@ -157,6 +157,15 @@ public PathFragment map(PathFragment execPath) {
     return isOutputPath(execPath, outputRoot) ? strip(execPath) : execPath;
   }
 
+  @Override
+  public int computeExecPathLengthDiff(DerivedArtifact artifact) {
+    String unmappedPath = artifact.getExecPathString();
+    // bazel-out/k8-fastbuild/... is mapped to bazel-out/${FIXED_CONFIG_SEGMENT}/...
+    int firstSlash = outputRoot.getPathString().length() + 1;
+    int secondSlash = unmappedPath.indexOf('/', firstSlash + 1);
+    return (secondSlash - firstSlash) - FIXED_CONFIG_SEGMENT.length();
+  }
+
   @Override
   public ArgChunk mapCustomStarlarkArgs(ArgChunk chunk) {
     if (!isStarlarkAction) {
@@ -174,7 +183,7 @@ public ArgChunk mapCustomStarlarkArgs(ArgChunk chunk) {
 
     // TODO: b/327187486 - This materializes strings when totalArgLength() is called. Can it
     //  compute the total arg length without creating garbage strings?
-    Iterable<String> args = chunk.arguments();
+    Iterable<String> args = chunk.arguments(this);
     return new SimpleArgChunk(() -> new CustomStarlarkArgsIterator(args.iterator(), argStripper));
   }
 
@@ -204,10 +213,14 @@ public FileRootApi mapRoot(Artifact artifact) {
     return PathMapper.super.mapRoot(artifact);
   }
 
-  private boolean isSupportedInputType(ActionInput artifact) {
-    return artifact instanceof DerivedArtifact
-        || artifact instanceof ParamFileActionInput
-        || artifact instanceof BasicActionInput;
+  private boolean isSupported(ActionInput artifact) {
+    if (artifact instanceof DerivedArtifact) {
+      return true;
+    }
+    if (artifact instanceof BasicActionInput || artifact instanceof VirtualActionInput) {
+      return isOutputPath(artifact, outputRoot);
+    }
+    return false;
   }
 
   private static final class CustomStarlarkArgsIterator implements Iterator<String> {
@@ -339,12 +352,6 @@ private static boolean isPathStrippable(
    * Strips the configuration prefix from an output artifact's exec path.
    */
   private static PathFragment strip(PathFragment execPath) {
-    if (execPath.subFragment(1, 2).getPathString().equals("tmp")) {
-      return execPath
-          .subFragment(0, 2)
-          .getRelative(FIXED_CONFIG_SEGMENT)
-          .getRelative(execPath.subFragment(3));
-    }
     return execPath
         .subFragment(0, 1)
         // Keep the config segment, but replace it with a fixed string to improve cacheability while