Reuse runfiles in tests with --reuse_sandbox_directories

Test actions are guaranteed to have a runfiles directory with the test name as part of the name. The path to the directory is unique between tests. If two tests (foo and bar) have the directory <source-root>/pkg/my_runfiles as part of their runfiles and this directory contains 1000 files, we would be symlinking the 1000 files for each test since the paths do not coincide.

To make sure we can reuse the runfiles directory we must rename the old runfiles directory for the action that was stashed to the path that is expected by the current test.

PiperOrigin-RevId: 608273594
Change-Id: Iea46fa3600d29d9a5747c93f104c0a5b23d1ea53

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

@@ -16,12 +16,15 @@
 
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.exec.TreeDeleter;
+import com.google.devtools.build.lib.sandbox.SandboxHelpers.SandboxOutputs;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Path;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.util.Collection;
+import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
@@ -35,6 +38,8 @@
 public class SandboxStash {
 
   public static final String SANDBOX_STASH_BASE = "sandbox_stash";
+  private static final String TEST_RUNNER_MNEMONIC = "TestRunner";
+  private static final String TEST_SRCDIR = "TEST_SRCDIR";
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
   /** An incrementing count of stashes to avoid filename clashes. */
@@ -53,22 +58,26 @@ public class SandboxStash {
   private final String workspaceName;
   private final Path sandboxBase;
 
+  private final Map<Path, String> stashPathToRunfilesDir = new ConcurrentHashMap<>();
+
   public SandboxStash(String workspaceName, Path sandboxBase) {
     this.workspaceName = workspaceName;
     this.sandboxBase = sandboxBase;
   }
 
-  static boolean takeStashedSandbox(Path sandboxPath, String mnemonic) {
+  static boolean takeStashedSandbox(
+      Path sandboxPath, String mnemonic, Map<String, String> environment, SandboxOutputs outputs) {
     if (instance == null) {
       return false;
     }
-    return instance.takeStashedSandboxInternal(sandboxPath, mnemonic);
+    return instance.takeStashedSandboxInternal(sandboxPath, mnemonic, environment, outputs);
   }
 
-  private boolean takeStashedSandboxInternal(Path sandboxPath, String mnemonic) {
+  private boolean takeStashedSandboxInternal(
+      Path sandboxPath, String mnemonic, Map<String, String> environment, SandboxOutputs outputs) {
     try {
       Path sandboxes = getSandboxStashDir(mnemonic, sandboxPath.getFileSystem());
-      if (sandboxes == null) {
+      if (sandboxes == null || isTestXmlGenerationOrCoverageSpawn(mnemonic, outputs)) {
         return false;
       }
       Collection<Path> stashes = sandboxes.getDirectoryEntries();
@@ -85,6 +94,14 @@ private boolean takeStashedSandboxInternal(Path sandboxPath, String mnemonic) {
           Path stashExecroot = stash.getChild("execroot");
           stashExecroot.renameTo(sandboxExecroot);
           stash.deleteTree();
+          if (isTestAction(mnemonic)) {
+            Path stashedRunfilesDir =
+                sandboxExecroot.getRelative(stashPathToRunfilesDir.get(stashExecroot));
+            Path currentRunfiles = sandboxExecroot.getRelative(getCurrentRunfilesDir(environment));
+            currentRunfiles.getParentDirectory().createDirectoryAndParents();
+            stashedRunfilesDir.renameTo(currentRunfiles);
+            stashPathToRunfilesDir.remove(stashExecroot);
+          }
           return true;
         } catch (FileNotFoundException e) {
           // Try the next one, somebody else took this one.
@@ -101,16 +118,18 @@ private boolean takeStashedSandboxInternal(Path sandboxPath, String mnemonic) {
   }
 
   /** Atomically moves the sandboxPath directory aside for later reuse. */
-  static void stashSandbox(Path path, String mnemonic) {
+  static void stashSandbox(
+      Path path, String mnemonic, Map<String, String> environment, SandboxOutputs outputs) {
     if (instance == null) {
       return;
     }
-    instance.stashSandboxInternal(path, mnemonic);
+    instance.stashSandboxInternal(path, mnemonic, environment, outputs);
   }
 
-  private void stashSandboxInternal(Path path, String mnemonic) {
+  private void stashSandboxInternal(
+      Path path, String mnemonic, Map<String, String> environment, SandboxOutputs outputs) {
     Path sandboxes = getSandboxStashDir(mnemonic, path.getFileSystem());
-    if (sandboxes == null) {
+    if (sandboxes == null || isTestXmlGenerationOrCoverageSpawn(mnemonic, outputs)) {
       return;
     }
     String stashName;
@@ -123,7 +142,11 @@ private void stashSandboxInternal(Path path, String mnemonic) {
     }
     try {
       stashPath.createDirectory();
-      path.getChild("execroot").renameTo(stashPath.getChild("execroot"));
+      Path stashPathExecroot = stashPath.getChild("execroot");
+      path.getChild("execroot").renameTo(stashPathExecroot);
+      if (isTestAction(mnemonic)) {
+        stashPathToRunfilesDir.put(stashPathExecroot, getCurrentRunfilesDir(environment));
+      }
     } catch (IOException e) {
       // Since stash names are unique, this IOException indicates some other problem with stashing,
       // so we turn it off.
@@ -244,4 +267,35 @@ static void clean(TreeDeleter treeDeleter, Path sandboxBase) {
       logger.atWarning().withCause(e).log("Failed to clean sandbox stash %s", stashDir);
     }
   }
+
+  /**
+   * Test actions are guaranteed to have a runfiles directory with the test name as part of the
+   * name. The path to the directory is unique between tests. If two tests (foo and bar) have the
+   * directory <source-root>/pkg/my_runfiles as part of their runfiles and this directory contains
+   * 1000 files, we would be symlinking the 1000 files for each test since the paths do not
+   * coincide. To make sure we can reuse the runfiles directory we must rename the old runfiles
+   * directory for the action that was stashed to the path that is expected by the current test.
+   */
+  private static boolean isTestAction(String mnemonic) {
+    return mnemonic.equals(TEST_RUNNER_MNEMONIC);
+  }
+
+  /**
+   * Test actions are split in two spawns. The first one runs the test and the second generates the
+   * XML output from the test log. We do not want the second spawn to reuse the stash because it
+   * doesn't contain the inputs needed to run the test; if it reused it, it would be expensive in
+   * two ways: it would have to clean up all the inputs, and it would destroy a valid stash that a
+   * different test could potentially use. If we are running coverage, there might be a third spawn
+   * for coverage where we apply the same reasoning.
+   *
+   * <p>We identify the second and third spawn because they have a single output.
+   */
+  private static boolean isTestXmlGenerationOrCoverageSpawn(
+      String mnemonic, SandboxOutputs outputs) {
+    return isTestAction(mnemonic) && outputs.files().size() == 1;
+  }
+
+  private static String getCurrentRunfilesDir(Map<String, String> environment) {
+    return environment.get("TEST_WORKSPACE") + "/" + environment.get(TEST_SRCDIR);
+  }
 }

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

@@ -75,7 +75,8 @@ public SymlinkedSandboxedSpawn(
   public void filterInputsAndDirsToCreate(
       Set<PathFragment> inputsToCreate, LinkedHashSet<PathFragment> dirsToCreate)
       throws IOException, InterruptedException {
-    boolean gotStash = SandboxStash.takeStashedSandbox(sandboxPath, mnemonic);
+    boolean gotStash =
+        SandboxStash.takeStashedSandbox(sandboxPath, mnemonic, getEnvironment(), outputs);
     sandboxExecRoot.createDirectoryAndParents();
     if (gotStash) {
       // When reusing an old sandbox, we do a full traversal of the parent directory of
@@ -101,7 +102,7 @@ protected void copyFile(Path source, Path target) throws IOException {
 
   @Override
   public void delete() {
-    SandboxStash.stashSandbox(sandboxPath, mnemonic);
+    SandboxStash.stashSandbox(sandboxPath, mnemonic, getEnvironment(), outputs);
     super.delete();
   }
 

src/test/shell/integration/sandboxing_test.sh

@@ -950,4 +950,89 @@ EOF
 
   bazel shutdown
 }
+
+function test_runfiles_from_tests_get_reused() {
+  mkdir pkg
+  touch pkg/file.txt
+  cat >pkg/reusing_test.bzl <<'EOF'
+def _reused_runfiles_test_impl(ctx):
+    output = ctx.actions.declare_file(ctx.label.name + ".sh")
+
+    runfiles = ctx.runfiles(files = ctx.files.file)
+    runfiles = runfiles.merge(runfiles)
+
+    test_code = """
+    #!/bin/bash
+    dir_inode_number=$(ls -di $TEST_SRCDIR | cut -f1 -d" ")
+    echo "The directory inode is $dir_inode_number"
+    file_inode_number=$(ls -i $TEST_SRCDIR/_main/pkg/file.txt | cut -f1 -d" ")
+    echo "The file inode is $file_inode_number"
+    """
+
+    ctx.actions.run_shell(
+        outputs = [output],
+        mnemonic = "myexample",
+        command = """
+        output_path={}
+        echo '{}' > $output_path
+        chmod 777 $output_path
+        """.format(output.path, test_code)
+    )
+
+    return [DefaultInfo(executable = output, runfiles = runfiles)]
+
+reused_runfiles_test = rule(
+    implementation = _reused_runfiles_test_impl,
+    test = True,
+    attrs = {
+        "file" : attr.label(allow_files=True,default="//pkg:file.txt"),
+    }
+)
+EOF
+
+  cat >pkg/BUILD <<'EOF'
+load(":reusing_test.bzl", "reused_runfiles_test")
+reused_runfiles_test(
+    name = "a",
+)
+reused_runfiles_test(
+    name = "b",
+)
+EOF
+
+  test_output="reuse_test_output.txt"
+  if is_bazel; then
+    bazel coverage --test_output=streamed \
+      --experimental_split_coverage_postprocessing=1 \
+      --experimental_fetch_all_coverage_outputs //pkg:a > ${test_output} \
+      || fail "Expected build to succeed"
+  else
+    bazel test --test_output=streamed //pkg:a > ${test_output} \
+      || fail "Expected build to succeed"
+  fi
+  dir_inode_a=$(awk '/The directory inode is/ {print $5}' ${test_output})
+  file_inode_a=$(awk '/The file inode is/ {print $5}' ${test_output})
+
+  if is_bazel; then
+    bazel coverage --test_output=streamed //pkg:b \
+      --experimental_split_coverage_postprocessing=1 \
+      --experimental_fetch_all_coverage_outputs > ${test_output} \
+      || fail "Expected build to succeed"
+  else
+    bazel test --test_output=streamed //pkg:b > ${test_output} \
+      || fail "Expected build to succeed"
+  fi
+  dir_inode_b=$(awk '/The directory inode is/ {print $5}' ${test_output})
+  file_inode_b=$(awk '/The file inode is/ {print $5}' ${test_output})
+
+  [[ ${dir_inode_a} == ${dir_inode_b} ]] \
+    || fail "Test //pkg:b didn't reuse runfiles directory"
+  [[ ${file_inode_a} == ${file_inode_b} ]] \
+    || fail "Test //pkg:b didn't reuse runfiles file"
+}
+
+function is_bazel() {
+  [ $TEST_WORKSPACE == "_main" ]
+}
+
 run_suite "sandboxing"