Refactor lost output ownership calculation.

`ImportantOutputHandler` takes on the responsibility of calculating owners in addition to just lost artifacts. Disparate ownership calculations in `CompletionContext` and `SpawnInputExpander` are removed.

PiperOrigin-RevId: 655645302
Change-Id: Ib1528a2fea6b721954ecbab1303bf1bb0d2272d5

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

@@ -147,7 +147,6 @@ java_library(
         "cache/VirtualActionInput.java",
     ],
     deps = [
-        ":action_input_helper",
         ":action_lookup_data",
         ":action_lookup_key",
         ":arg_chunk",

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

@@ -30,7 +30,6 @@
 import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
-import java.util.Collection;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -128,36 +127,6 @@ public FileArtifactValue getFileArtifactValue(Artifact artifact) {
     return importantInputMap.getInputMetadata(artifact);
   }
 
-  /** Returns owner mappings for artifact expansions contained in {@code inputsOfInterest}. */
-  public ActionInputDepOwners getDepOwners(Collection<ActionInput> inputsOfInterest) {
-    ActionInputDepOwnerMap depOwners = new ActionInputDepOwnerMap(inputsOfInterest);
-    treeArtifacts.forEach(
-        (tree, treeValue) -> {
-          for (Artifact child : treeValue.getChildren()) {
-            depOwners.addOwner(child, tree);
-          }
-        });
-    if (expandFilesets) {
-      for (Artifact fileset : filesets.keySet()) {
-        visitFileset(
-            fileset,
-            new ArtifactReceiver() {
-              @Override
-              public void accept(Artifact artifact) {
-                throw new AssertionError(artifact);
-              }
-
-              @Override
-              public void acceptFilesetMapping(
-                  Artifact fileset, PathFragment relName, Path targetFile) {
-                depOwners.addOwner(ActionInputHelper.fromPath(targetFile.asFragment()), fileset);
-              }
-            });
-      }
-    }
-    return depOwners;
-  }
-
   /** Visits the expansion of the given artifacts. */
   public void visitArtifacts(Iterable<Artifact> artifacts, ArtifactReceiver receiver) {
     for (Artifact artifact : artifacts) {

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

@@ -13,6 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.actions;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.skyframe.DetailedException;
@@ -36,12 +38,11 @@ public interface ImportantOutputHandler extends ActionContext {
    * @param expander used to expand {@linkplain Artifact#isDirectory directory artifacts} in {@code
    *     outputs}
    * @param metadataProvider provides metadata for artifacts in {@code outputs} and their expansions
-   * @return a map from digest to output for any artifacts that need to be regenerated via action
-   *     rewinding
+   * @return any artifacts that need to be regenerated via action rewinding
    * @throws ImportantOutputException for an issue processing the outputs, not including lost
-   *     outputs which are reported in the returned map
+   *     outputs which are reported in the returned {@link LostArtifacts}
    */
-  ImmutableMap<String, ActionInput> processOutputsAndGetLostArtifacts(
+  LostArtifacts processOutputsAndGetLostArtifacts(
       Iterable<Artifact> outputs, ArtifactExpander expander, InputMetadataProvider metadataProvider)
       throws ImportantOutputException, InterruptedException;
 
@@ -59,12 +60,11 @@ ImmutableMap<String, ActionInput> processOutputsAndGetLostArtifacts(
    *     runfiles}
    * @param metadataProvider provides metadata for artifacts in {@code runfiles} and their
    *     expansions
-   * @return a map from digest to output for any artifacts that need to be regenerated via action
-   *     rewinding
+   * @return any artifacts that need to be regenerated via action rewinding
    * @throws ImportantOutputException for an issue processing the runfiles, not including lost
-   *     outputs which are reported in the returned map
+   *     outputs which are reported in the returned {@link LostArtifacts}
    */
-  ImmutableMap<String, ActionInput> processRunfilesAndGetLostArtifacts(
+  LostArtifacts processRunfilesAndGetLostArtifacts(
       PathFragment runfilesDir,
       Map<PathFragment, Artifact> runfiles,
       ArtifactExpander expander,
@@ -92,6 +92,21 @@ void processTestOutputs(List<Path> testOutputs)
    */
   Duration LOG_THRESHOLD = Duration.ofMillis(100);
 
+  /**
+   * Represents artifacts that need to be regenerated via action rewinding, along with their owners.
+   */
+  record LostArtifacts(ImmutableMap<String, ActionInput> byDigest, ActionInputDepOwners owners) {
+
+    public LostArtifacts(ImmutableMap<String, ActionInput> byDigest, ActionInputDepOwners owners) {
+      this.byDigest = checkNotNull(byDigest);
+      this.owners = checkNotNull(owners);
+    }
+
+    public boolean isEmpty() {
+      return byDigest.isEmpty();
+    }
+  }
+
   /** Represents an exception encountered during processing of important outputs. */
   final class ImportantOutputException extends Exception implements DetailedException {
     private final FailureDetail failureDetail;

src/main/java/com/google/devtools/build/lib/exec/SpawnInputExpander.java

@@ -46,7 +46,6 @@
 import java.util.Map;
 import java.util.SortedMap;
 import java.util.TreeMap;
-import javax.annotation.Nullable;
 
 /**
  * A helper class for spawn strategies to turn runfiles suppliers into input mappings. This class
@@ -55,16 +54,6 @@
  */
 public final class SpawnInputExpander {
 
-  /** Accepts mappings from exec path to {@link ActionInput}. */
-  public interface InputSink {
-    void acceptMapping(PathFragment execPath, ActionInput input, @Nullable Artifact owner);
-
-    /** Adapts a {@link Map} to a {@link InputSink} which disregards owners. */
-    static InputSink fromMap(Map<PathFragment, ActionInput> map) {
-      return (execPath, input, owner) -> map.put(execPath, input);
-    }
-  }
-
   private final Path execRoot;
   private final RelativeSymlinkBehavior relSymlinkBehavior;
   private final boolean expandArchivedTreeArtifacts;
@@ -87,25 +76,24 @@ public SpawnInputExpander(
   }
 
   private static void addMapping(
-      InputSink inputSink,
+      Map<PathFragment, ActionInput> inputMap,
       PathFragment targetLocation,
       ActionInput input,
-      PathFragment baseDirectory,
-      @Nullable Artifact owner) {
+      PathFragment baseDirectory) {
     Preconditions.checkArgument(!targetLocation.isAbsolute(), targetLocation);
-    inputSink.acceptMapping(baseDirectory.getRelative(targetLocation), input, owner);
+    inputMap.put(baseDirectory.getRelative(targetLocation), input);
   }
 
   @VisibleForTesting
   void addSingleRunfilesTreeToInputs(
       RunfilesTree runfilesTree,
-      InputSink inputSink,
+      Map<PathFragment, ActionInput> inputMap,
       ArtifactExpander artifactExpander,
       PathMapper pathMapper,
       PathFragment baseDirectory)
       throws ForbiddenActionInputException {
     addSingleRunfilesTreeToInputs(
-        inputSink,
+        inputMap,
         runfilesTree.getExecPath(),
         runfilesTree.getMapping(),
         artifactExpander,
@@ -120,7 +108,7 @@ void addSingleRunfilesTreeToInputs(
    * figure out how not to call this method (or else how to make this method more palatable)
    */
   public void addSingleRunfilesTreeToInputs(
-      InputSink inputSink,
+      Map<PathFragment, ActionInput> inputMap,
       PathFragment root,
       Map<PathFragment, Artifact> mappings,
       ArtifactExpander artifactExpander,
@@ -133,11 +121,10 @@ public void addSingleRunfilesTreeToInputs(
       Artifact artifact = mapping.getValue();
       if (artifact == null) {
         addMapping(
-            inputSink,
+            inputMap,
             mapForRunfiles(pathMapper, root, location),
             VirtualActionInput.EMPTY_MARKER,
-            baseDirectory,
-            /* owner= */ null);
+            baseDirectory);
         continue;
       }
       Preconditions.checkArgument(!artifact.isMiddlemanArtifact(), artifact);
@@ -146,7 +133,7 @@ public void addSingleRunfilesTreeToInputs(
             expandArchivedTreeArtifacts ? null : artifactExpander.getArchivedTreeArtifact(artifact);
         if (archivedTreeArtifact != null) {
           // TODO(bazel-team): Add path mapping support for archived tree artifacts.
-          addMapping(inputSink, location, archivedTreeArtifact, baseDirectory, artifact);
+          addMapping(inputMap, location, archivedTreeArtifact, baseDirectory);
         } else {
           List<ActionInput> expandedInputs =
               ActionInputHelper.expandArtifacts(
@@ -156,12 +143,11 @@ public void addSingleRunfilesTreeToInputs(
                   /* keepMiddlemanArtifacts= */ false);
           for (ActionInput input : expandedInputs) {
             addMapping(
-                inputSink,
+                inputMap,
                 mapForRunfiles(pathMapper, root, location)
                     .getRelative(((TreeFileArtifact) input).getParentRelativePath()),
                 input,
-                baseDirectory,
-                artifact);
+                baseDirectory);
           }
         }
       } else if (artifact.isFileset()) {
@@ -172,39 +158,33 @@ public void addSingleRunfilesTreeToInputs(
           throw new IllegalStateException(e);
         }
         // TODO(bazel-team): Add path mapping support for filesets.
-        addFilesetManifest(location, artifact, filesetLinks, inputSink, baseDirectory);
+        addFilesetManifest(location, artifact, filesetLinks, inputMap, baseDirectory);
       } else {
         // TODO: b/7075837 - If we want to prohibit directory inputs, we can check if
         //  localArtifact is a directory and, if so, throw a ForbiddenActionInputException.
-        addMapping(
-            inputSink,
-            mapForRunfiles(pathMapper, root, location),
-            artifact,
-            baseDirectory,
-            /* owner= */ null);
+        addMapping(inputMap, mapForRunfiles(pathMapper, root, location), artifact, baseDirectory);
       }
     }
   }
 
   @VisibleForTesting
   void addFilesetManifests(
       Map<Artifact, ImmutableList<FilesetOutputSymlink>> filesetMappings,
-      InputSink inputSink,
+      Map<PathFragment, ActionInput> inputMap,
       PathFragment baseDirectory)
       throws ForbiddenRelativeSymlinkException {
     for (Map.Entry<Artifact, ImmutableList<FilesetOutputSymlink>> entry :
         filesetMappings.entrySet()) {
       Artifact fileset = entry.getKey();
-      addFilesetManifest(
-          fileset.getExecPath(), fileset, entry.getValue(), inputSink, baseDirectory);
+      addFilesetManifest(fileset.getExecPath(), fileset, entry.getValue(), inputMap, baseDirectory);
     }
   }
 
   private void addFilesetManifest(
       PathFragment location,
       Artifact filesetArtifact,
       ImmutableList<FilesetOutputSymlink> filesetLinks,
-      InputSink inputSink,
+      Map<PathFragment, ActionInput> inputMap,
       PathFragment baseDirectory)
       throws ForbiddenRelativeSymlinkException {
     Preconditions.checkArgument(filesetArtifact.isFileset(), filesetArtifact);
@@ -218,7 +198,7 @@ private void addFilesetManifest(
               ? VirtualActionInput.EMPTY_MARKER
               : ActionInputHelper.fromPath(execRoot.getRelative(value).asFragment());
       // TODO(bazel-team): Add path mapping support for filesets.
-      addMapping(inputSink, mapping.getKey(), artifact, baseDirectory, filesetArtifact);
+      addMapping(inputMap, mapping.getKey(), artifact, baseDirectory);
     }
   }
 
@@ -230,7 +210,6 @@ private void addInputs(
       PathMapper pathMapper,
       PathFragment baseDirectory)
       throws ForbiddenActionInputException {
-    InputSink inputSink = InputSink.fromMap(inputMap);
     // Actions that accept TreeArtifacts as inputs generally expect the directory corresponding
     // to the artifact to be created, even if it is empty. We explicitly keep empty TreeArtifacts
     // here to signal consumers that they should create the directory.
@@ -243,27 +222,19 @@ private void addInputs(
     for (ActionInput input : inputs) {
       if (input instanceof TreeFileArtifact) {
         addMapping(
-            inputSink,
+            inputMap,
             pathMapper
                 .map(((TreeFileArtifact) input).getParent().getExecPath())
                 .getRelative(((TreeFileArtifact) input).getParentRelativePath()),
             input,
-            baseDirectory,
-            // Owners are disregarded since we're aggregating into a map.
-            /* owner= */ null);
+            baseDirectory);
       } else if (isMiddlemanArtifact(input)) {
         RunfilesTree runfilesTree =
             inputMetadataProvider.getRunfilesMetadata(input).getRunfilesTree();
         addSingleRunfilesTreeToInputs(
-            runfilesTree, inputSink, artifactExpander, pathMapper, baseDirectory);
+            runfilesTree, inputMap, artifactExpander, pathMapper, baseDirectory);
       } else {
-        addMapping(
-            inputSink,
-            pathMapper.map(input.getExecPath()),
-            input,
-            baseDirectory,
-            // Owners are disregarded since we're aggregating into a map.
-            /* owner= */ null);
+        addMapping(inputMap, pathMapper.map(input.getExecPath()), input, baseDirectory);
       }
     }
   }
@@ -286,15 +257,14 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
       PathFragment baseDirectory)
       throws ForbiddenActionInputException {
     TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
-    InputSink inputSink = InputSink.fromMap(inputMap);
     addInputs(
         inputMap,
         spawn.getInputFiles(),
         artifactExpander,
         inputMetadataProvider,
         spawn.getPathMapper(),
         baseDirectory);
-    addFilesetManifests(spawn.getFilesetMappings(), inputSink, baseDirectory);
+    addFilesetManifests(spawn.getFilesetMappings(), inputMap, baseDirectory);
     return inputMap;
   }
 
@@ -377,7 +347,7 @@ public void walkInputs(
           public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
               throws ForbiddenRelativeSymlinkException {
             TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
-            addFilesetManifests(filesetMappings, InputSink.fromMap(inputMap), baseDirectory);
+            addFilesetManifests(filesetMappings, inputMap, baseDirectory);
             return inputMap;
           }
         });
@@ -469,11 +439,7 @@ public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
               throws ForbiddenActionInputException {
             TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
             addSingleRunfilesTreeToInputs(
-                runfilesTree,
-                InputSink.fromMap(inputMap),
-                artifactExpander,
-                pathMapper,
-                baseDirectory);
+                runfilesTree, inputMap, artifactExpander, pathMapper, baseDirectory);
             return inputMap;
           }
         });