Add ArtifactNestedSetKeys dependencies to rewind graph for propagatin…

…g actions.

If a propagating action ("mayInsensitivelyPropagateInputs" returns
true) has input artifacts in a nested set, those nested sets are not
invalidated as part of the rewinding if the action needs to be rewound.

This is not currently an issue because the only such action is
SkyframeFilesetManifestAction, which has its own special rewinding logic
with FilesetEntryFunction. However we plan to remove this, so we require
that this case is correctly handled.

Care needs to be taken that we don't break existing code to add
ArtifactNestedSetKeys to the rewind graph, which assumes there are no
such keys in the graph already. The new code to add keys is called near
this process so this interaction is more obvious.

PiperOrigin-RevId: 725983745
Change-Id: I08db9d85f8190a863c21cad17fc2d342c766ef52

src/main/java/com/google/devtools/build/lib/collect/nestedset/ArtifactNestedSetKey.java

@@ -110,6 +110,30 @@ public ImmutableList<Artifact> expandToArtifacts() {
         .toList();
   }
 
+  /**
+   * Augments the given rewind graph with all chains of {@link ArtifactNestedSetKey} nodes reachable
+   * from {@code failedKeyDeps} and non source artifacts in them.
+   *
+   * <p>The walk is terminated when a node is already in the rewind graph.
+   */
+  public static void addNestedSetChainsToRewindGraph(
+      MutableGraph<SkyKey> rewindGraph, ArtifactNestedSetKey key) {
+    if (rewindGraph.nodes().contains(key)) {
+      return;
+    }
+    for (Object child : key.children) {
+      if (child instanceof Artifact artifact) {
+        if (!artifact.isSourceArtifact()) {
+          rewindGraph.putEdge(key, Artifact.key(artifact));
+        }
+      } else {
+        ArtifactNestedSetKey nextNode = createInternal((Object[]) child);
+        addNestedSetChainsToRewindGraph(rewindGraph, nextNode);
+        rewindGraph.putEdge(key, nextNode);
+      }
+    }
+  }
+
   /**
    * Augments the given rewind graph with paths from {@code failedKey} to {@code lostArtifacts}
    * discoverable by following the {@link ArtifactNestedSetKey} nodes in {@code failedKeyDeps}.

src/main/java/com/google/devtools/build/lib/skyframe/rewinding/ActionRewindStrategy.java

@@ -23,12 +23,14 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ConcurrentHashMultiset;
+import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
+import com.google.common.collect.SetMultimap;
 import com.google.common.flogger.GoogleLogger;
 import com.google.common.graph.EndpointPair;
 import com.google.common.graph.ImmutableGraph;
@@ -220,6 +222,11 @@ private Reset prepareRewindPlan(
     Set<DerivedArtifact> lostArtifacts =
         getLostInputOwningDirectDeps(failedKey, failedKeyDeps, lostInputs, inputDepOwners);
 
+    // Additional nested sets we may need to invalidate that are the dependencies of an
+    // insensitively propagating action, associated with the key that depends on them.
+    SetMultimap<SkyKey, ArtifactNestedSetKey> nestedSetsForPropagatingActions =
+        HashMultimap.create();
+
     for (DerivedArtifact lostArtifact : lostArtifacts) {
       Map<ActionLookupData, Action> actionMap = getActionsForLostArtifact(lostArtifact, env);
       if (actionMap == null) {
@@ -237,7 +244,23 @@ private Reset prepareRewindPlan(
       // always a transitive dep.
       rewindGraph.putEdge(failedKey, Artifact.key(lostArtifact));
       depsToRewind.addAll(actions(newlyVisitedActions));
-      checkActions(newlyVisitedActions, env, rewindGraph, depsToRewind);
+      checkActions(
+          newlyVisitedActions, env, rewindGraph, depsToRewind, nestedSetsForPropagatingActions);
+    }
+
+    // addNestedSetPathsToRewindGraph, called after this loop, stops its walk when it finds a node
+    // that is already in the rewind graph.
+    // However, because this rewinds all NestedSet chains from a given root, early termination later
+    // on won't matter because all dependent NestedSets are already in the graph.
+    // This may seem excessive, but it is not expected that many NestedSets are actually involved in
+    // this walk, and that this only happens rarely.
+    // TODO(b/395634488): This should be solved in a more elegant way, but a solution is needed to
+    // unblock the simplifications to Fileset (b/394611260)
+    for (SkyKey rootKey : nestedSetsForPropagatingActions.keySet()) {
+      for (ArtifactNestedSetKey nestedSetKey : nestedSetsForPropagatingActions.get(rootKey)) {
+        ArtifactNestedSetKey.addNestedSetChainsToRewindGraph(rewindGraph, nestedSetKey);
+        rewindGraph.putEdge(rootKey, nestedSetKey);
+      }
     }
 
     // This needs to be done after the loop above because addArtifactDepsAndGetNewlyVisitedActions
@@ -488,7 +511,8 @@ private void checkActions(
       ImmutableList<ActionAndLookupData> actionsToCheck,
       Environment env,
       MutableGraph<SkyKey> rewindGraph,
-      ImmutableList.Builder<Action> depsToRewind)
+      ImmutableList.Builder<Action> depsToRewind,
+      SetMultimap<SkyKey, ArtifactNestedSetKey> nestedSetDeps)
       throws InterruptedException {
     ArrayDeque<ActionAndLookupData> uncheckedActions = new ArrayDeque<>(actionsToCheck);
     while (!uncheckedActions.isEmpty()) {
@@ -514,7 +538,12 @@ private void checkActions(
         // Rewinding this action won't recreate the missing input. We need to also rewind this
         // action's non-source inputs and the actions which created those inputs.
         addPropagatingActionDepsAndGetNewlyVisitedArtifactsAndActions(
-            rewindGraph, actionKey, action, artifactsToCheck, newlyDiscoveredActions);
+            rewindGraph,
+            actionKey,
+            action,
+            artifactsToCheck,
+            newlyDiscoveredActions,
+            nestedSetDeps);
       }
 
       for (ActionLookupData actionLookupData : newlyDiscoveredActions) {
@@ -579,7 +608,8 @@ private void addPropagatingActionDepsAndGetNewlyVisitedArtifactsAndActions(
       ActionLookupData actionKey,
       Action action,
       ArrayList<DerivedArtifact> newlyVisitedArtifacts,
-      ArrayList<ActionLookupData> newlyVisitedActions) {
+      ArrayList<ActionLookupData> newlyVisitedActions,
+      SetMultimap<SkyKey, ArtifactNestedSetKey> nestedSetDeps) {
 
     for (Artifact input : action.getInputs().toList()) {
       if (input.isSourceArtifact()) {
@@ -602,6 +632,12 @@ private void addPropagatingActionDepsAndGetNewlyVisitedArtifactsAndActions(
       rewindGraph.putEdge(actionKey, artifactKey);
     }
 
+    // Record the action's NestedSet inputs as dependencies to be rewound.
+    action
+        .getInputs()
+        .getNonLeaves()
+        .forEach(nestedSet -> nestedSetDeps.put(actionKey, ArtifactNestedSetKey.create(nestedSet)));
+
     // Rewinding ignores artifacts returned by Action#getAllowedDerivedInputs because:
     // 1) the set of actions with non-throwing implementations of getAllowedDerivedInputs,
     // 2) the set of actions that "mayInsensitivelyPropagateInputs",