Use an in-memory filesystem to track injected remote files

src/main/java/com/google/devtools/build/lib/remote/BUILD

@@ -102,6 +102,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:output_service",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
+        "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs",
         "//src/main/java/com/google/devtools/build/skyframe",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:failure_details_java_proto",

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

@@ -20,27 +20,31 @@
 import static com.google.common.collect.Streams.stream;
 
 import com.google.common.collect.ImmutableMap;
-import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputMap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.actions.cache.MetadataInjector;
+import com.google.devtools.build.lib.clock.Clock;
 import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
 import com.google.devtools.build.lib.vfs.DelegateFileSystem;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.Dirent;
 import com.google.devtools.build.lib.vfs.FileStatus;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.vfs.inmemoryfs.FileInfo;
+import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryContentInfo;
+import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.nio.channels.ReadableByteChannel;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -61,8 +65,7 @@ public class RemoteActionFileSystem extends DelegateFileSystem {
   private final ActionInputMap inputArtifactData;
   private final ImmutableMap<PathFragment, Artifact> outputMapping;
   private final RemoteActionInputFetcher inputFetcher;
-  private final Map<PathFragment, RemoteFileArtifactValue> injectedMetadata = new HashMap<>();
-  private final Map<PathFragment, TreeArtifactValue> injectedTreeMetadata = new HashMap<>();
+  private final RemoteInMemoryFileSystem remoteOutputTree;
 
   @Nullable private MetadataInjector metadataInjector = null;
 
@@ -80,9 +83,12 @@ public class RemoteActionFileSystem extends DelegateFileSystem {
     this.outputMapping =
         stream(outputArtifacts).collect(toImmutableMap(Artifact::getExecPath, a -> a));
     this.inputFetcher = checkNotNull(inputFetcher, "inputFetcher");
+    this.remoteOutputTree = new RemoteInMemoryFileSystem(getDigestFunction());
   }
 
-  /** 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;
   }
@@ -91,43 +97,58 @@ public void updateContext(MetadataInjector metadataInjector) {
     this.metadataInjector = metadataInjector;
   }
 
-  void injectTree(SpecialArtifact tree, TreeArtifactValue metadata) {
-    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);
-      }
+  void injectRemoteFile(PathFragment path, byte[] digest, long size, String actionId)
+      throws IOException {
+    if (!path.startsWith(outputBase)) {
+      return;
     }
-
-    injectedTreeMetadata.put(tree.getExecPath(), metadata);
-  }
-
-  void injectFile(ActionInput file, RemoteFileArtifactValue metadata) {
-    injectedMetadata.put(file.getExecPath(), metadata);
+    remoteOutputTree.injectRemoteFile(path, digest, size, actionId);
   }
 
-  void flush() {
+  void flush() throws IOException {
     checkNotNull(metadataInjector, "metadataInjector is null");
 
     for (Map.Entry<PathFragment, Artifact> entry : outputMapping.entrySet()) {
       PathFragment execPath = entry.getKey();
+      PathFragment path = execRoot.getRelative(execPath);
       Artifact output = entry.getValue();
       if (output.isTreeArtifact()) {
-        TreeArtifactValue metadata = injectedTreeMetadata.get(execPath);
-        if (metadata != null) {
-          metadataInjector.injectTree((SpecialArtifact) output, metadata);
+        SpecialArtifact parent = (SpecialArtifact) output;
+        TreeArtifactValue.Builder tree = TreeArtifactValue.newBuilder(parent);
+        if (remoteOutputTree.exists(path)) {
+          TreeArtifactValue.visitTree(remoteOutputTree.getPath(path),
+              ((parentRelativePath, type) -> {
+                if (type == Dirent.Type.DIRECTORY) {
+                  return;
+                }
+                RemoteFileInfo remoteFile = remoteOutputTree.getRemoteFileInfo(
+                    path.getRelative(parentRelativePath), /* followSymlinks= */ true);
+                if (remoteFile != null) {
+                  TreeFileArtifact child = TreeFileArtifact.createTreeOutput(parent,
+                      parentRelativePath);
+                  tree.putChild(child, createRemoteMetadata(remoteFile));
+                }
+              }));
         }
+
+        // TODO: Check directory content on the local fs to support mixed tree.
+
+        metadataInjector.injectTree(parent, tree.build());
       } else {
-        RemoteFileArtifactValue metadata = injectedMetadata.get(execPath);
-        if (metadata != null) {
-          metadataInjector.injectFile(output, metadata);
+        RemoteFileInfo remoteFile = remoteOutputTree.getRemoteFileInfo(path, /* followSymlinks= */
+            true);
+        if (remoteFile != null) {
+          metadataInjector.injectFile(output, createRemoteMetadata(remoteFile));
         }
       }
     }
   }
 
+  private RemoteFileArtifactValue createRemoteMetadata(RemoteFileInfo remoteFile) {
+    return RemoteFileArtifactValue.create(remoteFile.getFastDigest(),
+        remoteFile.getSize(), /* locationIndex= */ 1, remoteFile.getActionId());
+  }
+
   @Override
   public String getFileSystemType(PathFragment path) {
     return "remoteActionFS";
@@ -139,7 +160,7 @@ protected boolean delete(PathFragment path) throws IOException {
     if (m == null) {
       return super.delete(path);
     }
-    return true;
+    return remoteOutputTree.getPath(path).delete();
   }
 
   @Override
@@ -160,6 +181,7 @@ public void setLastModifiedTime(PathFragment path, long newTime) throws IOExcept
     if (m == null) {
       super.setLastModifiedTime(path, newTime);
     }
+    remoteOutputTree.setLastModifiedTime(path, newTime);
   }
 
   @Override
@@ -391,7 +413,13 @@ private RemoteFileArtifactValue getRemoteInputMetadata(PathFragment path) {
       return (RemoteFileArtifactValue) m;
     }
 
-    return injectedMetadata.get(execPath);
+    RemoteFileInfo remoteFile = remoteOutputTree.getRemoteFileInfo(path, /* followSymlinks= */
+        true);
+    if (remoteFile != null) {
+      return createRemoteMetadata(remoteFile);
+    }
+
+    return null;
   }
 
   private void downloadFileIfRemote(PathFragment path) throws IOException {
@@ -438,4 +466,93 @@ protected void createHardLink(PathFragment linkPath, PathFragment originalPath)
       throws IOException {
     super.createHardLink(linkPath, originalPath);
   }
+
+  static class RemoteInMemoryFileSystem extends InMemoryFileSystem {
+
+    public RemoteInMemoryFileSystem(DigestHashFunction hashFunction) {
+      super(hashFunction);
+    }
+
+
+    @Override
+    protected synchronized OutputStream getOutputStream(PathFragment path, boolean append)
+        throws IOException {
+      // To get an output stream from remote file, we need to first stage it.
+      throw new IllegalStateException("Shouldn't be called directly");
+    }
+
+    @Override
+    protected FileInfo newFile(Clock clock, PathFragment path) {
+      return new RemoteFileInfo(clock);
+    }
+
+    void injectRemoteFile(PathFragment path, byte[] digest, long size, String actionId)
+        throws IOException {
+      createDirectoryAndParents(path.getParentDirectory());
+      InMemoryContentInfo node = getOrCreateWritableInode(path);
+      // If a node was already existed and is not a remote file node (i.e. directory or symlink node
+      // ), throw an error.
+      if (!(node instanceof RemoteFileInfo)) {
+        throw new IOException("Could not inject into " + node);
+      }
+
+      RemoteFileInfo remoteFileInfo = (RemoteFileInfo) node;
+      remoteFileInfo.set(digest, size, actionId);
+    }
+
+    @Nullable
+    RemoteFileInfo getRemoteFileInfo(PathFragment path, boolean followSymlinks) {
+      InMemoryContentInfo node = inodeStatErrno(path, followSymlinks).inode();
+      if (!(node instanceof RemoteFileInfo)) {
+        return null;
+      }
+      return (RemoteFileInfo) node;
+    }
+  }
+
+  static class RemoteFileInfo extends FileInfo {
+
+    private byte[] digest;
+    private long size;
+    private String actionId;
+
+    RemoteFileInfo(Clock clock) {
+      super(clock);
+    }
+
+    private void set(byte[] digest, long size, String actionId) {
+      this.digest = digest;
+      this.size = size;
+      this.actionId = actionId;
+    }
+
+    @Override
+    public OutputStream getOutputStream(boolean append) throws IOException {
+      throw new IllegalStateException("Shouldn't be called directly");
+    }
+
+    @Override
+    public InputStream getInputStream() throws IOException {
+      throw new IllegalStateException("Shouldn't be called directly");
+    }
+
+    @Override
+    public byte[] getxattr(String name) throws IOException {
+      throw new IllegalStateException("Shouldn't be called directly");
+    }
+
+    @Override
+    public byte[] getFastDigest() {
+      return digest;
+    }
+
+    @Override
+    public long getSize() {
+      return size;
+    }
+
+    public String getActionId() {
+      return actionId;
+    }
+  }
 }

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

@@ -66,12 +66,9 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
-import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
-import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionRequirements;
-import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.actions.ForbiddenActionInputException;
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.Spawn;
@@ -109,7 +106,6 @@
 import com.google.devtools.build.lib.remote.util.Utils;
 import com.google.devtools.build.lib.remote.util.Utils.InMemoryOutput;
 import com.google.devtools.build.lib.server.FailureDetails.RemoteExecution;
-import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -751,55 +747,33 @@ private void createSymlinks(Iterable<SymlinkMetadata> symlinks) throws IOExcepti
     }
   }
 
-  private void injectRemoteArtifact(
-      RemoteAction action, ActionInput output, ActionResultMetadata metadata) throws IOException {
+  private void injectRemoteArtifacts(RemoteAction action, ActionResultMetadata metadata)
+      throws IOException {
     FileSystem actionFileSystem = action.getSpawnExecutionContext().getActionFileSystem();
     checkState(actionFileSystem instanceof RemoteActionFileSystem);
 
     RemoteActionExecutionContext context = action.getRemoteActionExecutionContext();
     RemoteActionFileSystem remoteActionFileSystem = (RemoteActionFileSystem) actionFileSystem;
 
-    Path path = remotePathResolver.outputPathToLocalPath(output);
-    if (output instanceof Artifact && ((Artifact) output).isTreeArtifact()) {
-      DirectoryMetadata directory = metadata.directory(path);
-      if (directory == null) {
-        // A declared output wasn't created. It might have been an optional output and if not
-        // SkyFrame will make sure to fail.
-        return;
-      }
+    for (Map.Entry<Path, DirectoryMetadata> entry : metadata.directories()) {
+      DirectoryMetadata directory = entry.getValue();
       if (!directory.symlinks().isEmpty()) {
         throw new IOException(
             "Symlinks in action outputs are not yet supported by "
                 + "--experimental_remote_download_outputs=minimal");
       }
-      SpecialArtifact parent = (SpecialArtifact) output;
-      TreeArtifactValue.Builder tree = TreeArtifactValue.newBuilder(parent);
+
       for (FileMetadata file : directory.files()) {
-        TreeFileArtifact child =
-            TreeFileArtifact.createTreeOutput(parent, file.path().relativeTo(parent.getPath()));
-        RemoteFileArtifactValue value =
-            RemoteFileArtifactValue.create(
-                DigestUtil.toBinaryDigest(file.digest()),
-                file.digest().getSizeBytes(),
-                /*locationIndex=*/ 1,
-                context.getRequestMetadata().getActionId());
-        tree.putChild(child, value);
+        remoteActionFileSystem.injectRemoteFile(file.path().asFragment(),
+            DigestUtil.toBinaryDigest(file.digest()), file.digest().getSizeBytes(),
+            context.getRequestMetadata().getActionId());
       }
-      remoteActionFileSystem.injectTree(parent, tree.build());
-    } else {
-      FileMetadata outputMetadata = metadata.file(path);
-      if (outputMetadata == null) {
-        // A declared output wasn't created. It might have been an optional output and if not
-        // SkyFrame will make sure to fail.
-        return;
-      }
-      remoteActionFileSystem.injectFile(
-          output,
-          RemoteFileArtifactValue.create(
-              DigestUtil.toBinaryDigest(outputMetadata.digest()),
-              outputMetadata.digest().getSizeBytes(),
-              /*locationIndex=*/ 1,
-              context.getRequestMetadata().getActionId()));
+    }
+
+    for (FileMetadata file : metadata.files()) {
+      remoteActionFileSystem.injectRemoteFile(file.path().asFragment(),
+          DigestUtil.toBinaryDigest(file.digest()), file.digest().getSizeBytes(),
+          context.getRequestMetadata().getActionId());
     }
   }
 
@@ -1154,9 +1128,10 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
           inMemoryOutputDigest = m.digest();
           inMemoryOutput = output;
         }
-        injectRemoteArtifact(action, output, metadata);
       }
 
+      injectRemoteArtifacts(action, metadata);
+
       try (SilentCloseable c = Profiler.instance().profile("Remote.downloadInMemoryOutput")) {
         if (inMemoryOutput != null) {
           ListenableFuture<byte[]> inMemoryOutputDownload =

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

@@ -34,6 +34,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.skyframe.SkyFunction.Environment;
+import java.io.IOException;
 import java.util.Map;
 import java.util.UUID;
 import javax.annotation.Nullable;
@@ -100,7 +101,7 @@ public void finalizeBuild(boolean buildSuccessful) {
   }
 
   @Override
-  public void flushActionFileSystem(FileSystem actionFileSystem) {
+  public void flushActionFileSystem(FileSystem actionFileSystem) throws IOException {
     ((RemoteActionFileSystem) actionFileSystem).flush();
   }