Add support for in-memory outputs to output reuse

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

@@ -38,7 +38,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 {
@@ -260,16 +260,30 @@ default String getFailureMessage() {
    *
    * <p>This behavior may be triggered with {@link
    * ExecutionRequirements#REMOTE_EXECUTION_INLINE_OUTPUTS}.
+   *
+   * <p>Implementing classes should override {@link #getInMemoryOutputBytes(ActionInput)} instead.
    */
   @Nullable
   default InputStream getInMemoryOutput(ActionInput output) {
+    ByteString bytes = getInMemoryOutputBytes(output);
+    if (bytes == null) {
+      return null;
+    }
+    return bytes.newInput();
+  }
+
+  /**
+   * Returns a {@link Spawn}'s output in-memory, if supported and available.
+   *
+   * <p>This behavior may be triggered with {@link
+   * ExecutionRequirements#REMOTE_EXECUTION_INLINE_OUTPUTS}.
+   */
+  @Nullable
+  default ByteString getInMemoryOutputBytes(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 +448,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 +468,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 getInMemoryOutputBytes(ActionInput output) {
       if (inMemoryOutputFile != null && inMemoryOutputFile.equals(output)) {
-        return inMemoryContents.newInput();
+        return inMemoryContents;
       }
       return null;
     }
@@ -603,8 +615,8 @@ public String getFailureMessage() {
 
     @Override
     @Nullable
-    public InputStream getInMemoryOutput(ActionInput output) {
-      return delegate.getInMemoryOutput(output);
+    public ByteString getInMemoryOutputBytes(ActionInput output) {
+      return delegate.getInMemoryOutputBytes(output);
     }
 
     @Override

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

@@ -1422,6 +1422,8 @@ 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()) {
         Path sourcePath =
@@ -1430,6 +1432,13 @@ public SpawnResult waitForAndReuseOutputs(RemoteAction action, LocalExecution pr
                 .getRemotePathResolver()
                 .outputPathToLocalPath(encodeBytestringUtf8(output));
         ActionInput outputArtifact = previousOutputs.get(sourcePath);
+        Path targetPath =
+            action.getRemotePathResolver().outputPathToLocalPath(encodeBytestringUtf8(output));
+        inMemoryOutputContent = previousSpawnResult.getInMemoryOutputBytes(outputArtifact);
+        if (inMemoryOutputContent != null) {
+          inMemoryOutputPath = targetPath.relativeTo(execRoot).getPathString();
+          continue;
+        }
         Path tmpPath = tempPathGenerator.generateTempPath();
         tmpPath.getParentDirectory().createDirectoryAndParents();
         try {
@@ -1442,8 +1451,6 @@ public SpawnResult waitForAndReuseOutputs(RemoteAction action, LocalExecution pr
             FileSystemUtils.copyFile(sourcePath, tmpPath);
           }
 
-          Path targetPath =
-              action.getRemotePathResolver().outputPathToLocalPath(encodeBytestringUtf8(output));
           realToTmpPath.put(targetPath, tmpPath);
         } catch (FileNotFoundException e) {
           // The spawn this action was deduplicated against failed to create an output file. If the
@@ -1484,6 +1491,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 getInMemoryOutputBytes(ActionInput output) {
+          if (output.getExecPathString().equals(finalInMemoryOutputPath)) {
+            return finalInMemoryOutputContent;
+          }
+          return null;
+        }
+      };
+    }
+
     return previousSpawnResult;
   }
 

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

@@ -91,6 +91,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.SortedMap;
@@ -755,6 +756,53 @@ public void pathMappedActionIsDeduplicated() throws Exception {
     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.copyFrom("in-memory", UTF_8))
+              .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().getInMemoryOutputBytes(inMemoryOutput).toStringUtf8())
+        .isEqualTo("in-memory");
+    assertThat(execRoot.getRelative(inMemoryOutput.getExecPath()).exists()).isFalse();
+    assertThat(secondCacheHandle.willStore()).isFalse();
+  }
+
   @Test
   public void deduplicatedActionWithNonZeroExitCodeIsACacheMiss() throws Exception {
     // arrange