Remove treeArtifactContents from OutputStore.

treeArtifactContents acts as a temporary cache prior to the full tree artifact being constructed and is optional - if not populated, ActionMetadataHandler reads tree artifact contents from disk. In fact, even if it is populated, disk reads are performed to verify expected contents.

Callers can use injectRemoteDirectory to establish known tree artifact contents instead. In this case, we assume that the contents are correct, so there is no need to check for a match on disk.

PiperOrigin-RevId: 311241850
diff --git a/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataHandler.java b/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataHandler.java
index d339f5a..6e06891 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataHandler.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataHandler.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.actions.cache;
 
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
@@ -34,7 +35,7 @@
   void setDigestForVirtualArtifact(Artifact artifact, byte[] digest);
 
   /** Retrieves the artifacts inside the TreeArtifact, without injecting its digest. */
-  Iterable<TreeFileArtifact> getExpandedOutputs(Artifact artifact);
+  ImmutableSet<TreeFileArtifact> getExpandedOutputs(Artifact artifact);
 
   /**
    * Returns true iff artifact was intentionally omitted (not saved).
diff --git a/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataInjector.java b/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataInjector.java
index ee95136..afbba6a 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataInjector.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/cache/MetadataInjector.java
@@ -15,10 +15,10 @@
 
 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.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.vfs.FileStatus;
-import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Map;
 
 /** Supports metadata injection of action outputs into skyframe. */
@@ -39,12 +39,11 @@
   /**
    * Inject the metadata of a tree artifact whose contents are stored remotely.
    *
-   * @param output an output directory.
-   * @param children the metadata of the files stored in the directory. The paths must be relative
-   *     to the path of {@code output}.
+   * @param output an output directory
+   * @param children the metadata of the files stored in the directory
    */
   void injectRemoteDirectory(
-      Artifact.SpecialArtifact output, Map<PathFragment, RemoteFileArtifactValue> children);
+      SpecialArtifact output, Map<TreeFileArtifact, RemoteFileArtifactValue> children);
 
   /**
    * Marks an {@link Artifact} as intentionally omitted.
@@ -55,12 +54,6 @@
   void markOmitted(ActionInput output);
 
   /**
-   * Registers the given output as contents of a TreeArtifact, without injecting its digest. Prefer
-   * {@link #injectDigest} when the digest is available.
-   */
-  void addExpandedTreeOutput(TreeFileArtifact output);
-
-  /**
    * Injects provided digest into the metadata handler, simultaneously caching lstat() data as well.
    */
   void injectDigest(ActionInput output, FileStatus statNoFollow, byte[] digest);
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
index 157d5d4..d03ccbe 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteCache.java
@@ -39,7 +39,10 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.ActionInputHelper;
 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.FileArtifactValue.RemoteFileArtifactValue;
@@ -611,20 +614,21 @@
             "Symlinks in action outputs are not yet supported by "
                 + "--experimental_remote_download_outputs=minimal");
       }
-      ImmutableMap.Builder<PathFragment, RemoteFileArtifactValue> childMetadata =
-          ImmutableMap.builder();
+      SpecialArtifact parent = (SpecialArtifact) output;
+      ImmutableMap.Builder<TreeFileArtifact, RemoteFileArtifactValue> childMetadata =
+          ImmutableMap.builderWithExpectedSize(directory.files.size());
       for (FileMetadata file : directory.files()) {
-        PathFragment p = file.path().relativeTo(output.getPath());
-        RemoteFileArtifactValue r =
+        TreeFileArtifact child =
+            ActionInputHelper.treeFileArtifact(parent, file.path().relativeTo(parent.getPath()));
+        RemoteFileArtifactValue value =
             new RemoteFileArtifactValue(
                 DigestUtil.toBinaryDigest(file.digest()),
                 file.digest().getSizeBytes(),
-                /* locationIndex= */ 1,
+                /*locationIndex=*/ 1,
                 actionId);
-        childMetadata.put(p, r);
+        childMetadata.put(child, value);
       }
-      metadataInjector.injectRemoteDirectory(
-          (Artifact.SpecialArtifact) output, childMetadata.build());
+      metadataInjector.injectRemoteDirectory(parent, childMetadata.build());
     } else {
       FileMetadata outputMetadata = metadata.file(execRoot.getRelative(output.getExecPathString()));
       if (outputMetadata == null) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java
index 25d6718..bc5a13d 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java
@@ -246,7 +246,6 @@
       // Calling code depends on this particular exception.
       throw new FileNotFoundException(artifact + " not found");
     }
-    // Fallthrough: the artifact must be a non-tree, non-middleman output artifact.
 
     // Don't store metadata for output artifacts that are not declared outputs of the action.
     if (!isKnownOutput(artifact)) {
@@ -264,6 +263,18 @@
     if (fileMetadata != null) {
       return metadataFromValue(fileMetadata);
     }
+    // This artifact was not injected directly to the store, but it may have been injected as part
+    // of a tree artifact.
+    // TODO(jhorvitz): We can skip this for action template expansion artifacts.
+    if (artifact.hasParent()) {
+      TreeArtifactValue tree = store.getTreeArtifactData(artifact.getParent());
+      if (tree != null) {
+        fileMetadata = tree.getChildValues().get(artifact);
+        if (fileMetadata != null) {
+          return metadataFromValue(fileMetadata);
+        }
+      }
+    }
 
     // No existing metadata; this can happen if the output metadata is not injected after a spawn
     // is executed. SkyframeActionExecutor.checkOutputs calls this method for every output file of
@@ -367,48 +378,7 @@
       }
     }
 
-    Set<TreeFileArtifact> registeredContents = store.getTreeArtifactContents(artifact);
-    if (registeredContents != null) {
-      // Check that our registered outputs matches on-disk outputs. Only perform this check
-      // when contents were explicitly registered.
-      // TODO(bazel-team): Provide a way for actions to register empty TreeArtifacts.
-
-      // By the time we're constructing TreeArtifactValues, use of the metadata handler
-      // should be single threaded and there should be no race condition.
-      // The current design of ActionMetadataHandler makes this hard to enforce.
-      Set<PathFragment> paths =
-          TreeArtifactValue.explodeDirectory(artifactPathResolver.toPath(artifact));
-      Set<TreeFileArtifact> diskFiles = ActionInputHelper.asTreeFileArtifacts(artifact, paths);
-      if (!diskFiles.equals(registeredContents)) {
-        // There might be more than one error here. We first look for missing output files.
-        Set<TreeFileArtifact> missingFiles = Sets.difference(registeredContents, diskFiles);
-        if (!missingFiles.isEmpty()) {
-          // Don't throw IOException--getMetadataMaybe() eats them.
-          // TODO(bazel-team): Report this error in a better way when called by checkOutputs()
-          // Currently it's hard to report this error without refactoring, since checkOutputs()
-          // likes to substitute its own error messages upon catching IOException, and falls
-          // through to unrecoverable error behavior on any other exception.
-          throw new IOException(
-              "Output file "
-                  + missingFiles.iterator().next()
-                  + " was registered, but not present on disk");
-        }
-
-        Set<TreeFileArtifact> extraFiles = Sets.difference(diskFiles, registeredContents);
-        // extraFiles cannot be empty
-        throw new IOException(
-            "File "
-                + extraFiles.iterator().next().getParentRelativePath()
-                + ", present in TreeArtifact "
-                + artifact
-                + ", was not registered");
-      }
-
-      value = constructTreeArtifactValue(registeredContents);
-    } else {
-      value = constructTreeArtifactValueFromFilesystem(artifact);
-    }
-
+    value = constructTreeArtifactValueFromFilesystem(artifact);
     store.putTreeArtifactData(artifact, value);
     return value;
   }
@@ -462,27 +432,13 @@
 
     Set<PathFragment> paths =
         TreeArtifactValue.explodeDirectory(artifactPathResolver.toPath(artifact));
-    // If you're reading tree artifacts from disk while tree artifact contents are being injected,
-    // something has gone terribly wrong.
-    Object previousContents = store.getTreeArtifactContents(artifact);
-    Preconditions.checkState(
-        previousContents == null,
-        "Race condition while constructing TreeArtifactValue: %s, %s",
-        artifact,
-        previousContents);
     return constructTreeArtifactValue(ActionInputHelper.asTreeFileArtifacts(artifact, paths));
   }
 
   @Override
-  public void addExpandedTreeOutput(TreeFileArtifact output) {
-    Preconditions.checkState(executionMode.get());
-    store.addTreeArtifactContents(output.getParent(), output);
-  }
-
-  @Override
-  public Iterable<TreeFileArtifact> getExpandedOutputs(Artifact artifact) {
-    Set<TreeFileArtifact> contents = store.getTreeArtifactContents(artifact);
-    return contents != null ? ImmutableSet.copyOf(contents) : ImmutableSet.of();
+  public ImmutableSet<TreeFileArtifact> getExpandedOutputs(Artifact artifact) {
+    TreeArtifactValue treeArtifact = store.getTreeArtifactData(artifact);
+    return treeArtifact != null ? treeArtifact.getChildren() : ImmutableSet.of();
   }
 
   @Override
@@ -557,22 +513,13 @@
 
   @Override
   public void injectRemoteDirectory(
-      SpecialArtifact output, Map<PathFragment, RemoteFileArtifactValue> children) {
+      SpecialArtifact output, Map<TreeFileArtifact, RemoteFileArtifactValue> children) {
     Preconditions.checkArgument(
-        isKnownOutput(output), output + " is not a declared output of this action");
+        isKnownOutput(output), "%s is not a declared output of this action", output);
     Preconditions.checkArgument(output.isTreeArtifact(), "output must be a tree artifact");
     Preconditions.checkState(
-        executionMode.get(), "Tried to inject %s outside of execution.", output);
-
-    ImmutableMap.Builder<TreeFileArtifact, FileArtifactValue> childFileValues =
-        ImmutableMap.builder();
-    for (Map.Entry<PathFragment, RemoteFileArtifactValue> child : children.entrySet()) {
-      childFileValues.put(
-          ActionInputHelper.treeFileArtifact(output, child.getKey()), child.getValue());
-    }
-
-    TreeArtifactValue treeArtifactValue = TreeArtifactValue.create(childFileValues.build());
-    store.putTreeArtifactData(output, treeArtifactValue);
+        executionMode.get(), "Tried to inject %s outside of execution", output);
+    store.putTreeArtifactData(output, TreeArtifactValue.create(children));
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/MinimalOutputStore.java b/src/main/java/com/google/devtools/build/lib/skyframe/MinimalOutputStore.java
index d12823b..1d9fcdc 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/MinimalOutputStore.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/MinimalOutputStore.java
@@ -14,7 +14,6 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.devtools.build.lib.actions.Artifact;
-import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 
 /**
@@ -31,7 +30,4 @@
       super.putArtifactData(artifact, value);
     }
   }
-
-  @Override
-  void addTreeArtifactContents(Artifact artifact, TreeFileArtifact contents) {}
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java b/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java
index 914b8b5..d69b309 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/OutputStore.java
@@ -13,11 +13,9 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.Artifact;
-import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import java.util.Set;
@@ -28,10 +26,9 @@
 /**
  * Storage layer for data associated with outputs of an action.
  *
- * <p>Data is mainly stored in three maps, {@link #artifactData}, {@link #treeArtifactData} and
- * {@link #treeArtifactContents}, all of which are keyed on an {@link Artifact}. For each of these
- * maps, this class exposes standard methods such as {@code get}, {@code put}, {@code add}, and
- * {@code getAll}.
+ * <p>Data is mainly stored in two maps, {@link #artifactData} and {@link #treeArtifactData}, both
+ * of which are keyed on an {@link Artifact}. For each of these maps, this class exposes standard
+ * methods such as {@code get}, {@code put}, {@code add}, and {@code getAll}.
  *
  * <p>This implementation aggressively stores all data. Subclasses may override mutating methods to
  * avoid storing unnecessary data.
@@ -44,9 +41,6 @@
   private final ConcurrentMap<Artifact, TreeArtifactValue> treeArtifactData =
       new ConcurrentHashMap<>();
 
-  private final ConcurrentMap<Artifact, Set<TreeFileArtifact>> treeArtifactContents =
-      new ConcurrentHashMap<>();
-
   private final Set<Artifact> injectedFiles = Sets.newConcurrentHashSet();
 
   @Nullable
@@ -79,22 +73,6 @@
     return ImmutableMap.copyOf(treeArtifactData);
   }
 
-  /**
-   * Returns a set of the given tree artifact's contents.
-   *
-   * <p>If the return value is {@code null}, this means nothing was injected, and the output
-   * TreeArtifact is to have its values read from disk instead.
-   */
-  @Nullable
-  final Set<TreeFileArtifact> getTreeArtifactContents(Artifact artifact) {
-    return treeArtifactContents.get(artifact);
-  }
-
-  void addTreeArtifactContents(Artifact artifact, TreeFileArtifact contents) {
-    Preconditions.checkArgument(artifact.isTreeArtifact(), artifact);
-    treeArtifactContents.computeIfAbsent(artifact, a -> Sets.newConcurrentHashSet()).add(contents);
-  }
-
   void injectRemoteFile(
       Artifact output, byte[] digest, long size, int locationIndex, String actionId) {
     injectOutputData(
@@ -116,7 +94,6 @@
   final void clear() {
     artifactData.clear();
     treeArtifactData.clear();
-    treeArtifactContents.clear();
     injectedFiles.clear();
   }
 
@@ -124,7 +101,6 @@
   final void remove(Artifact artifact) {
     artifactData.remove(artifact);
     treeArtifactData.remove(artifact);
-    treeArtifactContents.remove(artifact);
     injectedFiles.remove(artifact);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TreeArtifactValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/TreeArtifactValue.java
index 5b74906..2e3ca1c 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TreeArtifactValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TreeArtifactValue.java
@@ -69,14 +69,15 @@
    * Returns a TreeArtifactValue out of the given Artifact-relative path fragments and their
    * corresponding FileArtifactValues.
    */
-  static TreeArtifactValue create(Map<TreeFileArtifact, FileArtifactValue> childFileValues) {
+  static TreeArtifactValue create(
+      Map<TreeFileArtifact, ? extends FileArtifactValue> childFileValues) {
     if (childFileValues.isEmpty()) {
       return EMPTY;
     }
     Map<String, FileArtifactValue> digestBuilder =
         Maps.newHashMapWithExpectedSize(childFileValues.size());
     boolean remote = true;
-    for (Map.Entry<TreeFileArtifact, FileArtifactValue> e : childFileValues.entrySet()) {
+    for (Map.Entry<TreeFileArtifact, ? extends FileArtifactValue> e : childFileValues.entrySet()) {
       FileArtifactValue value = e.getValue();
       // TODO(buchgr): Enforce that all children in a tree artifact are either remote or local
       // once b/70354083 is fixed.
@@ -109,7 +110,7 @@
     return digest.clone();
   }
 
-  public Iterable<TreeFileArtifact> getChildren() {
+  public ImmutableSet<TreeFileArtifact> getChildren() {
     return childData.keySet();
   }
 
@@ -165,7 +166,7 @@
         }
 
         @Override
-        public Iterable<TreeFileArtifact> getChildren() {
+        public ImmutableSet<TreeFileArtifact> getChildren() {
           throw new UnsupportedOperationException();
         }