// 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.rewinding;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.joining;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ConcurrentHashMultiset;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableCollection;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multiset;
import com.google.common.collect.SetMultimap;
import com.google.common.flogger.GoogleLogger;
import com.google.common.graph.MutableGraph;
import com.google.devtools.build.lib.actions.Action;
import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionLookupData;
import com.google.devtools.build.lib.actions.ActionLookupValue;
import com.google.devtools.build.lib.actions.ActionTemplate;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.Artifact.DerivedArtifact;
import com.google.devtools.build.lib.actions.FilesetOutputSymlink;
import com.google.devtools.build.lib.actions.InputMetadataProvider;
import com.google.devtools.build.lib.actions.LostInputsActionExecutionException;
import com.google.devtools.build.lib.actions.RunfilesArtifactValue;
import com.google.devtools.build.lib.actions.RunfilesTree;
import com.google.devtools.build.lib.bugreport.BugReporter;
import com.google.devtools.build.lib.clock.BlazeClock;
import com.google.devtools.build.lib.collect.nestedset.ArtifactNestedSetKey;
import com.google.devtools.build.lib.events.ExtendedEventHandler;
import com.google.devtools.build.lib.profiler.AutoProfiler;
import com.google.devtools.build.lib.profiler.ProfilerTask;
import com.google.devtools.build.lib.remote.common.LostInputsEvent;
import com.google.devtools.build.lib.server.FailureDetails.ActionRewinding;
import com.google.devtools.build.lib.server.FailureDetails.ActionRewinding.Code;
import com.google.devtools.build.lib.skyframe.ActionUtils;
import com.google.devtools.build.lib.skyframe.ArtifactFunction;
import com.google.devtools.build.lib.skyframe.ArtifactFunction.ArtifactDependencies;
import com.google.devtools.build.lib.skyframe.SkyframeActionExecutor;
import com.google.devtools.build.lib.skyframe.TopLevelActionLookupKeyWrapper;
import com.google.devtools.build.lib.skyframe.proto.ActionRewind.ActionDescription;
import com.google.devtools.build.lib.skyframe.proto.ActionRewind.ActionRewindEvent;
import com.google.devtools.build.lib.skyframe.proto.ActionRewind.LostInput;
import com.google.devtools.build.lib.skyframe.rewinding.ActionRewindException.FallbackToBuildRewindingException;
import com.google.devtools.build.lib.skyframe.rewinding.ActionRewindException.GenericActionRewindException;
import com.google.devtools.build.lib.skyframe.serialization.analysis.RemoteAnalysisCacheReaderDepsProvider;
import com.google.devtools.build.skyframe.SkyFunction.Environment;
import com.google.devtools.build.skyframe.SkyFunction.Reset;
import com.google.devtools.build.skyframe.SkyKey;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import javax.annotation.Nullable;

/**
 * Given an action that failed to execute because of lost inputs which were generated by other
 * actions, finds the actions which generated them and the set of Skyframe nodes which must be
 * rewound in order to recreate the lost inputs.
 */
public final class ActionRewindStrategy {

  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
  @VisibleForTesting static final int MAX_REPEATED_LOST_INPUTS = 20;
  @VisibleForTesting static final int MAX_ACTION_REWIND_EVENTS = 5;
  private static final int MAX_LOST_INPUTS_RECORDED = 5;

  private final SkyframeActionExecutor skyframeActionExecutor;
  private final BugReporter bugReporter;

  private Map<SkyKey, Multiset<LostInputRecord>> currentBuildLostInputRecords =
      new ConcurrentHashMap<>();
  private final AtomicInteger lostInputsCount = new AtomicInteger();
  private final AtomicInteger lostOutputsCount = new AtomicInteger();
  private final AtomicInteger sameActionLostInputsCount = new AtomicInteger();
  private final AtomicInteger sameTopLevelKeyLostOutputsCount = new AtomicInteger();

  private final List<ActionRewindEvent> rewindEventSamples =
      Collections.synchronizedList(new ArrayList<>(MAX_ACTION_REWIND_EVENTS));
  private final AtomicInteger rewindEventSampleCounter = new AtomicInteger();

  private final Supplier<RemoteAnalysisCacheReaderDepsProvider> cachingDependenciesSupplier;

  public ActionRewindStrategy(
      SkyframeActionExecutor skyframeActionExecutor,
      BugReporter bugReporter,
      Supplier<RemoteAnalysisCacheReaderDepsProvider> cachingDependenciesSupplier) {
    this.skyframeActionExecutor = checkNotNull(skyframeActionExecutor);
    this.bugReporter = checkNotNull(bugReporter);
    this.cachingDependenciesSupplier = cachingDependenciesSupplier;
  }

  /** The result of rewind planning. */
  public static final class RewindPlanResult {
    @Nullable // null if there were missing dependencies
    private final Reset reset;

    private RewindPlanResult(@Nullable Reset reset) {
      this.reset = reset;
    }

    @Nullable
    public Reset toNullIfMissingDependenciesElseReset() {
      return reset;
    }
  }

  /**
   * Returns a {@link RewindPlanResult} specifying the Skyframe nodes to rewind to recreate the lost
   * inputs specified by {@code lostInputsException}.
   *
   * <p>If {@code #allowSkyframeRestarts} is true, the returned {@link RewindPlanResult} may return
   * null for {@link RewindPlanResult#toNullIfMissingDependenciesElseReset} if some dependencies are
   * not yet done. Otherwise, the returned {@link
   * RewindPlanResult#toNullIfMissingDependenciesElseReset} always returns a {@link Reset}.
   *
   * <p>Also prepares {@link SkyframeActionExecutor} for the rewind plan.
   *
   * <p>Lost outputs and their owning expansion artifacts are removed from {@code builtArtifacts},
   * as they should not be reported as built.
   *
   * @throws ActionRewindException if rewinding is disabled, or if any lost outputs have been seen
   *     by {@code failedKey} as lost before too many times
   */
  public RewindPlanResult prepareRewindPlanForLostTopLevelOutputs(
      TopLevelActionLookupKeyWrapper failedKey,
      Set<SkyKey> failedKeyDeps,
      ImmutableSetMultimap<String, ActionInput> lostOutputsByDigest,
      InputMetadataProvider metadataProvider,
      Set<Artifact> builtArtifacts,
      Environment env)
      throws ActionRewindException, InterruptedException {
    checkRewindingEnabled(lostOutputsByDigest, LostType.OUTPUT, env.getListener());

    ImmutableList<LostInputRecord> lostRecords = createLostInputRecords(lostOutputsByDigest);

    SetMultimap<ActionInput, Artifact> owners =
        calculateLostInputOwners(lostOutputsByDigest.values(), metadataProvider);

    for (ActionInput lostOutput : lostOutputsByDigest.values()) {
      builtArtifacts.remove(lostOutput);
      builtArtifacts.removeAll(owners.get(lostOutput));
    }

    ImmutableList.Builder<ActionAnalysisMetadata> depsToRewind = ImmutableList.builder();
    RewindPlanResult rewindPlanResult;
    try (var ignored =
        AutoProfiler.profiled(
            "Preparing rewind plan for %d lost outputs of %s"
                .formatted(lostRecords.size(), failedKey.actionLookupKey().getLabel()),
            ProfilerTask.ACTION_REWINDING)) {
      rewindPlanResult =
          prepareRewindPlan(
              failedKey, failedKeyDeps, lostOutputsByDigest, owners, env, depsToRewind);
    }
    Reset reset = rewindPlanResult.reset;
    if (reset == null) {
      // Skips steps that manipulate persistent state, structures retained by this class and
      // information tracked by SkyframeActionExecutor, if a restart is needed.
      //
      // This includes counting if an input has been lost too many times.
      return rewindPlanResult;
    }

    checkIfTopLevelOutputLostTooManyTimes(failedKey, lostRecords, lostOutputsByDigest);

    if (shouldRecordRewindEventSample()) {
      rewindEventSamples.add(createLostOutputRewindEvent(failedKey, reset, lostRecords));
    }

    for (ActionAnalysisMetadata dep : depsToRewind.build()) {
      skyframeActionExecutor.prepareDepForRewinding(failedKey, dep);
    }

    return rewindPlanResult;
  }

  /**
   * Returns a {@link RewindPlanResult} specifying the Skyframe nodes to rewind to recreate the lost
   * inputs specified by {@code lostInputsException}.
   *
   * <p>If {@code #allowSkyframeRestarts} is true, the returned {@link RewindPlanResult} may return
   * null for {@link RewindPlanResult#toNullIfMissingDependenciesElseReset} if some dependencies are
   * not yet done. Otherwise, the returned {@link
   * RewindPlanResult#toNullIfMissingDependenciesElseReset} always returns a {@link Reset}.
   *
   * <p>Also prepares {@link SkyframeActionExecutor} for the rewind plan and emits an {@link
   * ActionRewoundEvent} if necessary.
   *
   * @throws ActionRewindException if rewinding is disabled, or if any lost inputs have been seen by
   *     {@code failedKey} as lost before too many times
   */
  public RewindPlanResult prepareRewindPlanForLostInputs(
      ActionLookupData failedKey,
      Action failedAction,
      Set<SkyKey> failedActionDeps,
      LostInputsActionExecutionException e,
      InputMetadataProvider metadataProvider,
      Environment env,
      long actionStartTimeNanos)
      throws ActionRewindException, InterruptedException {
    ImmutableSetMultimap<String, ActionInput> lostInputsByDigest = e.getLostInputs();
    checkRewindingEnabled(lostInputsByDigest, LostType.INPUT, env.getListener());

    ImmutableList<LostInputRecord> lostInputRecords = createLostInputRecords(lostInputsByDigest);

    SetMultimap<ActionInput, Artifact> owners =
        calculateLostInputOwners(lostInputsByDigest.values(), metadataProvider);

    ImmutableList.Builder<ActionAnalysisMetadata> depsToRewind = ImmutableList.builder();
    RewindPlanResult rewindPlanResult;
    try (var ignored =
        AutoProfiler.profiled(
            "Preparing rewind plan for %d lost inputs of %s"
                .formatted(lostInputRecords.size(), failedAction.prettyPrint()),
            ProfilerTask.ACTION_REWINDING)) {
      rewindPlanResult =
          prepareRewindPlan(
              failedKey, failedActionDeps, lostInputsByDigest, owners, env, depsToRewind);
    }
    Reset rewindPlan = rewindPlanResult.reset;
    if (rewindPlan == null) {
      // Skips steps that manipulate persistent state, structures retained by this class and
      // information tracked by SkyframeActionExecutor, if a restart is needed.
      //
      // This includes counting if an input has been lost too many times.
      return rewindPlanResult;
    }

    checkIfActionLostInputTooManyTimes(
        failedKey, failedAction, lostInputRecords, lostInputsByDigest);

    if (shouldRecordRewindEventSample()) {
      rewindEventSamples.add(
          createLostInputRewindEvent(failedAction, rewindPlan, lostInputRecords));
    }

    if (e.isActionStartedEventAlreadyEmitted()) {
      env.getListener()
          .post(new ActionRewoundEvent(actionStartTimeNanos, BlazeClock.nanoTime(), failedAction));
    }
    skyframeActionExecutor.prepareForRewinding(failedKey, failedAction, depsToRewind.build());

    return rewindPlanResult;
  }

  private enum LostType {
    INPUT("inputs", Code.LOST_INPUT_REWINDING_DISABLED),
    OUTPUT("outputs", Code.LOST_OUTPUT_REWINDING_DISABLED);

    private final String description;
    private final Code codeWhenDisabled;

    LostType(String description, Code codeWhenDisabled) {
      this.description = description;
      this.codeWhenDisabled = codeWhenDisabled;
    }
  }

  private void checkRewindingEnabled(
      ImmutableSetMultimap<String, ActionInput> lostArtifacts,
      LostType lostType,
      ExtendedEventHandler listener)
      throws ActionRewindException {
    if (skyframeActionExecutor.rewindingEnabled()) {
      return;
    }
    if (skyframeActionExecutor.invocationRetriesEnabled()) {
      // If rewinding failed, Bazel may still be able to recover by retrying the invocation in
      // BlazeCommandDispatcher if retries are enabled. This requires emitting an event to inform
      // Bazel's remote module of the lost inputs.
      listener.post(new LostInputsEvent(lostArtifacts.keySet()));
      throw new FallbackToBuildRewindingException(
          lostArtifacts.entries().stream()
              .limit(MAX_LOST_INPUTS_RECORDED)
              .map(lost -> "%s (%s)".formatted(prettyPrint(lost.getValue()), lost.getKey()))
              .collect(
                  joining(
                      ", ",
                      "Lost %s no longer available remotely: ".formatted(lostType.description),
                      "")));
    }
    throw new GenericActionRewindException(
        "Unexpected lost %s (pass --rewind_lost_inputs to enable recovery): %s"
            .formatted(lostType.description, prettyPrint(lostArtifacts.values())),
        lostType.codeWhenDisabled);
  }

  private RewindPlanResult prepareRewindPlan(
      SkyKey failedKey,
      Set<SkyKey> failedKeyDeps,
      ImmutableSetMultimap<String, ActionInput> lostInputsByDigest,
      SetMultimap<ActionInput, Artifact> owners,
      Environment env,
      ImmutableList.Builder<ActionAnalysisMetadata> depsToRewind)
      throws InterruptedException {
    ImmutableList<ActionInput> lostInputs = lostInputsByDigest.values().asList();

    // This graph tracks which Skyframe nodes must be rewound and the dependency relationships
    // between them.
    MutableGraph<SkyKey> rewindGraph = Reset.newRewindGraphFor(failedKey);

    Set<DerivedArtifact> lostArtifacts =
        getLostInputOwningDirectDeps(failedKey, failedKeyDeps, lostInputs, owners);

    // 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();

    boolean missingDependencies = false;
    for (DerivedArtifact lostArtifact : lostArtifacts) {
      Map<ActionLookupData, ActionAnalysisMetadata> actionMap =
          getActionsForLostArtifact(lostArtifact, env);
      if (actionMap == null) {
        // Some deps of the artifact are not done. If allowSkyframeRestarts() is false, another
        // rewind must be in-flight, and there is no need to rewind the shared deps twice.
        //
        // TODO: b/420979437 - clearly document the scenario above.
        missingDependencies = true;
        continue;
      }
      ImmutableList<ActionAndLookupDataOrTemplate> newlyVisitedActions =
          addArtifactDepsAndGetNewlyVisitedActions(rewindGraph, lostArtifact, actionMap);

      // 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. 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(actionAnalysisMetadatas(newlyVisitedActions));
      switch (checkActions(
          newlyVisitedActions, env, rewindGraph, depsToRewind, nestedSetsForPropagatingActions)) {
        case SUCCESS:
          break;
        case MISSING_DEPENDENCIES:
          missingDependencies = true;
          break;
      }
    }

    // The default behavior, when allowSkyframeRestarts() is false, ignores missing dependencies.
    //
    // TODO: b/420979437 - clearly document why this is okay
    if (missingDependencies && allowSkyframeRestarts()) {
      // When this is reached, the calling SkyFunction will request a restart. It's possible that
      // the restarted SkyFunction won't observe a missing input and instead succeed without
      // rewinding. Then, the resulting deps could include transitive deps that aren't usually
      // present for that function.
      //
      // TODO: b/399637310 - consider using SkyKeyComputeState to avoid this.
      return new RewindPlanResult(/* reset= */ null);
    }

    // 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
    // short-circuits when a node is already in the rewind graph.
    ArtifactNestedSetKey.addNestedSetPathsToRewindGraph(
        rewindGraph, failedKey, failedKeyDeps, lostArtifacts);

    return new RewindPlanResult(Reset.of(rewindGraph));
  }

  /**
   * Creates a {@link Reset} to recover from undone indirect inputs that are unavailable due to
   * unsuccessful rewinding.
   *
   * <p>Undone direct dependencies are handled by Skyframe (see {@link
   * com.google.devtools.build.skyframe.SkyFunctionEnvironment.UndonePreviouslyRequestedDeps}). This
   * method only exists for artifacts whose {@link Artifact#key} is <em>not</em> a direct dependency
   * of {@code failedKey} because the artifact is behind an {@link ArtifactNestedSetKey}.
   *
   * <p>Used when an indirect dependency {@link Artifact#key} was rewound and completed with an
   * error, but intermediate nested set nodes were never rewound, resulting in an inconsistent state
   * where successful nodes depend on a node in error.
   *
   * <p>The returned {@link Reset} contains the {@link Artifact#key}, but that is expected to
   * already be in error, and attempting to rewind an error is no-op.
   */
  public Reset patchNestedSetGraphToPropagateError(
      ActionLookupData failedKey,
      Action failedAction,
      ImmutableSet<Artifact> undoneInputs,
      ImmutableSet<SkyKey> failedActionDeps) {
    checkState(
        skyframeActionExecutor.rewindingEnabled(), "Unexpected undone inputs: %s", undoneInputs);
    MutableGraph<SkyKey> rewindGraph = Reset.newRewindGraphFor(failedKey);
    ArtifactNestedSetKey.addNestedSetPathsToRewindGraph(
        rewindGraph, failedKey, failedActionDeps, undoneInputs);
    for (Artifact input : undoneInputs) {
      checkState(
          rewindGraph.nodes().contains(Artifact.key(input)),
          "Cannot find input %s under any nested set deps of %s",
          input,
          failedKey);
    }

    // An undone input may be observed either during input checking (before attempting action
    // execution) or during lost input handling (after attempting action execution). In the latter
    // case, it is necessary to obsolete the ActionExecutionState so that after rewinding, we will
    // check inputs again and discover the propagated exception. This call is a no-op in the former
    // case, since there is no ActionExecutionState to obsolete.
    skyframeActionExecutor.prepareForRewinding(
        failedKey, failedAction, /* depsToRewind= */ ImmutableList.of());

    return Reset.of(rewindGraph);
  }

  /**
   * Logs the first N action rewind events and clears the history of failed actions' lost inputs and
   * rewind plans.
   */
  public void reset(ExtendedEventHandler eventHandler) {
    PostableActionRewindingStats rewindingStats =
        new PostableActionRewindingStats(
            /* lostInputsCount= */ lostInputsCount.getAndSet(0),
            /* lostOutputsCount= */ lostOutputsCount.getAndSet(0),
            /* rewoundActionCount= */ skyframeActionExecutor.getRewoundActionCount(),
            /* sameActionLostInputsCount= */ sameActionLostInputsCount.getAndSet(0),
            /* sameTopLevelKeyLostOutputsCount= */ sameTopLevelKeyLostOutputsCount.getAndSet(0),
            /* actionRewindEvents= */ ImmutableList.copyOf(rewindEventSamples));
    eventHandler.post(rewindingStats);
    currentBuildLostInputRecords = new ConcurrentHashMap<>();
    rewindEventSamples.clear();
    rewindEventSampleCounter.set(0);
  }

  private void checkIfTopLevelOutputLostTooManyTimes(
      TopLevelActionLookupKeyWrapper failedKey,
      ImmutableList<LostInputRecord> currentAttemptLostOutputRecords,
      ImmutableSetMultimap<String, ActionInput> lostOutputsByDigest)
      throws ActionRewindException {
    lostOutputsCount.addAndGet(currentAttemptLostOutputRecords.size());

    Multiset<LostInputRecord> historyForThisTopLevelKey =
        currentBuildLostInputRecords.computeIfAbsent(
            failedKey, k -> ConcurrentHashMultiset.create());
    if (!historyForThisTopLevelKey.isEmpty()) {
      sameTopLevelKeyLostOutputsCount.incrementAndGet();
    }

    for (LostInputRecord lostInputRecord : currentAttemptLostOutputRecords) {
      String digest = lostInputRecord.lostInputDigest();
      int losses = historyForThisTopLevelKey.add(lostInputRecord, /* occurrences= */ 1) + 1;
      if (losses > MAX_REPEATED_LOST_INPUTS) {
        ActionInput lostOutput =
            Iterables.find(
                lostOutputsByDigest.get(digest),
                lost -> lost.getExecPathString().equals(lostInputRecord.lostInputPath));
        ActionRewindException e =
            new GenericActionRewindException(
                String.format(
                    "Lost output %s (digest %s), and rewinding was ineffective after %d attempts.",
                    prettyPrint(lostOutput), digest, MAX_REPEATED_LOST_INPUTS),
                ActionRewinding.Code.LOST_OUTPUT_TOO_MANY_TIMES);
        bugReporter.sendBugReport(e);
        throw e;
      } else if (losses > 1) {
        logger.atWarning().log(
            "Lost output again (losses=%s, output=%s, digest=%s, failedKey=%s)",
            losses, lostOutputsByDigest.get(digest), digest, failedKey);
      }
    }
  }

  private static String prettyPrint(Collection<ActionInput> inputs) {
    return inputs.stream().map(ActionRewindStrategy::prettyPrint).collect(joining(","));
  }

  private static String prettyPrint(ActionInput input) {
    return input instanceof Artifact a ? a.prettyPrint() : input.getExecPathString();
  }

  /**
   * Checks if the action identified by {@code failedKey} has lost any of the inputs in {@code
   * currentAttemptLostInputRecords} too many times.
   *
   * @throws ActionRewindException if any lost input has been seen by {@code failedKey} as lost
   *     before too many times
   */
  private void checkIfActionLostInputTooManyTimes(
      ActionLookupData failedKey,
      Action failedAction,
      ImmutableList<LostInputRecord> currentAttemptLostInputRecords,
      ImmutableSetMultimap<String, ActionInput> lostInputsByDigest)
      throws ActionRewindException {
    lostInputsCount.addAndGet(currentAttemptLostInputRecords.size());

    Multiset<LostInputRecord> historyForThisAction =
        currentBuildLostInputRecords.computeIfAbsent(
            failedKey, k -> ConcurrentHashMultiset.create());
    if (!historyForThisAction.isEmpty()) {
      sameActionLostInputsCount.incrementAndGet();
    }

    for (LostInputRecord lostInputRecord : currentAttemptLostInputRecords) {
      // The same action losing the same input more than once is unexpected [*]. The action should
      // have waited until the depended-on action which generates the lost input is (re)run before
      // trying again.
      //
      // Note that we could enforce a stronger check: if action A, which depends on an input N
      // previously detected as lost (by any action, not just A), discovers that N is still lost,
      // and action A started after the re-evaluation of N's generating action, then something has
      // gone wrong. Administering that check would be more complex (e.g., the start/completion
      // times of actions would need tracking), so we punt on it for now.
      //
      // [*], TODO(b/123993876): To mitigate a race condition (believed to be) caused by
      // non-topological Skyframe dirtying of depended-on nodes, this check fails the build only if
      // the same input is repeatedly lost.
      String digest = lostInputRecord.lostInputDigest();
      int losses = historyForThisAction.add(lostInputRecord, /* occurrences= */ 1) + 1;
      if (losses > MAX_REPEATED_LOST_INPUTS) {
        // This ensures coalesced shared actions aren't orphaned.
        skyframeActionExecutor.prepareForRewinding(
            failedKey, failedAction, /* depsToRewind= */ ImmutableList.of());
        ActionInput lostInput =
            Iterables.find(
                lostInputsByDigest.get(digest),
                lost -> lost.getExecPathString().equals(lostInputRecord.lostInputPath));
        String message =
            String.format(
                "lost input too many times (#%s) for the same action. lostInput: %s, "
                    + "lostInput digest: %s, failedAction: %.10000s",
                losses, lostInput, digest, failedAction);
        ActionRewindException e =
            new GenericActionRewindException(
                message, ActionRewinding.Code.LOST_INPUT_TOO_MANY_TIMES);
        bugReporter.sendBugReport(e);
        throw e;
      } else if (losses > 1) {
        logger.atInfo().log(
            "lost input again (#%s) for the same action. lostInput: %s, "
                + "lostInput digest: %s, failedAction: %.10000s",
            losses, lostInputsByDigest.get(digest), digest, failedAction);
      }
    }
  }

  private static ImmutableList<LostInputRecord> createLostInputRecords(
      ImmutableSetMultimap<String, ActionInput> lostInputsByDigest) {
    return lostInputsByDigest.entries().stream()
        .map(e -> new LostInputRecord(e.getKey(), e.getValue().getExecPathString()))
        .collect(toImmutableList());
  }

  /**
   * Calculates the aggregation artifacts (tree artifacts, filesets, runfiles) that own {@code
   * lostInputs}.
   */
  private static SetMultimap<ActionInput, Artifact> calculateLostInputOwners(
      ImmutableCollection<ActionInput> lostInputs, InputMetadataProvider inputArtifactData) {
    Set<ActionInput> lostInputsAndOwners = new HashSet<>();
    SetMultimap<ActionInput, Artifact> owners = HashMultimap.create();
    for (ActionInput lostInput : lostInputs) {
      lostInputsAndOwners.add(lostInput);
      if (lostInput instanceof Artifact artifact && artifact.hasParent()) {
        lostInputsAndOwners.add(artifact.getParent());
        owners.put(artifact, artifact.getParent());
      }
    }

    inputArtifactData
        .getFilesets()
        .forEach(
            (fileset, outputTree) -> {
              for (FilesetOutputSymlink link : outputTree.symlinks()) {
                if (lostInputsAndOwners.contains(link.target())) {
                  lostInputsAndOwners.add(fileset);
                  owners.put(link.target(), fileset);
                }
              }
            });

    // Runfiles trees may contain tree artifacts and filesets, but not vice versa. Runfiles are
    // processed last to ensure that any lost input owning tree artifacts and filesets are already
    // in lostInputsAndOwners.
    for (RunfilesTree runfilesTree : inputArtifactData.getRunfilesTrees()) {
      Artifact runfilesArtifact = (Artifact) inputArtifactData.getInput(runfilesTree.getExecPath());
      checkState(runfilesArtifact.isRunfilesTree(), runfilesArtifact);

      RunfilesArtifactValue runfilesValue = inputArtifactData.getRunfilesMetadata(runfilesArtifact);
      for (Artifact artifact : runfilesValue.getAllArtifacts()) {
        if (lostInputsAndOwners.contains(artifact)) {
          owners.put(artifact, runfilesArtifact);
        }
      }
    }

    return owners;
  }

  private Set<DerivedArtifact> getLostInputOwningDirectDeps(
      SkyKey failedKey,
      Set<SkyKey> failedKeyDeps,
      ImmutableList<ActionInput> lostInputs,
      SetMultimap<ActionInput, Artifact> owners) {
    // 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;

      Set<Artifact> directOwners = owners.get(lostInput);
      for (Artifact directOwner : directOwners) {
        checkDerived(directOwner);

        // Rewinding must invalidate all Skyframe paths from the failed action to the action which
        // generates the lost input. Intermediate nodes not on the shortest path to that action may
        // have values that depend on the output of that action. If these intermediate nodes are not
        // invalidated, then their values may become stale. Therefore, this method collects not only
        // the first action dep associated with the lost input, but all of them.

        Set<Artifact> transitiveOwners = owners.get(directOwner);
        for (Artifact transitiveOwner : transitiveOwners) {
          checkDerived(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 runfiles tree) that the
            // action directly depends on.
            lostInputOwningDirectDeps.add((DerivedArtifact) transitiveOwner);
            foundLostInputDepOwner = true;
          }
        }

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

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

        lostInputOwningDirectDeps.add((DerivedArtifact) lostInput);
        foundLostInputDepOwner = true;
      }

      if (!foundLostInputDepOwner) {
        // Rewinding can't do anything about a lost input that can't be associated with a direct dep
        // of the failed action. In this case, try resetting the failed action (and no other deps)
        // just in case that helps. If it does not help, then eventually the action will fail in
        // checkIfActionLostInputTooManyTimes.
        bugReporter.sendNonFatalBugReport(
            new IllegalStateException(
                String.format(
                    "Lost input not a dep of the failed action and can't be associated with such"
                        + " a dep. lostInput: %s, owners: %s, failedKey: %s",
                    lostInput, owners, failedKey)));
      }
    }
    return lostInputOwningDirectDeps;
  }

  private static void checkDerived(Artifact artifact) {
    checkState(!artifact.isSourceArtifact(), "Unexpected source artifact: %s", artifact);
  }

  private enum CheckActionsStatus {
    MISSING_DEPENDENCIES,
    SUCCESS
  }

  /**
   * Looks at each action in {@code actionsToCheck} and determines whether additional artifacts or
   * actions need to be rewound. If this finds more actions to rewind, those actions are recursively
   * checked too.
   */
  private CheckActionsStatus checkActions(
      ImmutableList<ActionAndLookupDataOrTemplate> actionsToCheck,
      Environment env,
      MutableGraph<SkyKey> rewindGraph,
      ImmutableList.Builder<ActionAnalysisMetadata> depsToRewind,
      SetMultimap<SkyKey, ArtifactNestedSetKey> nestedSetDeps)
      throws InterruptedException {
    boolean missingDependencies = false;
    var uncheckedActions = new ArrayDeque<ActionAndLookupData>(actionsToCheck.size());
    filterActionAndLookupDataTo(actionsToCheck, uncheckedActions);
    while (!uncheckedActions.isEmpty()) {
      ActionAndLookupData actionAndLookupData = uncheckedActions.removeFirst();
      ActionLookupData actionKey = actionAndLookupData.lookupData();
      Action action = actionAndLookupData.actionAnalysisMetadata();
      ArrayList<DerivedArtifact> artifactsToCheck = new ArrayList<>();
      ArrayList<ActionLookupData> newlyDiscoveredActions = new ArrayList<>();

      if (action.mayInsensitivelyPropagateInputs()) {
        // 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,
            nestedSetDeps);
      }

      for (ActionLookupData actionLookupData : newlyDiscoveredActions) {
        Action additionalAction =
            ActionUtils.getActionForLookupData(
                env, actionLookupData, /* crashIfActionOwnerMissing= */ !allowSkyframeRestarts());
        if (additionalAction == null) {
          missingDependencies = true;
          continue;
        }
        depsToRewind.add(additionalAction);
        uncheckedActions.add(new ActionAndLookupData(actionLookupData, additionalAction));
      }
      for (DerivedArtifact artifact : artifactsToCheck) {
        Map<ActionLookupData, ActionAnalysisMetadata> actionMap =
            getActionsForLostArtifact(artifact, env);
        if (actionMap == null) {
          missingDependencies = true;
          continue;
        }
        ImmutableList<ActionAndLookupDataOrTemplate> newlyVisitedActions =
            addArtifactDepsAndGetNewlyVisitedActions(rewindGraph, artifact, actionMap);
        depsToRewind.addAll(actionAnalysisMetadatas(newlyVisitedActions));
        filterActionAndLookupDataTo(newlyVisitedActions, uncheckedActions);
      }
    }
    return missingDependencies
        ? CheckActionsStatus.MISSING_DEPENDENCIES
        : CheckActionsStatus.SUCCESS;
  }

  private static void filterActionAndLookupDataTo(
      ImmutableList<ActionAndLookupDataOrTemplate> actions, ArrayDeque<ActionAndLookupData> queue) {
    for (ActionAndLookupDataOrTemplate action : actions) {
      switch (action) {
        case ActionAndLookupData actionAndLookupData -> queue.add(actionAndLookupData);
        case TemplateWrapper ignored -> {}
      }
    }
  }

  /**
   * For a propagating {@code action} with key {@code actionKey}, add its generated inputs' keys to
   * {@code rewindGraph}, add edges from {@code actionKey} to those keys, add any {@link Artifact}s
   * to {@code newlyVisitedArtifacts}, and add any {@link ActionLookupData}s to {@code
   * newlyVisitedActions}.
   */
  private void addPropagatingActionDepsAndGetNewlyVisitedArtifactsAndActions(
      MutableGraph<SkyKey> rewindGraph,
      ActionLookupData actionKey,
      Action action,
      ArrayList<DerivedArtifact> newlyVisitedArtifacts,
      ArrayList<ActionLookupData> newlyVisitedActions,
      SetMultimap<SkyKey, ArtifactNestedSetKey> nestedSetDeps) {

    for (Artifact input : action.getInputs().toList()) {
      if (input.isSourceArtifact()) {
        continue;
      }
      SkyKey artifactKey = Artifact.key(input);
      // Rewinding all derived inputs of propagating actions is overkill. Preferably, we'd want to
      // only rewind the inputs which correspond to the known lost outputs. The information to do
      // this is probably present in the data available to #prepareRewindPlan.
      //
      // Rewinding is expected to be rare, so refining this may not be necessary.
      boolean newlyVisited = rewindGraph.addNode(artifactKey);
      if (newlyVisited) {
        if (artifactKey instanceof Artifact) {
          newlyVisitedArtifacts.add((DerivedArtifact) artifactKey);
        } else if (artifactKey instanceof ActionLookupData actionLookupData) {
          newlyVisitedActions.add(actionLookupData);
        }
      }
      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",
    // should have no overlap. Log a bug report if we see such an action:
    if (action.discoversInputs()) {
      bugReporter.sendBugReport(
          new IllegalStateException(
              String.format(
                  "Action insensitively propagates and discovers inputs. actionKey: %s, action: "
                      + "%.10000s",
                  actionKey, action)));
    }
  }

  /**
   * For an artifact {@code artifact} with generating actions (and their associated {@link
   * ActionLookupData}) {@code actionMap}, add those actions' keys to {@code rewindGraph} and add
   * edges from {@code artifact} to those keys.
   *
   * <p>Returns a list of key+action pairs for each action whose key was newly added to the graph.
   */
  private static ImmutableList<ActionAndLookupDataOrTemplate>
      addArtifactDepsAndGetNewlyVisitedActions(
          MutableGraph<SkyKey> rewindGraph,
          Artifact artifact,
          Map<ActionLookupData, ActionAnalysisMetadata> actionMap) {

    ImmutableList.Builder<ActionAndLookupDataOrTemplate> newlyVisitedActions =
        ImmutableList.builderWithExpectedSize(actionMap.size());
    SkyKey artifactKey = Artifact.key(artifact);
    for (Map.Entry<ActionLookupData, ActionAnalysisMetadata> actionEntry : actionMap.entrySet()) {
      ActionLookupData actionKey = actionEntry.getKey();
      if (rewindGraph.addNode(actionKey)) {
        newlyVisitedActions.add(
            switch (actionEntry.getValue()) {
              case Action action -> new ActionAndLookupData(actionKey, action);
              case ActionTemplate<?> template -> new TemplateWrapper(template);
              default ->
                  throw new IllegalStateException(
                      "ActionAnalysisMetadata " + actionEntry.getValue() + " of unexpected type");
            });
      }
      if (!artifactKey.equals(actionKey)) {
        rewindGraph.putEdge(artifactKey, actionKey);
      }
    }
    return newlyVisitedActions.build();
  }

  /**
   * Returns the map of {@code lostInput}'s execution-phase dependencies (i.e. generating actions),
   * keyed by their {@link ActionLookupData} keys, or {@code null} if any of those dependencies are
   * not done.
   */
  @Nullable
  private Map<ActionLookupData, ActionAnalysisMetadata> getActionsForLostArtifact(
      DerivedArtifact lostInput, Environment env) throws InterruptedException {
    ImmutableSet<ActionLookupData> actionExecutionDeps = getActionExecutionDeps(lostInput, env);
    if (actionExecutionDeps == null) {
      return null;
    }

    Map<ActionLookupData, ActionAnalysisMetadata> actions =
        Maps.newHashMapWithExpectedSize(actionExecutionDeps.size());
    boolean missingAction = false;
    for (ActionLookupData dep : actionExecutionDeps) {
      ActionAnalysisMetadata actionAnalysisMetadata = getActionAnalysisMetadata(env, dep);
      if (actionAnalysisMetadata == null) {
        missingAction = true;
        continue;
      }
      actions.put(dep, actionAnalysisMetadata);
    }
    if (missingAction) {
      return null;
    }
    return actions;
  }

  @Nullable
  private ActionAnalysisMetadata getActionAnalysisMetadata(
      Environment env, ActionLookupData actionLookupData) throws InterruptedException {
    ActionLookupValue actionLookupValue =
        ArtifactFunction.getActionLookupValue(
            actionLookupData.getActionLookupKey(),
            env,
            /* crashIfActionOwnerMissing= */ !allowSkyframeRestarts());
    if (actionLookupValue == null) {
      return null;
    }
    return actionLookupValue.getActions().get(actionLookupData.getActionIndex());
  }

  /**
   * Returns the set of {@code lostInput}'s execution-phase dependencies (i.e. generating actions),
   * or {@code null} if any of those dependencies are not done.
   */
  @Nullable
  private ImmutableSet<ActionLookupData> getActionExecutionDeps(
      DerivedArtifact lostInput, Environment env) throws InterruptedException {
    if (!lostInput.isTreeArtifact()) {
      return ImmutableSet.of(lostInput.getGeneratingActionKey());
    }
    ArtifactDependencies artifactDependencies =
        ArtifactDependencies.discoverDependencies(
            lostInput, env, /* crashIfActionOwnerMissing= */ !allowSkyframeRestarts());
    if (artifactDependencies == null) {
      return null;
    }

    if (!artifactDependencies.isTemplateActionForTreeArtifact()) {
      return ImmutableSet.of(lostInput.getGeneratingActionKey());
    }

    // This ignores the ActionTemplateExpansionKey dependency of the template artifact because we
    // expect to never need to rewind that.
    ImmutableList<ActionLookupData> actionTemplateExpansionKeys =
        artifactDependencies.getActionTemplateExpansionKeys(env);
    if (actionTemplateExpansionKeys == null) {
      return null;
    }
    return ImmutableSet.<ActionLookupData>builder()
        .add(lostInput.getGeneratingActionKey())
        .addAll(actionTemplateExpansionKeys)
        .build();
  }

  private boolean shouldRecordRewindEventSample() {
    return rewindEventSampleCounter.getAndIncrement() < MAX_ACTION_REWIND_EVENTS;
  }

  private boolean allowSkyframeRestarts() {
    return cachingDependenciesSupplier.get().mode().isRetrievalEnabled();
  }

  private static ActionRewindEvent createLostOutputRewindEvent(
      TopLevelActionLookupKeyWrapper failedKey,
      Reset rewindPlan,
      ImmutableList<LostInputRecord> lostOutputRecords) {
    return createRewindEventBuilder(rewindPlan, lostOutputRecords)
        .setTopLevelActionLookupKeyDescription(failedKey.actionLookupKey().toString())
        .build();
  }

  private static ActionRewindEvent createLostInputRewindEvent(
      Action failedAction, Reset rewindPlan, ImmutableList<LostInputRecord> lostInputRecords) {
    return createRewindEventBuilder(rewindPlan, lostInputRecords)
        .setActionDescription(
            ActionDescription.newBuilder()
                .setType(failedAction.getMnemonic())
                .setRuleLabel(failedAction.getOwner().getLabel().toString()))
        .build();
  }

  private static ActionRewindEvent.Builder createRewindEventBuilder(
      Reset rewindPlan, ImmutableList<LostInputRecord> lostInputRecords) {
    return ActionRewindEvent.newBuilder()
        .addAllLostInputs(
            lostInputRecords.stream()
                .limit(MAX_LOST_INPUTS_RECORDED)
                .map(
                    lostInputRecord ->
                        LostInput.newBuilder()
                            .setPath(lostInputRecord.lostInputPath())
                            .setDigest(lostInputRecord.lostInputDigest())
                            .build())
                .collect(toImmutableList()))
        .setTotalLostInputsCount(lostInputRecords.size())
        .setInvalidatedNodesCount(rewindPlan.rewindGraph().nodes().size());
  }

  private record LostInputRecord(String lostInputDigest, String lostInputPath) {
    LostInputRecord {
      requireNonNull(lostInputDigest, "lostInputDigest");
      requireNonNull(lostInputPath, "lostInputPath");
    }
  }

  private sealed interface ActionAndLookupDataOrTemplate
      permits TemplateWrapper, ActionAndLookupData {
    ActionAnalysisMetadata actionAnalysisMetadata();
  }

  private record TemplateWrapper(ActionTemplate<?> actionAnalysisMetadata)
      implements ActionAndLookupDataOrTemplate {
    private TemplateWrapper {
      requireNonNull(actionAnalysisMetadata, "actionAnalysisMetadata");
    }
  }

  private record ActionAndLookupData(ActionLookupData lookupData, Action actionAnalysisMetadata)
      implements ActionAndLookupDataOrTemplate {
    private ActionAndLookupData {
      requireNonNull(lookupData, "lookupData");
      requireNonNull(actionAnalysisMetadata, "actionAnalysisMetadata");
    }
  }

  private static List<ActionAnalysisMetadata> actionAnalysisMetadatas(
      List<ActionAndLookupDataOrTemplate> newlyVisitedActions) {
    return Lists.transform(
        newlyVisitedActions, ActionAndLookupDataOrTemplate::actionAnalysisMetadata);
  }
}
