Inject all metadata at once after the action is executed.

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

@@ -16,6 +16,8 @@
 package com.google.devtools.build.lib.remote;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.Streams.stream;
 
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputMap;
@@ -56,25 +58,29 @@ public class RemoteActionFileSystem extends DelegateFileSystem {
   private final PathFragment execRoot;
   private final PathFragment outputBase;
   private final ActionInputMap inputArtifactData;
+  private final Map<PathFragment, Artifact> outputMapping;
   private final RemoteActionInputFetcher inputFetcher;
   private final Map<PathFragment, RemoteFileArtifactValue> injectedMetadata = new HashMap<>();
+  private final Map<PathFragment, TreeArtifactValue> injectedTreeMetadata = new HashMap<>();
 
-  @Nullable private MetadataInjector metadataInjector = null;
+  @Nullable
+  private MetadataInjector metadataInjector = null;
 
-  RemoteActionFileSystem(
-      FileSystem localDelegate,
-      PathFragment execRootFragment,
-      String relativeOutputPath,
-      ActionInputMap inputArtifactData,
-      RemoteActionInputFetcher inputFetcher) {
+  RemoteActionFileSystem(FileSystem localDelegate, PathFragment execRootFragment,
+      String relativeOutputPath, ActionInputMap inputArtifactData,
+      Iterable<Artifact> outputArtifacts, RemoteActionInputFetcher inputFetcher) {
     super(localDelegate);
     this.execRoot = checkNotNull(execRootFragment, "execRootFragment");
     this.outputBase = execRoot.getRelative(checkNotNull(relativeOutputPath, "relativeOutputPath"));
     this.inputArtifactData = checkNotNull(inputArtifactData, "inputArtifactData");
+    this.outputMapping = stream(outputArtifacts).collect(
+        toImmutableMap(Artifact::getExecPath, a -> a));
     this.inputFetcher = checkNotNull(inputFetcher, "inputFetcher");
   }
 
-  /** Returns true if {@code path} is a file that's stored remotely. */
+  /**
+   * Returns true if {@code path} is a file that's stored remotely.
+   */
   boolean isRemote(Path path) {
     return getRemoteInputMetadata(path.asFragment()) != null;
   }
@@ -84,27 +90,39 @@ public void updateContext(MetadataInjector metadataInjector) {
   }
 
   void injectTree(SpecialArtifact tree, TreeArtifactValue metadata) {
-    checkNotNull(metadataInjector, "metadataInject is null");
-
-    for (Map.Entry<TreeFileArtifact, FileArtifactValue> entry :
-        metadata.getChildValues().entrySet()) {
+    for (Map.Entry<TreeFileArtifact, FileArtifactValue> entry : metadata.getChildValues()
+        .entrySet()) {
       FileArtifactValue childMetadata = entry.getValue();
       if (childMetadata instanceof RemoteFileArtifactValue) {
         TreeFileArtifact child = entry.getKey();
         injectedMetadata.put(child.getExecPath(), (RemoteFileArtifactValue) childMetadata);
       }
     }
 
-    metadataInjector.injectTree(tree, metadata);
+    injectedTreeMetadata.put(tree.getExecPath(), metadata);
   }
 
   void injectFile(ActionInput file, RemoteFileArtifactValue metadata) {
-    checkNotNull(metadataInjector, "metadataInject is null");
-
     injectedMetadata.put(file.getExecPath(), metadata);
+  }
 
-    if (file instanceof Artifact) {
-      metadataInjector.injectFile((Artifact) file, metadata);
+  void flush() {
+    checkNotNull(metadataInjector, "metadataInjector is null");
+
+    for (Map.Entry<PathFragment, Artifact> entry : outputMapping.entrySet()) {
+      PathFragment execPath = entry.getKey();
+      Artifact output = entry.getValue();
+      if (output.isTreeArtifact()) {
+        TreeArtifactValue metadata = injectedTreeMetadata.get(execPath);
+        if (metadata != null) {
+          metadataInjector.injectTree((SpecialArtifact) output, metadata);
+        }
+      } else {
+        RemoteFileArtifactValue metadata = injectedMetadata.get(execPath);
+        if (metadata != null) {
+          metadataInjector.injectFile(output, metadata);
+        }
+      }
     }
   }
 
@@ -381,8 +399,8 @@ private void downloadFileIfRemote(PathFragment path) throws IOException {
         inputFetcher.downloadFile(delegateFs.getPath(path), m);
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
-        throw new IOException(
-            String.format("Received interrupt while fetching file '%s'", path), e);
+        throw new IOException(String.format("Received interrupt while fetching file '%s'", path),
+            e);
       }
     }
   }

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

@@ -70,6 +70,7 @@ public FileSystem createActionFileSystem(
         execRootFragment,
         relativeOutputPath,
         inputArtifactData,
+        outputArtifacts,
         actionInputFetcher);
   }
 
@@ -98,6 +99,11 @@ public void finalizeBuild(boolean buildSuccessful) {
     // Intentionally left empty.
   }
 
+  @Override
+  public void flushActionFileSystem(FileSystem actionFileSystem) {
+    ((RemoteActionFileSystem) actionFileSystem).flush();
+  }
+
   @Override
   public void finalizeAction(Action action, MetadataHandler metadataHandler) {
     // Intentionally left empty.
@@ -142,7 +148,8 @@ public ArtifactPathResolver createPathResolverForArtifactValues(
       Map<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
     FileSystem remoteFileSystem =
         new RemoteActionFileSystem(
-            fileSystem, execRoot, relativeOutputPath, actionInputMap, actionInputFetcher);
+            fileSystem, execRoot, relativeOutputPath, actionInputMap, ImmutableList.of(),
+            actionInputFetcher);
     return ArtifactPathResolver.createPathResolver(remoteFileSystem, fileSystem.getPath(execRoot));
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java

@@ -1200,6 +1200,8 @@ private ActionExecutionValue actuallyCompleteAction(
         Preconditions.checkState(action.inputsDiscovered(),
             "Action %s successfully executed, but inputs still not known", action);
 
+        flushActionFileSystem(actionExecutionContext.getActionFileSystem(), outputService);
+
         if (!checkOutputs(
             action,
             metadataHandler,
@@ -1511,6 +1513,13 @@ private static LostInputsCheck lostInputsCheck(
         : () -> outputService.checkActionFileSystemForLostInputs(actionFileSystem, action);
   }
 
+  private static void flushActionFileSystem(@Nullable FileSystem actionFileSystem,
+      @Nullable OutputService outputService) {
+    if (outputService != null && actionFileSystem != null) {
+      outputService.flushActionFileSystem(actionFileSystem);
+    }
+  }
+
   /**
    * Validates that all action outputs were created or intentionally omitted. This can result in
    * chmod calls on the output files; see {@link ActionMetadataHandler}.

src/main/java/com/google/devtools/build/lib/vfs/OutputService.java

@@ -192,14 +192,23 @@ default void updateActionFileSystemContext(
       FileSystem actionFileSystem,
       Environment env,
       MetadataInjector injector,
-      ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {}
+      ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
+  }
 
   /**
    * Checks the filesystem returned by {@link #createActionFileSystem} for errors attributable to
    * lost inputs.
    */
   default void checkActionFileSystemForLostInputs(FileSystem actionFileSystem, Action action)
-      throws LostInputsActionExecutionException {}
+      throws LostInputsActionExecutionException {
+  }
+
+  /**
+   * Flush the internal state of filesystem returned by {@link #createActionFileSystem} after action
+   * execution, before skyframe checking the action outputs.
+   */
+  default void flushActionFileSystem(FileSystem actionFileSystem) {
+  }
 
   default boolean supportsPathResolverForArtifactValues() {
     return false;

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

@@ -26,6 +26,7 @@
 import build.bazel.remote.execution.v2.Digest;
 import com.google.bytestream.ByteStreamProto.WriteRequest;
 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.eventbus.EventBus;
@@ -363,6 +364,7 @@ public void remoteFileShouldNotBeUploaded_actionFs() throws Exception {
             execRoot.asFragment(),
             outputRoot.getRoot().asPath().relativeTo(execRoot).getPathString(),
             outputs,
+            ImmutableList.of(artifact),
             actionInputFetcher);
     Path remotePath = remoteFs.getPath(artifact.getPath().getPathString());
     assertThat(remotePath.getFileSystem()).isEqualTo(remoteFs);

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

@@ -20,6 +20,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.hash.HashCode;
 import com.google.common.util.concurrent.Futures;
 import com.google.devtools.build.lib.actions.ActionInputMap;
@@ -153,15 +154,23 @@ public void testDeleteLocalFile() throws Exception {
   }
 
   private RemoteActionFileSystem newRemoteActionFileSystem(ActionInputMap inputs) {
+    return newRemoteActionFileSystem(inputs, ImmutableList.of());
+  }
+
+  private RemoteActionFileSystem newRemoteActionFileSystem(ActionInputMap inputs,
+      Iterable<Artifact> outputs) {
     return new RemoteActionFileSystem(
         fs,
         execRoot.asFragment(),
         outputRoot.getRoot().asPath().relativeTo(execRoot).getPathString(),
         inputs,
+        outputs,
         inputFetcher);
   }
 
-  /** Returns a remote artifact and puts its metadata into the action input map. */
+  /**
+   * Returns a remote artifact and puts its metadata into the action input map.
+   */
   private Artifact createRemoteArtifact(
       String pathFragment, String contents, ActionInputMap inputs) {
     Path p = outputRoot.getRoot().asPath().getRelative(pathFragment);