[GOOGLE][Skymeld] Make orphaned artifact detection available for Skym…

…eld.

Internal change.

PiperOrigin-RevId: 588730288
Change-Id: Ib9e13e4271e2ec555d14aa743acca33c1b66b814

src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java

@@ -140,6 +140,7 @@ public class ExecutionTool {
   private BlazeExecutor executor;
   private final ActionInputPrefetcher prefetcher;
   private final ImmutableSet<ExecutorLifecycleListener> executorLifecycleListeners;
+
   private final SpawnStrategyRegistry spawnStrategyRegistry;
   private final ModuleActionContextRegistry actionContextRegistry;
 
@@ -349,7 +350,8 @@ public void prepareForExecution(Stopwatch executionTimer)
     for (ExecutorLifecycleListener executorLifecycleListener : executorLifecycleListeners) {
       try (SilentCloseable c =
           Profiler.instance().profile(executorLifecycleListener + ".executionPhaseStarting")) {
-        executorLifecycleListener.executionPhaseStarting(null, () -> null, null);
+        executorLifecycleListener.executionPhaseStarting(
+            null, () -> null, skyframeExecutor.getEphemeralCheckIfOutputConsumed());
       }
     }
     try (SilentCloseable c = Profiler.instance().profile("configureResourceManager")) {

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

@@ -22,8 +22,6 @@
 import com.google.devtools.build.skyframe.ExecutionPhaseSkyKey;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.Set;
 
 /**
@@ -69,16 +67,25 @@ private ArtifactNestedSetKey(Object[] children) {
    * Returns a list of this key's direct dependencies, including {@link Artifact#key} for leaves and
    * {@link ArtifactNestedSetKey} for non-leaves.
    */
-  public List<SkyKey> getDirectDepKeys() {
-    List<SkyKey> depKeys = new ArrayList<>(children.length);
+  public ImmutableList<SkyKey> getDirectDepKeys() {
+    ImmutableList.Builder<SkyKey> depKeys = ImmutableList.builderWithExpectedSize(children.length);
     for (Object child : children) {
       if (child instanceof Artifact) {
         depKeys.add(Artifact.key((Artifact) child));
       } else {
         depKeys.add(createInternal((Object[]) child));
       }
     }
-    return depKeys;
+    return depKeys.build();
+  }
+
+  /** Applies a consumer function to the direct artifacts of this nested set. */
+  public void applyToDirectArtifacts(DirectArtifactConsumer function) throws InterruptedException {
+    for (Object child : children) {
+      if (child instanceof Artifact) {
+        function.accept((Artifact) child);
+      }
+    }
   }
 
   @Override
@@ -175,4 +182,10 @@ public boolean equals(Object that) {
   public String toString() {
     return String.format("ArtifactNestedSetKey[%s]@%s", children.length, hashCode());
   }
+
+  /** A consumer to be applied to each direct artifact. */
+  @FunctionalInterface
+  public interface DirectArtifactConsumer {
+    void accept(Artifact artifact) throws InterruptedException;
+  }
 }

src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD

@@ -52,7 +52,9 @@ java_library(
 
 java_library(
     name = "artifact_nested_set_key",
-    srcs = ["ArtifactNestedSetKey.java"],
+    srcs = [
+        "ArtifactNestedSetKey.java",
+    ],
     deps = [
         ":nestedset",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",

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

@@ -159,6 +159,7 @@ public final class ActionExecutionFunction implements SkyFunction {
   private final BlazeDirectories directories;
   private final Supplier<TimestampGranularityMonitor> tsgm;
   private final BugReporter bugReporter;
+  private final Supplier<ConsumedArtifactsTracker> consumedArtifactsTrackerSupplier;
 
   // TODO(b/314282963): Remove this after the rollout.
   private final Supplier<Boolean> clearNestedSetAfterActionExecution;
@@ -170,6 +171,7 @@ public ActionExecutionFunction(
       BlazeDirectories directories,
       Supplier<TimestampGranularityMonitor> tsgm,
       BugReporter bugReporter,
+      Supplier<ConsumedArtifactsTracker> consumedArtifactsTrackerSupplier,
       Supplier<Boolean> clearNestedSetAfterActionExecution) {
     this.actionRewindStrategy = checkNotNull(actionRewindStrategy);
     this.skyframeActionExecutor = checkNotNull(skyframeActionExecutor);
@@ -178,6 +180,7 @@ public ActionExecutionFunction(
     this.tsgm = checkNotNull(tsgm);
     this.bugReporter = checkNotNull(bugReporter);
     this.clearNestedSetAfterActionExecution = clearNestedSetAfterActionExecution;
+    this.consumedArtifactsTrackerSupplier = consumedArtifactsTrackerSupplier;
   }
 
   @Override
@@ -291,7 +294,13 @@ private SkyValue computeInternal(SkyKey skyKey, Environment env)
 
     if (!state.hasArtifactData()) {
       Iterable<SkyKey> depKeys =
-          getInputDepKeys(allInputs, action.getSchedulingDependencies(), state);
+          getInputDepKeys(
+              action,
+              consumedArtifactsTrackerSupplier.get(),
+              allInputs,
+              action.getSchedulingDependencies(),
+              state);
+
       SkyframeLookupResult inputDeps = env.getValuesAndExceptions(depKeys);
       if (previousExecution == null) {
         // Do we actually need to find our metadata?
@@ -363,7 +372,12 @@ private SkyValue computeInternal(SkyKey skyKey, Environment env)
           actionStartTime,
           env,
           allInputs,
-          getInputDepKeys(allInputs, action.getSchedulingDependencies(), state),
+          getInputDepKeys(
+              action,
+              /* consumedArtifactsTracker= */ null,
+              allInputs,
+              action.getSchedulingDependencies(),
+              state),
           state);
     } catch (ActionExecutionException e) {
       // In this case we do not report the error to the action reporter because we have already
@@ -385,20 +399,75 @@ private SkyValue computeInternal(SkyKey skyKey, Environment env)
       action.getInputs().clearMemo();
       allInputs.clearMemo();
     }
+    // After the action execution is finalized, unregister the outputs from the consumed set to save
+    // memory.
+    // Note: This can theoretically lead to infinite action rewinding if we're unlucky enough.
+    // Consider an action foo whose outputs A and B are needed by 2 separate actions consumerA and
+    // consumerB. If these 2 actions trigger rewinding alternately, at the correct timing, e.g.:
+    // 1. consumerA requests for A. A is registered. foo produces only A since B isn't registered. A
+    // is de-registered. consumerA isn't executed yet.
+    // 2. consumerB requests for B. B is registered. foo is rewound and produces only B since A
+    // isn't registered. B is de-registered. consumerB isn't executed yet.
+    // 3. Before consumerA enters execution, A falls out of the CAS. consumerA sees that A is
+    // missing and triggers rewinding for A. Repeat step (1).
+    // 4. Before consumerB enters execution, B falls out of the CAS. consumerB sees that B is
+    // missing and triggers rewinding for B. Repeat step (2).
+    if (consumedArtifactsTrackerSupplier.get() != null) {
+      consumedArtifactsTrackerSupplier
+          .get()
+          .unregisterOutputsAfterExecutionDone(action.getOutputs());
+    }
+
     return result;
   }
 
   private static Iterable<SkyKey> getInputDepKeys(
+      Action action,
+      ConsumedArtifactsTracker consumedArtifactsTracker,
       NestedSet<Artifact> allInputs,
       NestedSet<Artifact> schedulingDependencies,
-      InputDiscoveryState state) {
+      InputDiscoveryState state)
+      throws InterruptedException {
+
+    // Register the action's inputs and scheduling deps as "consumed" in the build.
+    // As a general rule, we do it before requesting for the evaluation of these artifacts. This
+    // would provide a good estimate of which outputs are consumed.
+    //
+    // The exception to this rule is middleman actions: it's possible that the output of a
+    // middleman action is only reachable via middleman actions in the action graph. In that case,
+    // we don't want to store any of the underlying artifacts, until we've discovered that there's
+    // a non-middleman action in its path.
+    if (!state.checkedForConsumedArtifactRegistration && consumedArtifactsTracker != null) {
+      // Special case: middleman actions.
+      // Delay registering the artifacts under this middleman action until we know that the
+      // middleman artifact itself is consumed by other non-middleman actions.
+      if (action.getActionType().isMiddleman()) {
+        consumedArtifactsTracker.skipRegisteringArtifactsUnderMiddleman(action.getPrimaryOutput());
+        // Skip the ArtifactNestedSetFunction altogether for this special case. Any artifact
+        // requested via ArtifactNestedSetFunction would be registered as consumed.
+        return Iterables.concat(
+            Artifact.keys(allInputs.toList()), Artifact.keys(schedulingDependencies.toList()));
+      } else {
+        // Only registering the leaves here, since the Artifacts under non-leaves will be registered
+        // in ArtifactNestedSetFunction. Similarly for the non-singleton Scheduling Dependencies.
+        for (Artifact input : allInputs.getLeaves()) {
+          consumedArtifactsTracker.registerConsumedArtifact(input);
+        }
+        if (schedulingDependencies.isSingleton()) {
+          consumedArtifactsTracker.registerConsumedArtifact(schedulingDependencies.getSingleton());
+        }
+      }
+      state.checkedForConsumedArtifactRegistration = true;
+    }
+
     // We "unwrap" the NestedSet and evaluate the first layer of direct Artifacts here in order to
     // save memory:
     // - This top layer costs 1 extra ArtifactNestedSetKey node.
     // - It's uncommon that 2 actions share the exact same set of inputs
     //   => the top layer offers little in terms of reusability.
     // More details: b/143205147.
     FluentIterable<SkyKey> result = FluentIterable.from(Artifact.keys(allInputs.getLeaves()));
+
     if (schedulingDependencies.isSingleton()) {
       result = result.append(Artifact.key(schedulingDependencies.getSingleton()));
     }
@@ -1486,6 +1555,8 @@ static class InputDiscoveryState implements SkyKeyComputeState {
     boolean preparedInputDiscovery = false;
     boolean actionInputCollectedEventSent = false;
 
+    boolean checkedForConsumedArtifactRegistration = false;
+
     /**
      * Stores the ArtifactNestedSetKeys created from the inputs of this actions. Objective: avoid
      * creating a new ArtifactNestedSetKey for the same NestedSet each time we run

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.collect.nestedset.ArtifactNestedSetKey;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
@@ -25,7 +26,7 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.SkyframeLookupResult;
-import java.util.List;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -48,11 +49,22 @@
  */
 final class ArtifactNestedSetFunction implements SkyFunction {
 
+  private final Supplier<ConsumedArtifactsTracker> consumedArtifactsTrackerSupplier;
+
+  ArtifactNestedSetFunction(Supplier<ConsumedArtifactsTracker> consumedArtifactsTrackerSupplier) {
+    this.consumedArtifactsTrackerSupplier = consumedArtifactsTrackerSupplier;
+  }
+
   @Override
   @Nullable
   public SkyValue compute(SkyKey skyKey, Environment env)
       throws InterruptedException, ArtifactNestedSetFunctionException {
-    List<SkyKey> depKeys = ((ArtifactNestedSetKey) skyKey).getDirectDepKeys();
+    ArtifactNestedSetKey artifactNestedSetKey = (ArtifactNestedSetKey) skyKey;
+    if (consumedArtifactsTrackerSupplier.get() != null) {
+      artifactNestedSetKey.applyToDirectArtifacts(
+          (x) -> consumedArtifactsTrackerSupplier.get().registerConsumedArtifact(x));
+    }
+    ImmutableList<SkyKey> depKeys = artifactNestedSetKey.getDirectDepKeys();
     SkyframeLookupResult depsEvalResult = env.getValuesAndExceptions(depKeys);
 
     NestedSetBuilder<Pair<SkyKey, Exception>> transitiveExceptionsBuilder =

src/main/java/com/google/devtools/build/lib/skyframe/BUILD

@@ -115,6 +115,7 @@ java_library(
         ":configured_target_key",
         ":configured_target_progress_receiver",
         ":configured_value_creation_exception",
+        ":consumed_artifacts_tracker",
         ":containing_package_lookup_function",
         ":containing_package_lookup_value",
         ":coverage_report_value",
@@ -124,6 +125,7 @@ java_library(
         ":diff_awareness_manager",
         ":directory_listing_function",
         ":directory_listing_state_value",
+        ":ephemeral_check_if_output_consumed",
         ":exclusive_test_build_driver_value",
         ":execution_finished_event",
         ":file_function",
@@ -700,6 +702,7 @@ java_library(
     deps = [
         ":artifact_function",
         ":artifact_nested_set_value",
+        ":consumed_artifacts_tracker",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset:artifact_nested_set_key",
@@ -1108,6 +1111,18 @@ java_library(
     ],
 )
 
+java_library(
+    name = "consumed_artifacts_tracker",
+    srcs = ["ConsumedArtifactsTracker.java"],
+    deps = [
+        ":ephemeral_check_if_output_consumed",
+        "//src/main/java/com/google/devtools/build/lib/actions",
+        "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
+        "//src/main/java/com/google/devtools/build/skyframe",
+        "//third_party:guava",
+    ],
+)
+
 java_library(
     name = "containing_package_lookup_function",
     srcs = ["ContainingPackageLookupFunction.java"],