Skip to content

Commit

Permalink
Lazily flatten the mandatory inputs NestedSet.
Browse files Browse the repository at this point in the history
It is only necessary to query this set in the exceptional case of a missing artifact. Avoid flattening it until it becomes necessary.

PiperOrigin-RevId: 383145412
  • Loading branch information
justinhorvitz authored and copybara-github committed Jul 5, 2021
1 parent d435634 commit e9740fc
Showing 1 changed file with 36 additions and 28 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.google.common.base.Joiner;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicates;
import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableCollection;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -113,6 +114,7 @@
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.IntFunction;
import java.util.function.Predicate;
import java.util.function.Supplier;
import javax.annotation.Nullable;
import net.starlark.java.eval.StarlarkSemantics;
Expand Down Expand Up @@ -268,11 +270,8 @@ private SkyValue computeInternal(SkyKey skyKey, Environment env)
env.getValues(state.allInputs.packageLookupsRequested);
Preconditions.checkState(!env.valuesMissing(), "%s %s", action, state);
}
CheckInputResults checkedInputs = null;
@Nullable
ImmutableSet<Artifact> mandatoryInputs =
action.discoversInputs() ? action.getMandatoryInputs().toSet() : null;

CheckInputResults checkedInputs = null;
NestedSet<Artifact> allInputs = state.allInputs.getAllInputs();

Iterable<SkyKey> depKeys = getInputDepKeys(allInputs, state);
Expand All @@ -290,9 +289,7 @@ private SkyValue computeInternal(SkyKey skyKey, Environment env)
try {
if (previousExecution == null && !state.hasArtifactData()) {
// Do we actually need to find our metadata?
checkedInputs =
checkInputs(
env, action, inputDeps, allInputs, mandatoryInputs, depKeys, actionLookupData);
checkedInputs = checkInputs(env, action, inputDeps, allInputs, depKeys, actionLookupData);
}
} catch (ActionExecutionException e) {
// Remove action from state map in case it's there (won't be unless it discovers inputs).
Expand Down Expand Up @@ -572,7 +569,6 @@ private ActionInputDepOwners createAugmentedInputDepOwners(
action,
inputDeps,
allInputs,
action.discoversInputs() ? action.getMandatoryInputs().toSet() : null,
requestedSkyKeys,
lostInputsAndOwnersSoFar,
actionLookupDataForError);
Expand Down Expand Up @@ -1094,7 +1090,6 @@ private CheckInputResults checkInputs(
ArtifactNestedSetEvalException>>
inputDeps,
NestedSet<Artifact> allInputs,
ImmutableSet<Artifact> mandatoryInputs,
Iterable<SkyKey> requestedSkyKeys,
ActionLookupData actionLookupDataForError)
throws ActionExecutionException, InterruptedException {
Expand All @@ -1103,7 +1098,6 @@ private CheckInputResults checkInputs(
action,
inputDeps,
allInputs,
mandatoryInputs,
requestedSkyKeys,
ActionInputMap::new,
CheckInputResults::new,
Expand All @@ -1124,7 +1118,6 @@ private ActionInputDepOwnerMap getInputDepOwners(
ArtifactNestedSetEvalException>>
inputDeps,
NestedSet<Artifact> allInputs,
ImmutableSet<Artifact> mandatoryInputs,
Iterable<SkyKey> requestedSkyKeys,
Collection<ActionInput> lostInputs,
ActionLookupData actionLookupDataForError)
Expand All @@ -1138,7 +1131,6 @@ private ActionInputDepOwnerMap getInputDepOwners(
action,
inputDeps,
allInputs,
mandatoryInputs,
requestedSkyKeys,
ignoredInputDepsSize -> new ActionInputDepOwnerMap(lostInputs),
(actionInputMapSink,
Expand All @@ -1150,6 +1142,29 @@ private ActionInputDepOwnerMap getInputDepOwners(
actionLookupDataForError);
}

private static Predicate<Artifact> makeMandatoryInputPredicate(Action action) {
if (!action.discoversInputs()) {
return Predicates.alwaysTrue();
}

return new Predicate<Artifact>() {
// Lazily flatten the NestedSet in case the predicate is never needed. It's only used in the
// exceptional case of a missing artifact.
private ImmutableSet<Artifact> mandatoryInputs = null;

@Override
public boolean test(Artifact input) {
if (!input.isSourceArtifact()) {
return true;
}
if (mandatoryInputs == null) {
mandatoryInputs = action.getMandatoryInputs().toSet();
}
return mandatoryInputs.contains(input);
}
};
}

/**
* May return {@code null} if {@code allowValuesMissingEarlyReturn} and {@code
* env.valuesMissing()} are true and no inputs result in {@link ActionExecutionException}s.
Expand All @@ -1165,14 +1180,13 @@ private <S extends ActionInputMapSink, R> R accumulateInputs(
ArtifactNestedSetEvalException>>
inputDeps,
NestedSet<Artifact> allInputs,
ImmutableSet<Artifact> mandatoryInputs,
Iterable<SkyKey> requestedSkyKeys,
IntFunction<S> actionInputMapSinkFactory,
AccumulateInputResultsFactory<S, R> accumulateInputResultsFactory,
boolean allowValuesMissingEarlyReturn,
ActionLookupData actionLookupDataForError)
throws ActionExecutionException, InterruptedException {

Predicate<Artifact> isMandatoryInput = makeMandatoryInputPredicate(action);
ActionExecutionFunctionExceptionHandler actionExecutionFunctionExceptionHandler =
new ActionExecutionFunctionExceptionHandler(
Suppliers.memoize(
Expand All @@ -1191,7 +1205,7 @@ private <S extends ActionInputMapSink, R> R accumulateInputs(
}),
inputDeps,
action,
mandatoryInputs,
isMandatoryInput,
requestedSkyKeys,
env.valuesMissing());

Expand All @@ -1218,7 +1232,7 @@ private <S extends ActionInputMapSink, R> R accumulateInputs(
for (Artifact input : allInputsList) {
SkyValue value = ArtifactNestedSetFunction.getInstance().getValueForKey(Artifact.key(input));
if (value == null) {
if (mandatoryInputs == null || mandatoryInputs.contains(input)) {
if (isMandatoryInput.test(input)) {
BugReport.sendBugReport(
new IllegalStateException(
String.format(
Expand All @@ -1228,7 +1242,7 @@ private <S extends ActionInputMapSink, R> R accumulateInputs(
continue;
}
if (value instanceof MissingArtifactValue) {
if (actionExecutionFunctionExceptionHandler.isMandatory(input)) {
if (isMandatoryInput.test(input)) {
actionExecutionFunctionExceptionHandler.accumulateMissingFileArtifactValue(
input, (MissingArtifactValue) value);
continue;
Expand Down Expand Up @@ -1438,7 +1452,7 @@ private final class ActionExecutionFunctionExceptionHandler {
SourceArtifactException, ActionExecutionException, ArtifactNestedSetEvalException>>
inputDeps;
private final Action action;
private final Set<Artifact> mandatoryInputs;
private final Predicate<Artifact> isMandatoryInput;
private final Iterable<SkyKey> requestedSkyKeys;
private final boolean valuesMissing;
List<LabelCause> missingArtifactCauses = Lists.newArrayListWithCapacity(0);
Expand All @@ -1454,13 +1468,13 @@ private final class ActionExecutionFunctionExceptionHandler {
ArtifactNestedSetEvalException>>
inputDeps,
Action action,
Set<Artifact> mandatoryInputs,
Predicate<Artifact> isMandatoryInput,
Iterable<SkyKey> requestedSkyKeys,
boolean valuesMissing) {
this.skyKeyToDerivedArtifactSetForExceptions = skyKeyToDerivedArtifactSetForExceptions;
this.inputDeps = inputDeps;
this.action = action;
this.mandatoryInputs = mandatoryInputs;
this.isMandatoryInput = isMandatoryInput;
this.requestedSkyKeys = requestedSkyKeys;
this.valuesMissing = valuesMissing;
}
Expand Down Expand Up @@ -1540,7 +1554,7 @@ private void handleSourceArtifactExceptionFromSkykey(SkyKey key, SourceArtifactE
return;
}

if (isMandatory((Artifact) key)) {
if (isMandatoryInput.test((Artifact) key)) {
missingArtifactCauses.add(
createLabelCauseNullOwnerOk(
(Artifact) key,
Expand Down Expand Up @@ -1586,15 +1600,9 @@ void maybeThrowException() throws ActionExecutionException {
}
}

boolean isMandatory(Artifact input) {
return !input.isSourceArtifact()
|| mandatoryInputs == null
|| mandatoryInputs.contains(input);
}

private void handleActionExecutionExceptionPerArtifact(
Artifact input, ActionExecutionException e) {
if (isMandatory(input)) {
if (isMandatoryInput.test(input)) {
// Prefer a catastrophic exception as the one we propagate.
if (firstActionExecutionException == null
|| (!firstActionExecutionException.isCatastrophe() && e.isCatastrophe())) {
Expand Down

0 comments on commit e9740fc

Please sign in to comment.