Add support for in-memory outputs to output reuse

Closes #23422. PiperOrigin-RevId: 669296792 Change-Id: If3a2519cbf57cc82f1f733900b93705065b37e63
bazelbuild · Aug 30, 2024 · 31d6663 · 31d6663
1 parent fc057d2
commit 31d6663
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 36 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/SpawnResult.java b/src/main/java/com/google/devtools/build/lib/actions/SpawnResult.java
@@ -29,7 +29,6 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
-import java.io.InputStream;
 import java.time.Instant;
 import java.util.Locale;
 import javax.annotation.Nullable;
@@ -38,7 +37,7 @@
 @SuppressWarnings("GoodTime") // Use ints instead of Durations to improve build time (cl/505728570)
 public interface SpawnResult {
 
-  int POSIX_TIMEOUT_EXIT_CODE = /*SIGNAL_BASE=*/ 128 + /*SIGALRM=*/ 14;
+  int POSIX_TIMEOUT_EXIT_CODE = /* SIGNAL_BASE= */ 128 + /* SIGALRM= */ 14;
 
   /** The status of the attempted Spawn execution. */
   enum Status {
@@ -262,14 +261,11 @@ default String getFailureMessage() {
    * ExecutionRequirements#REMOTE_EXECUTION_INLINE_OUTPUTS}.
    */
   @Nullable
-  default InputStream getInMemoryOutput(ActionInput output) {
+  default ByteString getInMemoryOutput(ActionInput output) {
     return null;
   }
 
-  String getDetailMessage(
-      String message,
-      boolean catastrophe,
-      boolean forciblyRunRemotely);
+  String getDetailMessage(String message, boolean catastrophe, boolean forciblyRunRemotely);
 
   /** Returns a file path to the action metadata log. */
   @Nullable
@@ -434,11 +430,8 @@ public String getFailureMessage() {
 
     @Override
     public String getDetailMessage(
-        String message,
-        boolean catastrophe,
-        boolean forciblyRunRemotely) {
-      TerminationStatus status = new TerminationStatus(
-          exitCode(), status() == Status.TIMEOUT);
+        String message, boolean catastrophe, boolean forciblyRunRemotely) {
+      TerminationStatus status = new TerminationStatus(exitCode(), status() == Status.TIMEOUT);
       String reason = "(" + status.toShortString() + ")"; // e.g. "(Exit 1)"
       String explanation = Strings.isNullOrEmpty(message) ? "" : ": " + message;
 
@@ -457,17 +450,18 @@ public String getDetailMessage(
         explanation += " (Remote action was terminated due to Out of Memory.)";
       }
       if (status() != Status.TIMEOUT && forciblyRunRemotely) {
-        explanation += " Action tagged as local was forcibly run remotely and failed - it's "
-            + "possible that the action simply doesn't work remotely";
+        explanation +=
+            " Action tagged as local was forcibly run remotely and failed - it's "
+                + "possible that the action simply doesn't work remotely";
       }
       return reason + explanation;
     }
 
     @Nullable
     @Override
-    public InputStream getInMemoryOutput(ActionInput output) {
+    public ByteString getInMemoryOutput(ActionInput output) {
       if (inMemoryOutputFile != null && inMemoryOutputFile.equals(output)) {
-        return inMemoryContents.newInput();
+        return inMemoryContents;
       }
       return null;
     }
@@ -603,7 +597,7 @@ public String getFailureMessage() {
 
     @Override
     @Nullable
-    public InputStream getInMemoryOutput(ActionInput output) {
+    public ByteString getInMemoryOutput(ActionInput output) {
       return delegate.getInMemoryOutput(output);
     }
 

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/StarlarkAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/StarlarkAction.java
@@ -46,6 +46,7 @@
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.server.FailureDetails.StarlarkAction.Code;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import com.google.protobuf.ByteString;
 import java.io.BufferedReader;
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -312,9 +313,9 @@ private InputStream getUnusedInputListInputStream(
       // Note: SpawnActionContext guarantees that the first list entry exists and corresponds to the
       // executed spawn.
       Artifact unusedInputsListArtifact = unusedInputsList.get();
-      InputStream inputStream = spawnResults.get(0).getInMemoryOutput(unusedInputsListArtifact);
-      if (inputStream != null) {
-        return inputStream;
+      ByteString content = spawnResults.get(0).getInMemoryOutput(unusedInputsListArtifact);
+      if (content != null) {
+        return content.newInput();
       }
       // Fallback to reading from disk.
       try {

diff --git a/src/main/java/com/google/devtools/build/lib/includescanning/SpawnIncludeScanner.java b/src/main/java/com/google/devtools/build/lib/includescanning/SpawnIncludeScanner.java
@@ -49,6 +49,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Symlinks;
 import com.google.devtools.build.lib.vfs.SyscallCache;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Collection;
@@ -341,6 +342,7 @@ public Collection<Inclusion> extractInclusions(
    *     Otherwise "null"
    * @throws ExecException if scanning fails
    */
+  @Nullable
   private static InputStream spawnGrep(
       Artifact input,
       PathFragment outputExecPath,
@@ -402,7 +404,11 @@ private static InputStream spawnGrep(
     }
 
     SpawnResult result = results.getFirst();
-    return result.getInMemoryOutput(output);
+    ByteString includesContent = result.getInMemoryOutput(output);
+    if (includesContent != null) {
+      return includesContent.newInput();
+    }
+    return null;
   }
 
   private static void dump(ActionExecutionContext fromContext, ActionExecutionContext toContext) {

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
@@ -1458,12 +1458,20 @@ public SpawnResult waitForAndReuseOutputs(RemoteAction action, LocalExecution pr
         previousExecution.action.getSpawn().getOutputFiles().stream()
             .collect(toImmutableMap(output -> execRoot.getRelative(output.getExecPath()), o -> o));
     Map<Path, Path> realToTmpPath = new HashMap<>();
+    ByteString inMemoryOutputContent = null;
+    String inMemoryOutputPath = null;
     try {
       for (String output : action.getCommand().getOutputPathsList()) {
         String reencodedOutput = reencodeExternalToInternal(output);
         Path sourcePath =
             previousExecution.action.getRemotePathResolver().outputPathToLocalPath(reencodedOutput);
         ActionInput outputArtifact = previousOutputs.get(sourcePath);
+        Path targetPath = action.getRemotePathResolver().outputPathToLocalPath(reencodedOutput);
+        inMemoryOutputContent = previousSpawnResult.getInMemoryOutput(outputArtifact);
+        if (inMemoryOutputContent != null) {
+          inMemoryOutputPath = targetPath.relativeTo(execRoot).getPathString();
+          continue;
+        }
         Path tmpPath = tempPathGenerator.generateTempPath();
         tmpPath.getParentDirectory().createDirectoryAndParents();
         try {
@@ -1475,8 +1483,6 @@ public SpawnResult waitForAndReuseOutputs(RemoteAction action, LocalExecution pr
           } else {
             FileSystemUtils.copyFile(sourcePath, tmpPath);
           }
-
-          Path targetPath = action.getRemotePathResolver().outputPathToLocalPath(reencodedOutput);
           realToTmpPath.put(targetPath, tmpPath);
         } catch (FileNotFoundException e) {
           // The spawn this action was deduplicated against failed to create an output file. If the
@@ -1517,6 +1523,21 @@ public SpawnResult waitForAndReuseOutputs(RemoteAction action, LocalExecution pr
       throw e;
     }
 
+    if (inMemoryOutputPath != null) {
+      String finalInMemoryOutputPath = inMemoryOutputPath;
+      ByteString finalInMemoryOutputContent = inMemoryOutputContent;
+      return new SpawnResult.DelegateSpawnResult(previousSpawnResult) {
+        @Override
+        @Nullable
+        public ByteString getInMemoryOutput(ActionInput output) {
+          if (output.getExecPathString().equals(finalInMemoryOutputPath)) {
+            return finalInMemoryOutputContent;
+          }
+          return null;
+        }
+      };
+    }
+
     return previousSpawnResult;
   }
 

diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
@@ -96,6 +96,7 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.SkyframeLookupResult;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
@@ -1508,16 +1509,11 @@ e, createFailureDetail("OutErr copy failure", Code.COPY_OUT_ERR_FAILURE)),
   }
 
   @Nullable
-  private byte[] getDotDContents(SpawnResult spawnResult) throws EnvironmentalExecException {
+  private byte[] getDotDContents(SpawnResult spawnResult) {
     if (getDotdFile() != null) {
-      InputStream in = spawnResult.getInMemoryOutput(getDotdFile());
-      if (in != null) {
-        try {
-          return ByteStreams.toByteArray(in);
-        } catch (IOException e) {
-          throw new EnvironmentalExecException(
-              e, createFailureDetail("Reading in-memory .d file failed", Code.D_FILE_READ_FAILURE));
-        }
+      ByteString content = spawnResult.getInMemoryOutput(getDotdFile());
+      if (content != null) {
+        return content.toByteArray();
       }
     }
     return null;

diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java
@@ -82,6 +82,7 @@
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.view.proto.Deps;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.ExtensionRegistry;
 import java.io.IOException;
 import java.io.InputStream;
@@ -842,11 +843,11 @@ private static Deps.Dependencies readExecutorJdeps(
       Artifact outputDepsProto,
       ActionExecutionContext actionExecutionContext)
       throws IOException {
-    InputStream inMemoryOutput = spawnResult.getInMemoryOutput(outputDepsProto);
+    ByteString inMemoryOutput = spawnResult.getInMemoryOutput(outputDepsProto);
     try (InputStream inputStream =
         inMemoryOutput == null
             ? actionExecutionContext.getInputPath(outputDepsProto).getInputStream()
-            : inMemoryOutput) {
+            : inMemoryOutput.newInput()) {
       return Deps.Dependencies.parseFrom(inputStream, ExtensionRegistry.getEmptyRegistry());
     }
   }

diff --git a/src/test/java/com/google/devtools/build/lib/actions/SpawnResultTest.java b/src/test/java/com/google/devtools/build/lib/actions/SpawnResultTest.java
@@ -66,7 +66,7 @@ public void getTimeoutMessageNoTime() {
   }
 
   @Test
-  public void inMemoryContents() throws Exception {
+  public void inMemoryContents() {
     ActionInput output = ActionInputHelper.fromPath("/foo/bar");
     ByteString contents = ByteString.copyFromUtf8("hello world");
 
@@ -78,7 +78,7 @@ public void inMemoryContents() throws Exception {
             .setInMemoryOutput(output, contents)
             .build();
 
-    assertThat(ByteString.readFrom(r.getInMemoryOutput(output))).isEqualTo(contents);
+    assertThat(r.getInMemoryOutput(output)).isEqualTo(contents);
     assertThat(r.getInMemoryOutput(null)).isEqualTo(null);
     assertThat(r.getInMemoryOutput(ActionInputHelper.fromPath("/does/not/exist"))).isEqualTo(null);
   }

diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnCacheTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnCacheTest.java
@@ -93,6 +93,7 @@
 import com.google.devtools.build.lib.vfs.SyscallCache;
 import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
 import com.google.devtools.common.options.Options;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Set;
@@ -882,6 +883,53 @@ public boolean mayModifySpawnOutputsAfterExecution() {
     assertThat(secondCacheHandle.willStore()).isFalse();
   }
 
+  @Test
+  public void pathMappedActionWithInMemoryOutputIsDeduplicated() throws Exception {
+    // arrange
+    RemoteSpawnCache cache = createRemoteSpawnCache();
+
+    SimpleSpawn firstSpawn = simplePathMappedSpawn("k8-fastbuild");
+    FakeActionInputFileCache firstFakeFileCache = new FakeActionInputFileCache(execRoot);
+    firstFakeFileCache.createScratchInput(firstSpawn.getInputFiles().getSingleton(), "xyz");
+    SpawnExecutionContext firstPolicy =
+        createSpawnExecutionContext(firstSpawn, execRoot, firstFakeFileCache, outErr);
+
+    SimpleSpawn secondSpawn = simplePathMappedSpawn("k8-opt");
+    FakeActionInputFileCache secondFakeFileCache = new FakeActionInputFileCache(execRoot);
+    secondFakeFileCache.createScratchInput(secondSpawn.getInputFiles().getSingleton(), "xyz");
+    SpawnExecutionContext secondPolicy =
+        createSpawnExecutionContext(secondSpawn, execRoot, secondFakeFileCache, outErr);
+
+    RemoteExecutionService remoteExecutionService = cache.getRemoteExecutionService();
+    Mockito.doCallRealMethod().when(remoteExecutionService).waitForAndReuseOutputs(any(), any());
+    // Simulate a very slow upload to the remote cache to ensure that the second spawn is
+    // deduplicated rather than a cache hit. This is a slight hack, but also avoid introducing
+    // concurrency to this test.
+    Mockito.doNothing().when(remoteExecutionService).uploadOutputs(any(), any(), any());
+
+    // act
+    try (CacheHandle firstCacheHandle = cache.lookup(firstSpawn, firstPolicy)) {
+      firstCacheHandle.store(
+          new SpawnResult.Builder()
+              .setExitCode(0)
+              .setStatus(Status.SUCCESS)
+              .setRunnerName("test")
+              .setInMemoryOutput(
+                  firstSpawn.getOutputFiles().getFirst(), ByteString.copyFromUtf8("in-memory"))
+              .build());
+    }
+    CacheHandle secondCacheHandle = cache.lookup(secondSpawn, secondPolicy);
+
+    // assert
+    ActionInput inMemoryOutput = secondSpawn.getOutputFiles().getFirst();
+    assertThat(secondCacheHandle.hasResult()).isTrue();
+    assertThat(secondCacheHandle.getResult().getRunnerName()).isEqualTo("deduplicated");
+    assertThat(secondCacheHandle.getResult().getInMemoryOutput(inMemoryOutput).toStringUtf8())
+        .isEqualTo("in-memory");
+    assertThat(execRoot.getRelative(inMemoryOutput.getExecPath()).exists()).isFalse();
+    assertThat(secondCacheHandle.willStore()).isFalse();
+  }
+
   @Test
   public void deduplicatedActionWithNonZeroExitCodeIsACacheMiss() throws Exception {
     // arrange