Keep runfiles tree IDs in memory for tools and multiple test attempts

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

@@ -120,6 +120,11 @@ public Artifact getRepoMappingManifestForLogging() {
     public boolean isLegacyExternalRunfiles() {
       return wrapped.isLegacyExternalRunfiles();
     }
+
+    @Override
+    public boolean isMappingCached() {
+      return wrapped.isMappingCached();
+    }
   }
 
   /**

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

@@ -75,4 +75,7 @@ public interface RunfilesTree {
 
   /** Whether this runfiles tree materializes external runfiles also at their legacy locations. */
   boolean isLegacyExternalRunfiles();
+
+  /** Whether {@link #getMapping()} is cached due to potential reuse within a single build. */
+  boolean isMappingCached();
 }

src/main/java/com/google/devtools/build/lib/analysis/RunfilesSupport.java

@@ -211,6 +211,11 @@ public boolean isLegacyExternalRunfiles() {
       return runfiles.isLegacyExternalRunfiles();
     }
 
+    @Override
+    public boolean isMappingCached() {
+      return cachedMapping != null;
+    }
+
     @Override
     public RunfileSymlinksMode getSymlinksMode() {
       return runfileSymlinksMode;

src/main/java/com/google/devtools/build/lib/exec/CompactSpawnLogContext.java

@@ -326,7 +326,8 @@ private int logInputs(
         additionalDirectoryIds.build(),
         inputMetadataProvider,
         fileSystem,
-        /* shared= */ false);
+        /* shared= */ false,
+        spawn.getMnemonic());
   }
 
   /**
@@ -343,7 +344,8 @@ private int logTools(
         ImmutableList.of(),
         inputMetadataProvider,
         fileSystem,
-        /* shared= */ true);
+        /* shared= */ true,
+        spawn.getMnemonic());
   }
 
   /**
@@ -353,6 +355,7 @@ private int logTools(
    * @param additionalDirectoryIds the entry IDs of additional {@link ExecLogEntry.Directory}
    *     entries to include as direct members
    * @param shared whether this nested set is likely to be a transitive member of other sets
+   * @param mnemonic the mnemonic of the spawn
    * @return the entry ID of the {@link ExecLogEntry.InputSet} describing the nested set, or 0 if
    *     the nested set is empty.
    */
@@ -361,7 +364,8 @@ private int logInputSet(
       Collection<Integer> additionalDirectoryIds,
       InputMetadataProvider inputMetadataProvider,
       FileSystem fileSystem,
-      boolean shared)
+      boolean shared,
+      String mnemonic)
       throws IOException, InterruptedException {
     if (set.isEmpty() && additionalDirectoryIds.isEmpty()) {
       return 0;
@@ -381,7 +385,8 @@ private int logInputSet(
                     /* additionalDirectoryIds= */ ImmutableList.of(),
                     inputMetadataProvider,
                     fileSystem,
-                    /* shared= */ true));
+                    /* shared= */ true,
+                    mnemonic));
           }
 
           for (ActionInput input : set.getLeaves()) {
@@ -393,10 +398,11 @@ private int logInputSet(
                       runfilesTree,
                       inputMetadataProvider,
                       fileSystem,
-                      // If the nested set containing the runfiles tree isn't shared (i.e., it
-                      // contains inputs, not tools), the runfiles are also likely not shared. This
-                      // avoids storing the runfiles tree of a test.
-                      shared));
+                      // Runfiles in non-test spawns are from tools and thus potentially reused
+                      // between spawns. Runfiles of tests cannot be shared unless the test runs
+                      // multiple times in the same build. In this case, the runfiles tree caches
+                      // its mapping.
+                      !"TestRunner".equals(mnemonic) || runfilesTree.isMappingCached()));
               continue;
             }
 
@@ -561,7 +567,9 @@ private int logRunfilesTree(
                   fileSystem,
                   // The runfiles tree itself is shared, but the nested set is unique to the tree as
                   // it contains the executable.
-                  /* shared= */ false));
+                  /* shared= */ false,
+                  // The mnemonic is only used for logging runfiles trees and they are not nested.
+                  /* mnemonic= */ ""));
           builder.setSymlinksId(
               logSymlinkEntries(
                   runfilesTree.getSymlinksForLogging(), inputMetadataProvider, fileSystem));

src/main/protobuf/spawn.proto

@@ -373,7 +373,7 @@ message ExecLogEntry {
     // Entry ID of the set of inputs. Unset means empty.
     uint32 input_set_id = 4;
 
-    // Entry ID of the set of tool inputs. Unset means empty.
+    // Entry ID of the subset of inputs that are tools. Unset means empty.
     uint32 tool_set_id = 5;
 
     // The set of outputs.

src/test/java/com/google/devtools/build/lib/analysis/util/FakeRunfilesTree.java

@@ -122,4 +122,9 @@ public Artifact getRepoMappingManifestForLogging() {
   public boolean isLegacyExternalRunfiles() {
     return runfiles.isLegacyExternalRunfiles();
   }
+
+  @Override
+  public boolean isMappingCached() {
+    return false;
+  }
 }

src/test/java/com/google/devtools/build/lib/exec/CompactSpawnLogContextTest.java

@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.exec;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.testutil.TestConstants.PRODUCT_NAME;
 
 import com.github.luben.zstd.ZstdInputStream;
 import com.google.common.collect.ImmutableList;
@@ -22,6 +23,8 @@
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.BuildConfigurationEvent;
+import com.google.devtools.build.lib.actions.RunfilesTree;
+import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
 import com.google.devtools.build.lib.analysis.actions.SymlinkAction;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos;
@@ -35,6 +38,7 @@
 import com.google.devtools.build.lib.testutil.TestConstants;
 import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.SyscallCache;
 import com.google.devtools.common.options.Options;
 import com.google.testing.junit.testparameterinjector.TestParameter;
@@ -130,17 +134,8 @@ public void testSymlinkAction() throws IOException, InterruptedException {
 
     SpawnLogContext context = createSpawnLogContext();
     context.logSymlinkAction(symlinkAction);
-    context.close();
-
-    var entries = new ArrayList<Protos.ExecLogEntry>();
-    try (InputStream in = logPath.getInputStream();
-        ZstdInputStream zstdIn = new ZstdInputStream(in)) {
-      Protos.ExecLogEntry entry;
-      while ((entry = Protos.ExecLogEntry.parseDelimitedFrom(zstdIn)) != null) {
-        entries.add(entry);
-      }
-    }
 
+    var entries = closeAndReadCompactLog(context);
     assertThat(entries)
         .containsExactly(
             Protos.ExecLogEntry.newBuilder()
@@ -160,6 +155,79 @@ public void testSymlinkAction() throws IOException, InterruptedException {
                 .build());
   }
 
+  @Test
+  public void testRunfilesTreeReusedForTool() throws Exception {
+    Artifact tool = ActionsTestUtil.createArtifact(rootDir, "data.txt");
+    writeFile(tool, "abc");
+    Artifact toolRunfilesMiddleman = ActionsTestUtil.createArtifact(middlemanDir, "runfiles");
+
+    PathFragment runfilesRoot = outputDir.getExecPath().getRelative("foo.runfiles");
+    RunfilesTree runfilesTree = createRunfilesTree(runfilesRoot, tool);
+
+    Artifact firstInput = ActionsTestUtil.createArtifact(rootDir, "first_input");
+    writeFile(firstInput, "def");
+    Artifact secondInput = ActionsTestUtil.createArtifact(rootDir, "second_input");
+    writeFile(secondInput, "ghi");
+
+    Spawn firstSpawn =
+        defaultSpawnBuilder()
+            .withTool(toolRunfilesMiddleman)
+            .withInputs(firstInput, toolRunfilesMiddleman)
+            .build();
+    Spawn secondSpawn =
+        defaultSpawnBuilder()
+            .withTool(toolRunfilesMiddleman)
+            .withInputs(secondInput, toolRunfilesMiddleman)
+            .build();
+
+    SpawnLogContext context = createSpawnLogContext();
+
+    context.logSpawn(
+        firstSpawn,
+        createInputMetadataProvider(toolRunfilesMiddleman, runfilesTree, firstInput),
+        createInputMap(runfilesTree, firstInput),
+        fs,
+        defaultTimeout(),
+        defaultSpawnResult());
+    context.logSpawn(
+        secondSpawn,
+        createInputMetadataProvider(toolRunfilesMiddleman, runfilesTree, secondInput),
+        createInputMap(runfilesTree, secondInput),
+        fs,
+        defaultTimeout(),
+        defaultSpawnResult());
+
+    var entries = closeAndReadCompactLog(context);
+    assertThat(entries.stream().filter(Protos.ExecLogEntry::hasRunfilesTree)).hasSize(1);
+
+    closeAndAssertLog(
+        context,
+        defaultSpawnExecBuilder()
+            .addInputs(
+                File.newBuilder()
+                    .setPath(PRODUCT_NAME + "-out/k8-fastbuild/bin/foo.runfiles/_main/data.txt")
+                    .setDigest(getDigest("abc"))
+                    .setIsTool(true))
+            .addInputs(
+                File.newBuilder()
+                    .setPath("first_input")
+                    .setDigest(getDigest("def"))
+                    .setIsTool(false))
+            .build(),
+        defaultSpawnExecBuilder()
+            .addInputs(
+                File.newBuilder()
+                    .setPath(PRODUCT_NAME + "-out/k8-fastbuild/bin/foo.runfiles/_main/data.txt")
+                    .setDigest(getDigest("abc"))
+                    .setIsTool(true))
+            .addInputs(
+                File.newBuilder()
+                    .setPath("second_input")
+                    .setDigest(getDigest("ghi"))
+                    .setIsTool(false))
+            .build());
+  }
+
   @Override
   protected SpawnLogContext createSpawnLogContext(ImmutableMap<String, String> platformProperties)
       throws IOException, InterruptedException {
@@ -178,7 +246,7 @@ protected SpawnLogContext createSpawnLogContext(ImmutableMap<String, String> pla
 
   @Override
   protected void closeAndAssertLog(SpawnLogContext context, SpawnExec... expected)
-      throws IOException, InterruptedException {
+      throws IOException {
     context.close();
 
     ArrayList<SpawnExec> actual = new ArrayList<>();
@@ -192,4 +260,19 @@ protected void closeAndAssertLog(SpawnLogContext context, SpawnExec... expected)
 
     assertThat(actual).containsExactlyElementsIn(expected).inOrder();
   }
+
+  private ImmutableList<Protos.ExecLogEntry> closeAndReadCompactLog(SpawnLogContext context)
+      throws IOException {
+    context.close();
+
+    ImmutableList.Builder<Protos.ExecLogEntry> entries = ImmutableList.builder();
+    try (InputStream in = logPath.getInputStream();
+        ZstdInputStream zstdIn = new ZstdInputStream(in)) {
+      Protos.ExecLogEntry entry;
+      while ((entry = Protos.ExecLogEntry.parseDelimitedFrom(zstdIn)) != null) {
+        entries.add(entry);
+      }
+    }
+    return entries.build();
+  }
 }

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

@@ -2718,6 +2718,11 @@ public Artifact getRepoMappingManifestForLogging() {
       public boolean isLegacyExternalRunfiles() {
         return false;
       }
+
+      @Override
+      public boolean isMappingCached() {
+        return false;
+      }
     };
   }