Remove the need for discoverInputsStage2(). Much like every other Skyframe

function, discoverInputs() can be implemented to return null upon encountering
missing ActionExecutionValues.

RELNOTES: None.
PiperOrigin-RevId: 206913969

4 files changed, 30 insertions, 74 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java b/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java
index 44db64bf46..4156175bb4 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java
@@ -42,7 +42,6 @@ import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.Symlinks;
-import com.google.devtools.build.skyframe.SkyFunction;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Map;
@@ -194,12 +193,6 @@ public abstract class AbstractAction implements Action, ActionApi {
   }
 
   @Override
-  public Iterable<Artifact> discoverInputsStage2(SkyFunction.Environment env)
-      throws ActionExecutionException, InterruptedException {
-    return null;
-  }
-
-  @Override
   public Iterable<Artifact> getAllowedDerivedInputs() {
     throw new IllegalStateException(
         "Method must be overridden for actions that may have unknown inputs.");
diff --git a/src/main/java/com/google/devtools/build/lib/actions/Action.java b/src/main/java/com/google/devtools/build/lib/actions/Action.java
index 4c27425f91..e8a75be78a 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/Action.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/Action.java
@@ -18,7 +18,6 @@ import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ConditionallyThreadCompatible;
 import com.google.devtools.build.lib.vfs.FileSystem;
 import com.google.devtools.build.lib.vfs.Path;
-import com.google.devtools.build.skyframe.SkyFunction;
 import java.io.IOException;
 import javax.annotation.Nullable;
 
@@ -139,27 +138,16 @@ public interface Action extends ActionExecutionMetadata {
   boolean isVolatile();
 
   /**
-   * Method used to find inputs before execution for an action that
-   * {@link ActionExecutionMetadata#discoversInputs}. Returns null if action's inputs will be
-   * discovered during execution proper.
+   * Method used to find inputs before execution for an action that {@link
+   * ActionExecutionMetadata#discoversInputs}. Returns the set of discovered inputs (may be the
+   * empty set) or null if this action declared additional Skyframe dependencies that must be
+   * computed before it can make a decision.
    */
   @Nullable
   Iterable<Artifact> discoverInputs(ActionExecutionContext actionExecutionContext)
       throws ActionExecutionException, InterruptedException;
 
   /**
-   * Used in combination with {@link #discoverInputs} if inputs need to be found before execution in
-   * multiple steps. Returns null if two-stage input discovery isn't necessary.
-   *
-   * <p>Any deps requested here must not change unless one of the action's inputs changes.
-   * Otherwise, changes to nodes that should cause re-execution of actions might be prevented by the
-   * action cache.
-   */
-  @Nullable
-  Iterable<Artifact> discoverInputsStage2(SkyFunction.Environment env)
-      throws ActionExecutionException, InterruptedException;
-
-  /**
    * Returns the set of artifacts that can possibly be inputs. It will be called iff inputsKnown()
    * is false for the given action instance and there is a related cache entry in the action cache.
    *
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
index 5b7c66faab..a73edf0697 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
@@ -140,9 +140,8 @@ public class CppCompileAction extends AbstractAction
   private Iterable<Artifact> additionalInputs = null;
 
   /**
-   * Set when a two-stage input discovery is used.
-   *
-   * <p>Used only during action execution.
+   * Used only during input discovery, when input discovery requires other actions
+   * to be executed first.
    */
   private Set<Artifact> usedModules = null;
 
@@ -160,13 +159,11 @@ public class CppCompileAction extends AbstractAction
    *   <li><i>Action caching.</i> It is set when restoring from the action cache. It is queried
    *       immediately after restoration to populate the {@link
    *       com.google.devtools.build.lib.skyframe.ActionExecutionValue}.
-   *   <li><i>Input discovery</i>It is set by {@link #discoverInputsStage2}. It is queried to
+   *   <li><i>Input discovery</i>It is set by {@link #discoverInputs}. It is queried to
    *       populate the {@link com.google.devtools.build.lib.skyframe.ActionExecutionValue}.
    *   <li><i>Compilation</i>Compilation reads this field to know what needs to be staged.
    * </ul>
    */
-  // TODO(djasper): investigate releasing memory used by this field as early as possible, for
-  // example, by including these values in additionalInputs.
   private ImmutableList<Artifact> discoveredModules = null;
 
   /**
@@ -439,45 +436,38 @@ public class CppCompileAction extends AbstractAction
       throws ActionExecutionException, InterruptedException {
     Preconditions.checkArgument(!sourceFile.isFileType(CppFileTypes.CPP_MODULE));
 
-    additionalInputs = findUsedHeaders(actionExecutionContext);
-    if (!shouldScanIncludes) {
-      return additionalInputs;
-    }
+    if (additionalInputs == null) {
+      additionalInputs = findUsedHeaders(actionExecutionContext);
 
-    if (!shouldScanDotdFiles()) {
-      additionalInputs = filterDiscoveredHeaders(actionExecutionContext, additionalInputs);
+      if (!shouldScanIncludes) {
+        return additionalInputs;
+      }
+
+      if (!shouldScanDotdFiles()) {
+        // If we aren't looking at .d files later, remove undeclared inputs now.
+        additionalInputs = filterDiscoveredHeaders(actionExecutionContext, additionalInputs);
+      }
     }
 
-    if (!shouldPruneModules) {
+    if (!shouldScanIncludes || !shouldPruneModules) {
       return additionalInputs;
     }
 
-    usedModules =
-        ccCompilationContext.getUsedModules(usePic, ImmutableSet.copyOf(additionalInputs));
-    return Iterables.concat(additionalInputs, usedModules);
-  }
-
-  /** @return null when either {@link #usedModules} was null or on Skyframe lookup failure */
-  @Nullable
-  @Override
-  public Iterable<Artifact> discoverInputsStage2(SkyFunction.Environment env)
-      throws InterruptedException {
     if (usedModules == null) {
-      // No modules were used in this compilation, no need to do any work.
-      return null;
+      usedModules =
+          ccCompilationContext.getUsedModules(usePic, ImmutableSet.copyOf(additionalInputs));
     }
-
-    Set<Artifact> transitivelyUsedModules = computeTransitivelyUsedModules(env, usedModules);
+    Set<Artifact> transitivelyUsedModules =
+        computeTransitivelyUsedModules(
+            actionExecutionContext.getEnvironmentForDiscoveringInputs(), usedModules);
     if (transitivelyUsedModules == null) {
-      // Not all used modules available yet. ActionExecutionValues have been requested. Return so
-      // that this function can be re-executed when ready.
       return null;
     }
 
     discoveredModules = ImmutableList.copyOf(Sets.union(usedModules, transitivelyUsedModules));
     topLevelModules = ImmutableList.copyOf(Sets.difference(usedModules, transitivelyUsedModules));
     usedModules = null;
-    return transitivelyUsedModules;
+    return Iterables.concat(additionalInputs, discoveredModules);
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
index aaa9221098..4e61f8c6bd 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
@@ -435,8 +435,12 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
           state.discoveredInputs =
               skyframeActionExecutor.discoverInputs(
                   action, metadataHandler, env, state.actionFileSystem);
-          Preconditions.checkState(state.discoveredInputs != null,
-              "discoverInputs() returned null on action %s", action);
+          Preconditions.checkState(
+              env.valuesMissing() == (state.discoveredInputs == null),
+              "discoverInputs() must return null iff requesting more dependencies.");
+          if (state.discoveredInputs == null) {
+            return null;
+          }
         } catch (MissingDepException e) {
           Preconditions.checkState(env.valuesMissing(), action);
           return null;
@@ -448,24 +452,6 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
         return null;
       }
 
-      // Stage 1 finished, let's do stage 2. The stage 1 of input discovery will have added some
-      // files with addDiscoveredInputs() and then have waited for those files to be available
-      // by returning null if env.valuesMissing() returned true. So stage 2 can now access those
-      // inputs to discover even more inputs and then potentially also wait for those to be
-      // available.
-      if (state.discoveredInputsStage2 == null) {
-        state.discoveredInputsStage2 = action.discoverInputsStage2(env);
-        if (env.valuesMissing()) {
-          return null;
-        }
-      }
-      if (state.discoveredInputsStage2 != null) {
-        addDiscoveredInputs(
-            state.inputArtifactData, state.expandedArtifacts, state.discoveredInputsStage2, env);
-        if (env.valuesMissing()) {
-          return null;
-        }
-      }
       metadataHandler =
           new ActionMetadataHandler(
               state.inputArtifactData,
@@ -821,7 +807,6 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
     Map<Artifact, Collection<Artifact>> expandedArtifacts = null;
     Token token = null;
     Iterable<Artifact> discoveredInputs = null;
-    Iterable<Artifact> discoveredInputsStage2 = null;
     ActionExecutionValue value = null;
     FileSystem actionFileSystem = null;