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

// Copyright 2018 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.google.devtools.build.lib.skyframe;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import com.google.devtools.build.lib.actions.Action;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionLookupData;
import com.google.devtools.build.lib.actions.AlreadyReportedActionExecutionException;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.LostInputsExecException.LostInputsActionExecutionException;
import com.google.devtools.build.lib.skyframe.ActionExecutionFunction.ActionExecutionFunctionException;
import com.google.devtools.build.skyframe.SkyFunction.Environment;
import com.google.devtools.build.skyframe.SkyFunction.Restart;
import com.google.devtools.build.skyframe.SkyKey;
import java.util.HashSet;
import java.util.Set;
import javax.annotation.Nullable;

/**
 * Given an action that failed to execute because of missing inputs which were generated by other
 * actions, this finds the Skyframe nodes corresponding to those inputs and the actions which
 * generated them.
 */
class ActionRewindStrategy {

  /**
   * Returns a {@link RewindPlan} specifying:
   *
   * <ol>
   *   <li>the Skyframe nodes to restart to recreate the lost inputs specified by {@code
   *       lostInputsException}
   *   <li>the actions whose execution state (in {@link SkyframeActionExecutor}) must be reset
   * </ol>
   *
   * @throws ActionExecutionFunctionException if any lost inputs are not the outputs of previously
   *     executed actions
   */
  // TODO(mschaller): support special/tree artifact types
  RewindPlan getRewindPlan(
      Action failedAction,
      Iterable<SkyKey> inputDepKeys,
      LostInputsActionExecutionException lostInputsException,
      Environment env)
      throws ActionExecutionFunctionException, InterruptedException {
    ImmutableList<ActionInput> lostInputs = lostInputsException.getLostInputs();
    for (ActionInput actionInput : lostInputs) {
      if (!(actionInput instanceof Artifact) || ((Artifact) actionInput).isSourceArtifact()) {
        throw new ActionExecutionFunctionException(
            new AlreadyReportedActionExecutionException(lostInputsException));
      }
    }

    // Find the action execution nodes that lost inputs depend on.
    HashSet<SkyKey> depsToRestart = Sets.newHashSet();
    ImmutableSet<SkyKey> inputDepKeysSet = ImmutableSet.copyOf(inputDepKeys);
    ImmutableList.Builder<Action> actionsToReset = ImmutableList.builder();
    for (ActionInput lostInput : lostInputs) {
      Preconditions.checkState(
          inputDepKeysSet.contains(lostInput),
          "Lost input not a dep of action.\nLost input: %s\nDeps: %s\nAction: %s",
          lostInput,
          inputDepKeysSet,
          failedAction);
      depsToRestart.add((Artifact) lostInput);

      Set<SkyKey> depsOfArtifact = getDepsOfArtifact((Artifact) lostInput, env);
      if (depsOfArtifact == null) {
        // Some deps of the artifact are not done. Another rewind must be in-flight, and there is no
        // need to restart the shared deps twice.
        continue;
      }
      // Restart the execution-phase dependencies of the artifact.
      depsToRestart.addAll(depsOfArtifact);
    }

    // SkyframeActionExecutor must re-execute the actions being restarted, so we must tell it to
    // evict its cached action results.
    actionsToReset.add(failedAction);
    for (SkyKey depToRestart : depsToRestart) {
      if (!(depToRestart instanceof ActionLookupData)) {
        continue;
      }
      actionsToReset.add(
          ActionExecutionFunction.getActionForLookupData(env, (ActionLookupData) depToRestart));
    }

    return new RewindPlan(
        Restart.selfAnd(ImmutableList.copyOf(depsToRestart)), actionsToReset.build());
  }

  /**
   * Returns the set of {@code lostInput}'s execution-phase dependencies, or {@code null} if any of
   * those dependencies are not done.
   */
  @Nullable
  private Set<SkyKey> getDepsOfArtifact(Artifact lostInput, Environment env)
      throws InterruptedException {
    ArtifactFunction.ArtifactDependencies artifactDependencies =
        ArtifactFunction.ArtifactDependencies.discoverDependencies(lostInput, env);
    if (artifactDependencies == null) {
      return null;
    }
    Preconditions.checkState(
        !artifactDependencies.isTemplateActionForTreeArtifact(),
        "Rewinding template actions not yet supported: %s",
        artifactDependencies);
    // TODO(mschaller): extend ArtifactDependencies to support template actions (and other special
    // cases)
    return ImmutableSet.of(artifactDependencies.getNontemplateActionExecutionKey());
  }

  static class RewindPlan {
    private final Restart nodesToRestart;
    private final ImmutableList<Action> actionsToReset;

    RewindPlan(Restart nodesToRestart, ImmutableList<Action> actionsToReset) {
      this.nodesToRestart = nodesToRestart;
      this.actionsToReset = actionsToReset;
    }

    Restart getNodesToRestart() {
      return nodesToRestart;
    }

    ImmutableList<Action> getActionsToReset() {
      return actionsToReset;
    }
  }
}