Wire up internal-inline-outputs requirement with RE API

In-memory outputs marked with the `internal-inline-outputs` execution requirement are now added to the `inline_output_files` hint field of the remote execution and remote cache API, potentially saving a round trip on files that are unconditionally read by Bazel client after spawn execution.

Support for output file inlining has been added to the remote worker implementation.

Fixes bazelbuild#8421.

Closes bazelbuild#23428.

PiperOrigin-RevId: 673938116
Change-Id: Ieaea53502c43f9fc90b99f72dd4cdb4711239dd3

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

@@ -71,6 +71,7 @@
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
@@ -270,14 +271,18 @@ public ListenableFuture<String> getAuthority() {
 
   @Override
   public ListenableFuture<ActionResult> downloadActionResult(
-      RemoteActionExecutionContext context, ActionKey actionKey, boolean inlineOutErr) {
+      RemoteActionExecutionContext context,
+      ActionKey actionKey,
+      boolean inlineOutErr,
+      Set<String> inlineOutputFiles) {
     GetActionResultRequest request =
         GetActionResultRequest.newBuilder()
             .setInstanceName(options.remoteInstanceName)
             .setDigestFunction(digestUtil.getDigestFunction())
             .setActionDigest(actionKey.getDigest())
             .setInlineStderr(inlineOutErr)
             .setInlineStdout(inlineOutErr)
+            .addAllInlineOutputFiles(inlineOutputFiles)
             .build();
     return Utils.refreshIfUnauthenticatedAsync(
         () ->

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

@@ -61,6 +61,7 @@
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.regex.Matcher;
@@ -146,7 +147,10 @@ public static CachedActionResult disk(ActionResult actionResult) {
   }
 
   public CachedActionResult downloadActionResult(
-      RemoteActionExecutionContext context, ActionKey actionKey, boolean inlineOutErr)
+      RemoteActionExecutionContext context,
+      ActionKey actionKey,
+      boolean inlineOutErr,
+      Set<String> inlineOutputFiles)
       throws IOException, InterruptedException {
     var spawnExecutionContext = context.getSpawnExecutionContext();
 
@@ -180,7 +184,8 @@ public CachedActionResult downloadActionResult(
                     spawnExecutionContext.report(SPAWN_CHECKING_REMOTE_CACHE_EVENT);
                   }
                   return Futures.transform(
-                      downloadActionResultFromRemote(context, actionKey, inlineOutErr),
+                      downloadActionResultFromRemote(
+                          context, actionKey, inlineOutErr, inlineOutputFiles),
                       CachedActionResult::remote,
                       directExecutor());
                 } else {
@@ -194,10 +199,13 @@ public CachedActionResult downloadActionResult(
   }
 
   private ListenableFuture<ActionResult> downloadActionResultFromRemote(
-      RemoteActionExecutionContext context, ActionKey actionKey, boolean inlineOutErr) {
+      RemoteActionExecutionContext context,
+      ActionKey actionKey,
+      boolean inlineOutErr,
+      Set<String> inlineOutputFiles) {
     checkState(remoteCacheClient != null && context.getReadCachePolicy().allowRemoteCache());
     return Futures.transformAsync(
-        remoteCacheClient.downloadActionResult(context, actionKey, inlineOutErr),
+        remoteCacheClient.downloadActionResult(context, actionKey, inlineOutErr, inlineOutputFiles),
         (actionResult) -> {
           if (actionResult == null) {
             return immediateFuture(null);

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

@@ -768,11 +768,19 @@ public RemoteActionResult lookupCache(RemoteAction action)
         action.getRemoteActionExecutionContext().getReadCachePolicy().allowAnyCache(),
         "spawn doesn't accept cached result");
 
+    ImmutableSet<String> inlineOutputFiles = ImmutableSet.of();
+    PathFragment inMemoryOutputPath = getInMemoryOutputPath(action.getSpawn());
+    if (inMemoryOutputPath != null) {
+      inlineOutputFiles =
+          ImmutableSet.of(action.getRemotePathResolver().localPathToOutputPath(inMemoryOutputPath));
+    }
+
     CachedActionResult cachedActionResult =
         remoteCache.downloadActionResult(
             action.getRemoteActionExecutionContext(),
             action.getActionKey(),
-            /* inlineOutErr= */ false);
+            /* inlineOutErr= */ false,
+            inlineOutputFiles);
 
     if (cachedActionResult == null) {
       return null;
@@ -958,11 +966,13 @@ public static class FileMetadata {
       private final Path path;
       private final Digest digest;
       private final boolean isExecutable;
+      private final ByteString contents;
 
-      private FileMetadata(Path path, Digest digest, boolean isExecutable) {
+      private FileMetadata(Path path, Digest digest, boolean isExecutable, ByteString contents) {
         this.path = path;
         this.digest = digest;
         this.isExecutable = isExecutable;
+        this.contents = contents;
       }
 
       public Path path() {
@@ -976,6 +986,10 @@ public Digest digest() {
       public boolean isExecutable() {
         return isExecutable;
       }
+
+      public ByteString content() {
+        return contents;
+      }
     }
 
     static class DirectoryMetadata {
@@ -1039,7 +1053,10 @@ private static DirectoryMetadata parseDirectory(
     for (FileNode file : dir.getFilesList()) {
       filesBuilder.add(
           new FileMetadata(
-              parent.getRelative(file.getName()), file.getDigest(), file.getIsExecutable()));
+              parent.getRelative(file.getName()),
+              file.getDigest(),
+              file.getIsExecutable(),
+              ByteString.EMPTY));
     }
 
     ImmutableList.Builder<SymlinkMetadata> symlinksBuilder = ImmutableList.builder();
@@ -1105,7 +1122,11 @@ static ActionResultMetadata parseActionResultMetadata(
               reencodeExternalToInternal(outputFile.getPath()));
       files.put(
           localPath,
-          new FileMetadata(localPath, outputFile.getDigest(), outputFile.getIsExecutable()));
+          new FileMetadata(
+              localPath,
+              outputFile.getDigest(),
+              outputFile.getIsExecutable(),
+              outputFile.getContents()));
     }
 
     var symlinkMap = new HashMap<Path, SymlinkMetadata>();
@@ -1181,7 +1202,7 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
     var expireAtEpochMilli = Instant.now().plus(remoteOptions.remoteCacheTtl).toEpochMilli();
 
     ActionInput inMemoryOutput = null;
-    AtomicReference<byte[]> inMemoryOutputData = new AtomicReference<>(null);
+    AtomicReference<ByteString> inMemoryOutputData = new AtomicReference<>(null);
     PathFragment inMemoryOutputPath = getInMemoryOutputPath(action.getSpawn());
     if (inMemoryOutputPath != null) {
       for (ActionInput output : action.getSpawn().getOutputFiles()) {
@@ -1227,15 +1248,25 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
         }
 
         if (isInMemoryOutputFile) {
-          downloadsBuilder.add(
-              transform(
-                  remoteCache.downloadBlob(
-                      context, inMemoryOutputPath.getPathString(), file.digest()),
-                  data -> {
-                    inMemoryOutputData.set(data);
-                    return null;
-                  },
-                  directExecutor()));
+          if (file.contents.isEmpty()) {
+            // As the contents field doesn't have presence information, we use the digest size to
+            // distinguish between an empty file and one that wasn't inlined.
+            if (file.digest.getSizeBytes() == 0) {
+              inMemoryOutputData.set(ByteString.EMPTY);
+            } else {
+              downloadsBuilder.add(
+                  transform(
+                      remoteCache.downloadBlob(
+                          context, inMemoryOutputPath.getPathString(), file.digest()),
+                      data -> {
+                        inMemoryOutputData.set(ByteString.copyFrom(data));
+                        return null;
+                      },
+                      directExecutor()));
+            }
+          } else {
+            inMemoryOutputData.set(file.contents);
+          }
         }
       }
     }
@@ -1369,7 +1400,7 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
     }
 
     if (inMemoryOutput != null && inMemoryOutputData.get() != null) {
-      return new InMemoryOutput(inMemoryOutput, ByteString.copyFrom(inMemoryOutputData.get()));
+      return new InMemoryOutput(inMemoryOutput, inMemoryOutputData.get());
     }
 
     return null;
@@ -1811,6 +1842,12 @@ public RemoteActionResult executeRemotely(
     if (remoteOptions.remoteExecutionPriority != 0) {
       requestBuilder.getExecutionPolicyBuilder().setPriority(remoteOptions.remoteExecutionPriority);
     }
+    PathFragment inMemoryOutputPath = getInMemoryOutputPath(action.getSpawn());
+    if (inMemoryOutputPath != null) {
+      requestBuilder.addInlineOutputFiles(
+          reencodeInternalToExternal(
+              action.getRemotePathResolver().localPathToOutputPath(inMemoryOutputPath)));
+    }
 
     ExecuteRequest request = requestBuilder.build();
 

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

@@ -25,6 +25,7 @@
 import build.bazel.remote.execution.v2.RequestMetadata;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.analysis.platform.PlatformUtils;
@@ -153,7 +154,11 @@ public ExecutionResult execute(
     try (SilentCloseable c =
         Profiler.instance().profile(ProfilerTask.REMOTE_CACHE_CHECK, "check cache hit")) {
       cachedActionResult =
-          remoteCache.downloadActionResult(context, actionKey, /* inlineOutErr= */ true);
+          remoteCache.downloadActionResult(
+              context,
+              actionKey,
+              /* inlineOutErr= */ true,
+              /* inlineOutputFiles= */ ImmutableSet.of());
     }
     ActionResult actionResult = null;
     if (cachedActionResult != null) {

src/main/java/com/google/devtools/build/lib/remote/common/RemoteCacheClient.java

@@ -24,6 +24,7 @@
 import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.Set;
 
 /**
  * An interface for a remote caching protocol.
@@ -76,11 +77,16 @@ public int hashCode() {
    * @param actionKey The digest of the {@link Action} that generated the action result.
    * @param inlineOutErr A hint to the server to inline the stdout and stderr in the {@code
    *     ActionResult} message.
+   * @param inlineOutputFiles A hint to the server to inline the specified output files in the
+   *     {@code ActionResult} message.
    * @return A Future representing pending download of an action result. If an action result for
    *     {@code actionKey} cannot be found the result of the Future is {@code null}.
    */
   ListenableFuture<ActionResult> downloadActionResult(
-      RemoteActionExecutionContext context, ActionKey actionKey, boolean inlineOutErr);
+      RemoteActionExecutionContext context,
+      ActionKey actionKey,
+      boolean inlineOutErr,
+      Set<String> inlineOutputFiles);
 
   /**
    * Uploads an action result for the {@code actionKey}.

src/main/java/com/google/devtools/build/lib/remote/disk/DiskCacheClient.java

@@ -122,9 +122,8 @@ public DiskCacheClient(
    * deliberately use the mtime because the atime is more likely to be externally modified and may
    * be unavailable on some filesystems.
    *
-   * <p>Prefer calling {@link #downloadBlob} or {@link #downloadActionResult} instead, which will
-   * automatically update the mtime. This method should only be called by the remote worker
-   * implementation.
+   * <p>Prefer calling {@link #downloadBlob} instead, which will automatically update the mtime.
+   * This method should only be called by the remote worker implementation.
    *
    * @throws IOException if an I/O error other than a missing file occurs.
    */

src/main/java/com/google/devtools/build/lib/remote/http/HttpCacheClient.java

@@ -87,6 +87,7 @@
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
@@ -617,7 +618,10 @@ public ListenableFuture<String> getAuthority() {
 
   @Override
   public ListenableFuture<ActionResult> downloadActionResult(
-      RemoteActionExecutionContext context, ActionKey actionKey, boolean inlineOutErr) {
+      RemoteActionExecutionContext context,
+      ActionKey actionKey,
+      boolean inlineOutErr,
+      Set<String> inlineOutputFiles) {
     return retrier.executeAsync(
         () ->
             Utils.downloadAsActionResult(

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

@@ -52,6 +52,7 @@
 import com.google.bytestream.ByteStreamProto.WriteResponse;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Maps;
 import com.google.common.eventbus.EventBus;
@@ -779,10 +780,12 @@ public void getActionResult(
     GrpcCacheClient client = newClient(remoteOptions);
     RemoteCache remoteCache =
         new RemoteCache(client, /* diskCacheClient= */ null, remoteOptions, DIGEST_UTIL);
-    remoteCache.downloadActionResult(
-        context,
-        DIGEST_UTIL.asActionKey(DIGEST_UTIL.computeAsUtf8("key")),
-        /* inlineOutErr= */ false);
+    var unused =
+        remoteCache.downloadActionResult(
+            context,
+            DIGEST_UTIL.asActionKey(DIGEST_UTIL.computeAsUtf8("key")),
+            /* inlineOutErr= */ false,
+            /* inlineOutputFiles= */ ImmutableSet.of());
   }
 
   @Test
@@ -1099,7 +1102,11 @@ public void getActionResult(
         });
     assertThat(
             getFromFuture(
-                client.downloadActionResult(context, actionKey, /* inlineOutErr= */ false)))
+                client.downloadActionResult(
+                    context,
+                    actionKey,
+                    /* inlineOutErr= */ false,
+                    /* inlineOutputFiles= */ ImmutableSet.of())))
         .isNull();
   }
 

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

@@ -174,7 +174,8 @@ public void downloadActionResult_reportsSpawnCheckingCacheEvent() throws Excepti
         remoteCache.downloadActionResult(
             remoteActionExecutionContext,
             digestUtil.asActionKey(digestUtil.computeAsUtf8("key")),
-            /* inlineOutErr= */ false);
+            /* inlineOutErr= */ false,
+            /* inlineOutputFiles= */ ImmutableSet.of());
 
     verify(remoteActionExecutionContext.getSpawnExecutionContext())
         .report(SpawnCheckingCacheEvent.create("remote-cache"));