Skip to content

Commit

Permalink
Change rewinding's ArtifactNestedSetKey search strategy to improve …
Browse files Browse the repository at this point in the history 
…worst-case performance.

Instead of doing a search for every (lost artifact, `ArtifactNestedSetKey`) pair, do a single "bulk" search looking for all lost artifacts. When there are many lost artifacts, this is much more efficient. Additionally, we can prune visitations of nodes shared by multiple `NestedSet`s.

PiperOrigin-RevId: 663981449
Change-Id: Ic06f65e4f1dc9ca16a9bdd4a9c520b3ecd44be57
  • Loading branch information
@justinhorvitz @copybara-github
justinhorvitz authored and copybara-github committed Aug 17, 2024
1 parent 0fa8f9f commit ffdf41a
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 121 deletions.
74 changes: 41 additions & 33 deletions src/main/java/com/google/devtools/build/lib/collect/nestedset/ArtifactNestedSetKey.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,14 @@
import static com.google.common.base.Preconditions.checkArgument;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Sets;
import com.google.common.collect.Iterables;
import com.google.common.graph.MutableGraph;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.skyframe.SkyFunctions;
import com.google.devtools.build.skyframe.ExecutionPhaseSkyKey;
import com.google.devtools.build.skyframe.SkyFunctionName;
import com.google.devtools.build.skyframe.SkyKey;
import java.util.HashSet;
import java.util.Set;

/**
Expand Down Expand Up @@ -109,46 +111,52 @@ public ImmutableList<Artifact> expandToArtifacts() {
}

/**
* Finds a path from this key to the given {@link Artifact}, which must be contained beneath the
* node represented by this key.
* Augments the given rewind graph with paths from {@code failedKey} to {@code lostArtifacts}
* discoverable by following the {@link ArtifactNestedSetKey} nodes in {@code failedKeyDeps}.
*
* <p>The returned list represents a Skyframe dependency chain from this key to the target
* artifact. The first element is always {@code this}, and the last element contains the target
* artifact as a direct child. The list may contain only a single element if the target artifact
* is a direct child of the node represented by this key.
* <p>{@code rewindGraph} must not contain any {@link ArtifactNestedSetKey} nodes prior to calling
* this method.
*/
public ImmutableList<ArtifactNestedSetKey> findPathToArtifact(Artifact target) {
ImmutableList.Builder<ArtifactNestedSetKey> path = ImmutableList.builder();
boolean found = findPathToArtifact(target, children, path, Sets.newIdentityHashSet());
checkArgument(found, "Artifact not found: %s", target);
path.add(this);
return path.build().reverse();
public static void addNestedSetPathsToRewindGraph(
MutableGraph<SkyKey> rewindGraph,
SkyKey failedKey,
Set<SkyKey> failedKeyDeps,
Set<? extends Artifact> lostArtifacts) {
var seen = new HashSet<ArtifactNestedSetKey>();
for (var nestedSetDep : Iterables.filter(failedKeyDeps, ArtifactNestedSetKey.class)) {
if (searchForLostArtifacts(nestedSetDep, rewindGraph, lostArtifacts, seen)) {
rewindGraph.putEdge(failedKey, nestedSetDep);
}
}
}

private static boolean findPathToArtifact(
Artifact target,
Object[] currentNode,
ImmutableList.Builder<ArtifactNestedSetKey> path,
Set<Object[]> seen) {
if (!seen.add(currentNode)) {
return false;
private static boolean searchForLostArtifacts(
ArtifactNestedSetKey node,
MutableGraph<SkyKey> rewindGraph,
Set<? extends Artifact> lostArtifacts,
Set<ArtifactNestedSetKey> seen) {
if (rewindGraph.nodes().contains(node)) {
return true;
}
// Check all leaves before entering recursion to favor finding a shorter path.
for (Object child : currentNode) {
if (child.equals(target)) {
return true;
}
if (!seen.add(node)) {
return false;
}
for (Object child : currentNode) {
if (!(child instanceof Object[] nextNode)) {
continue;
}
if (findPathToArtifact(target, nextNode, path, seen)) {
path.add(createInternal(nextNode));
return true;
boolean anyFound = false;
for (Object child : node.children) {
if (child instanceof Artifact artifact) {
if (lostArtifacts.contains(artifact)) {
rewindGraph.putEdge(node, Artifact.key(artifact));
anyFound = true;
}
} else {
ArtifactNestedSetKey nextNode = createInternal((Object[]) child);
if (searchForLostArtifacts(nextNode, rewindGraph, lostArtifacts, seen)) {
rewindGraph.putEdge(node, nextNode);
anyFound = true;
}
}
}
return false;
return anyFound;
}

@Override
Expand Down
6 changes: 3 additions & 3 deletions src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1309,7 +1309,7 @@ private <S extends ActionInputMapSink, R> R accumulateInputs(
}

if (!undoneInputs.isEmpty()) {
throw new UndoneInputsException(ImmutableList.copyOf(undoneInputs), inputDepKeys);
throw new UndoneInputsException(ImmutableSet.copyOf(undoneInputs), inputDepKeys);
}

// If there were no errors, we don't go through the scheduling dependencies because the only
Expand Down Expand Up @@ -1549,10 +1549,10 @@ public boolean isCatastrophic() {
* completed with an error.
*/
private static final class UndoneInputsException extends Exception {
private final ImmutableList<Artifact> undoneInputs;
private final ImmutableSet<Artifact> undoneInputs;
private final ImmutableSet<SkyKey> inputDepKeys;

UndoneInputsException(ImmutableList<Artifact> undoneInputs, ImmutableSet<SkyKey> inputDepKeys) {
UndoneInputsException(ImmutableSet<Artifact> undoneInputs, ImmutableSet<SkyKey> inputDepKeys) {
this.undoneInputs = undoneInputs;
this.inputDepKeys = inputDepKeys;
}
Expand Down
117 changes: 32 additions & 85 deletions src/main/java/com/google/devtools/build/lib/skyframe/rewinding/ActionRewindStrategy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,6 @@
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.MultimapBuilder;
import com.google.common.flogger.GoogleLogger;
import com.google.common.graph.EndpointPair;
import com.google.common.graph.ImmutableGraph;
Expand Down Expand Up @@ -207,7 +205,7 @@ public Reset prepareRewindPlanForLostInputs(

private Reset prepareRewindPlan(
SkyKey failedKey,
Set<SkyKey> failedActionDeps,
Set<SkyKey> failedKeyDeps,
ImmutableMap<String, ActionInput> lostInputsByDigest,
ActionInputDepOwners inputDepOwners,
Environment env,
Expand All @@ -219,23 +217,10 @@ private Reset prepareRewindPlan(
// between them.
MutableGraph<SkyKey> rewindGraph = Reset.newRewindGraphFor(failedKey);

// With NSOS, not all input artifacts' keys are direct deps of the action. This maps input
// artifacts to their containing direct dep ArtifactNestedSetKey(s).
Multimap<Artifact, ArtifactNestedSetKey> nestedSetKeys = expandNestedSetKeys(failedActionDeps);

Set<DerivedArtifact> lostArtifacts =
getLostInputOwningDirectDeps(
failedKey,
lostInputs,
inputDepOwners,
ImmutableSet.<SkyKey>builder()
.addAll(failedActionDeps)
.addAll(Artifact.keys(nestedSetKeys.keySet()))
.build());
getLostInputOwningDirectDeps(failedKey, failedKeyDeps, lostInputs, inputDepOwners);

for (DerivedArtifact lostArtifact : lostArtifacts) {
SkyKey artifactKey = Artifact.key(lostArtifact);

Map<ActionLookupData, Action> actionMap = getActionsForLostArtifact(lostArtifact, env);
if (actionMap == null) {
// Some deps of the artifact are not done. Another rewind must be in-flight, and there is no
Expand All @@ -245,18 +230,21 @@ private Reset prepareRewindPlan(
ImmutableList<ActionAndLookupData> newlyVisitedActions =
addArtifactDepsAndGetNewlyVisitedActions(rewindGraph, lostArtifact, actionMap);

for (ArtifactNestedSetKey nestedSetKey : nestedSetKeys.get(lostArtifact)) {
addNestedSetToRewindGraph(rewindGraph, failedKey, lostArtifact, nestedSetKey);
}
// Note that artifactKey must be rewound. We do this after
// Note that Artifact.key(lostArtifact) must be rewound. We do this after
// addArtifactDepsAndGetNewlyVisitedActions so that it can track if actions are already known
// to be in the graph. With NSOS, it is possible that artifactKey is not actually a direct dep
// of the action, but this edge is benign since it's always a transitive dep.
rewindGraph.putEdge(failedKey, artifactKey);
// to be in the graph. It is possible that the Artifact.key() is not actually a direct dep of
// the action if it is below an ArtifactNestedSetKey, but this edge is benign since it's
// always a transitive dep.
rewindGraph.putEdge(failedKey, Artifact.key(lostArtifact));
depsToRewind.addAll(actions(newlyVisitedActions));
checkActions(newlyVisitedActions, env, rewindGraph, depsToRewind);
}

// This needs to be done after the loop above because addArtifactDepsAndGetNewlyVisitedActions
// short-circuits when a node is already in the rewind graph.
ArtifactNestedSetKey.addNestedSetPathsToRewindGraph(
rewindGraph, failedKey, failedKeyDeps, lostArtifacts);

return Reset.of(rewindGraph);
}

Expand All @@ -279,22 +267,19 @@ private Reset prepareRewindPlan(
public Reset patchNestedSetGraphToPropagateError(
ActionLookupData failedKey,
Action failedAction,
ImmutableList<Artifact> undoneInputs,
ImmutableSet<Artifact> undoneInputs,
ImmutableSet<SkyKey> failedActionDeps) {
checkState(
skyframeActionExecutor.rewindingEnabled(), "Unexpected undone inputs: %s", undoneInputs);
MutableGraph<SkyKey> rewindGraph = Reset.newRewindGraphFor(failedKey);
Multimap<Artifact, ArtifactNestedSetKey> nestedSetKeys = expandNestedSetKeys(failedActionDeps);
ArtifactNestedSetKey.addNestedSetPathsToRewindGraph(
rewindGraph, failedKey, failedActionDeps, undoneInputs);
for (Artifact input : undoneInputs) {
Collection<ArtifactNestedSetKey> containingNestedSets = nestedSetKeys.get(input);
checkState(
!containingNestedSets.isEmpty(),
rewindGraph.nodes().contains(Artifact.key(input)),
"Cannot find input %s under any nested set deps of %s",
input,
failedKey);
for (ArtifactNestedSetKey nestedSetKey : containingNestedSets) {
addNestedSetToRewindGraph(rewindGraph, failedKey, input, nestedSetKey);
}
}

// An undone input may be observed either during input checking (before attempting action
Expand Down Expand Up @@ -418,9 +403,20 @@ private static ImmutableList<LostInputRecord> createLostInputRecords(

private Set<DerivedArtifact> getLostInputOwningDirectDeps(
SkyKey failedKey,
Set<SkyKey> failedKeyDeps,
ImmutableList<ActionInput> lostInputs,
ActionInputDepOwners inputDepOwners,
Set<SkyKey> failedActionDeps) {
ActionInputDepOwners inputDepOwners) {
// Not all input artifacts' keys are direct deps - they may be below an ArtifactNestedSetKey.
// Expand all ArtifactNestedSetKey deps to get a flat set with all input artifact keys.
Set<SkyKey> expandedDeps = new HashSet<>();
for (SkyKey dep : failedKeyDeps) {
if (dep instanceof ArtifactNestedSetKey nestedSetDep) {
expandedDeps.addAll(Artifact.keys(nestedSetDep.expandToArtifacts()));
} else {
expandedDeps.add(dep);
}
}

Set<DerivedArtifact> lostInputOwningDirectDeps = new HashSet<>();
for (ActionInput lostInput : lostInputs) {
boolean foundLostInputDepOwner = false;
Expand All @@ -439,7 +435,7 @@ private Set<DerivedArtifact> getLostInputOwningDirectDeps(
for (Artifact transitiveOwner : transitiveOwners) {
checkDerived(transitiveOwner);

if (failedActionDeps.contains(Artifact.key(transitiveOwner))) {
if (expandedDeps.contains(Artifact.key(transitiveOwner))) {
// The lost input is included in an aggregation artifact (e.g. a tree artifact or
// fileset) that is included by an aggregation artifact (e.g. a middleman) that the
// action directly depends on.
Expand All @@ -448,16 +444,15 @@ private Set<DerivedArtifact> getLostInputOwningDirectDeps(
}
}

if (failedActionDeps.contains(Artifact.key(owner))) {
if (expandedDeps.contains(Artifact.key(owner))) {
// The lost input is included in an aggregation artifact (e.g. a tree artifact, fileset,
// or middleman) that the action directly depends on.
lostInputOwningDirectDeps.add((DerivedArtifact) owner);
foundLostInputDepOwner = true;
}
}

if (lostInput instanceof Artifact artifact
&& failedActionDeps.contains(Artifact.key(artifact))) {
if (lostInput instanceof Artifact artifact && expandedDeps.contains(Artifact.key(artifact))) {
checkDerived((Artifact) lostInput);

lostInputOwningDirectDeps.add((DerivedArtifact) lostInput);
Expand Down Expand Up @@ -817,52 +812,4 @@ static ActionAndLookupData create(ActionLookupData lookupData, Action action) {
private static List<Action> actions(List<ActionAndLookupData> newlyVisitedActions) {
return Lists.transform(newlyVisitedActions, ActionAndLookupData::action);
}

/**
* Constructs a mapping from input artifact to all direct dep {@link ArtifactNestedSetKey}s that
* transitively contain the artifact.
*
* <p>More formally, a key-value pair {@code (Artifact k, ArtifactNestedSetKey v)} is present in
* the returned map iff {@code deps.contains(v) && v.expandToArtifacts().contains(k)}.
*
* <p>When {@link com.google.devtools.build.lib.skyframe.ActionExecutionFunction} requests input
* deps, it unwraps a single layer of {@linkplain Action#getInputs the action's inputs}, thus
* requesting an {@link ArtifactNestedSetKey} for each of {@code
* action.getInputs().getNonLeaves()}.
*/
private static Multimap<Artifact, ArtifactNestedSetKey> expandNestedSetKeys(Set<SkyKey> deps) {
Multimap<Artifact, ArtifactNestedSetKey> map =
MultimapBuilder.hashKeys().arrayListValues().build();
for (ArtifactNestedSetKey key : Iterables.filter(deps, ArtifactNestedSetKey.class)) {
for (Artifact artifact : key.expandToArtifacts()) {
map.put(artifact, key);
}
}
return map;
}

/**
* Adds a skyframe dependency chain from {@code failedKey} to {@code lostArtifactKey} to the
* rewind graph.
*
* <p>Each edge along the path from {@code failedKey} to {@code lostArtifactKey} is added to the
* rewind graph. This necessarily includes the initial edge {@code (failedKey, directDep)}.
*
* <p>Although {@code lostArtifact} may be reachable via multiple distinct paths, it is only
* necessary to rewind one such path to ensure successful completion of {@code failedKey}. Other
* failing actions that depend on {@code lostArtifact} via a different path may initiate their own
* rewind strategy.
*/
private static void addNestedSetToRewindGraph(
MutableGraph<SkyKey> rewindGraph,
SkyKey failedKey,
Artifact lostArtifact,
ArtifactNestedSetKey directDep) {
SkyKey current = failedKey;
for (ArtifactNestedSetKey key : directDep.findPathToArtifact(lostArtifact)) {
rewindGraph.putEdge(current, key);
current = key;
}
rewindGraph.putEdge(current, Artifact.key(lostArtifact));
}
}

0 comments on commit ffdf41a

Please sign in to comment.