Declare dependencies on discovered inputs before execution instead of after.

As a side effect, we no longer restart ActionExecutionFunction after the action has executed unless the action ran locally (and therefore didn't discover its inputs beforehand). This also means that we no longer need to store an out-of-Skyframe cache for discovered includes except when checking the action cache. Since my suspicion is that the out-of-Skyframe cache will have a minimal performance impact once it is just being used for the action cache, I may delete it in a follow-up cl.

After this change, we will overapproximate the set of includes because we depend on all includes, rather than just the ones that the action was found to depend on after execution.

This is a prerequisite for Skyframe-native include scanning in that Skyframe-native include scanning will need to add at least as many Skyframe nodes and edges. If we end up punting on it, then we may want to revert this change. But for now I think it's worth having. I'll run some more numbers to see what the actual performance impact is.

--
MOS_MIGRATED_REVID=88492955

5 files changed, 214 insertions, 47 deletions

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 23d841deba..f0a67102c6 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
@@ -89,13 +89,13 @@ public class CppCompileAction extends AbstractAction implements IncludeScannable
     /**
      * Returns the set of files to be added for an included file (as returned in the .d file)
      */
-    Iterable<Artifact> getInputsForIncludedFile(
+    Collection<Artifact> getInputsForIncludedFile(
         Artifact includedFile, ArtifactResolver artifactResolver);
   }
 
   public static final IncludeResolver VOID_INCLUDE_RESOLVER = new IncludeResolver() {
     @Override
-    public Iterable<Artifact> getInputsForIncludedFile(Artifact includedFile,
+    public Collection<Artifact> getInputsForIncludedFile(Artifact includedFile,
         ArtifactResolver artifactResolver) {
       return ImmutableList.of();
     }
@@ -308,16 +308,38 @@ public class CppCompileAction extends AbstractAction implements IncludeScannable
   public Collection<Artifact> discoverInputs(ActionExecutionContext actionExecutionContext)
       throws ActionExecutionException, InterruptedException {
     Executor executor = actionExecutionContext.getExecutor();
-    Collection<Artifact> returnValue = null;
+    Collection<Artifact> initialResult;
     try {
-      returnValue = executor.getContext(CppCompileActionContext.class)
+      initialResult = executor.getContext(CppCompileActionContext.class)
           .findAdditionalInputs(this, actionExecutionContext);
     } catch (ExecException e) {
       throw e.toActionExecutionException("Include scanning of rule '" + getOwner().getLabel() + "'",
           executor.getVerboseFailures(), this);
     }
-    this.additionalInputs = returnValue == null ? ImmutableList.<Artifact>of() : returnValue;
-    return returnValue;
+    if (initialResult == null) {
+      // We will find inputs during execution. Store an empty list to show we did try to discover
+      // inputs and return null to inform the caller that inputs will be discovered later.
+      this.additionalInputs = ImmutableList.of();
+      return null;
+    }
+    this.additionalInputs = initialResult;
+    // In some cases, execution backends need extra files for each included file. Add them
+    // to the set of inputs the caller may need to be aware of.
+    Collection<Artifact> result = new HashSet<>();
+    ArtifactResolver artifactResolver =
+        executor.getContext(IncludeScanningContext.class).getArtifactResolver();
+    for (Artifact artifact : initialResult) {
+      result.addAll(includeResolver.getInputsForIncludedFile(artifact, artifactResolver));
+    }
+    for (Artifact artifact : getInputs()) {
+      result.addAll(includeResolver.getInputsForIncludedFile(artifact, artifactResolver));
+    }
+    if (result.isEmpty()) {
+      result = initialResult;
+    } else {
+      result.addAll(initialResult);
+    }
+    return result;
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionContext.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionContext.java
index 2a0f6a1b51..87d0a2b031 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionContext.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionContext.java
@@ -41,7 +41,13 @@ public interface CppCompileActionContext extends ActionContext {
     byte[] getContents() throws IOException;
   }
 
-  /** Does include scanning to find the list of files needed to execute the action. */
+  /**
+   * Does include scanning to find the list of files needed to execute the action.
+   *
+   * <p>Returns null if additional inputs will only be found during action execution, not before.
+   * </p>
+   */
+  @Nullable
   public Collection<Artifact> findAdditionalInputs(CppCompileAction action,
       ActionExecutionContext actionExecutionContext)
       throws ExecException, InterruptedException, ActionExecutionException;
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 812d1fa89b..e03102a9cb 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
@@ -14,9 +14,11 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.base.Function;
+import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionCacheChecker.Token;
@@ -49,11 +51,12 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentMap;
 
 /**
  * A builder for {@link ActionExecutionValue}s.
  */
-public class ActionExecutionFunction implements SkyFunction {
+public class ActionExecutionFunction implements SkyFunction, CompletionReceiver {
 
   private static final Predicate<Artifact> IS_SOURCE_ARTIFACT = new Predicate<Artifact>() {
     @Override
@@ -64,11 +67,13 @@ public class ActionExecutionFunction implements SkyFunction {
 
   private final SkyframeActionExecutor skyframeActionExecutor;
   private final TimestampGranularityMonitor tsgm;
+  private ConcurrentMap<Action, ContinuationState> stateMap;
 
   public ActionExecutionFunction(SkyframeActionExecutor skyframeActionExecutor,
       TimestampGranularityMonitor tsgm) {
     this.skyframeActionExecutor = skyframeActionExecutor;
     this.tsgm = tsgm;
+    stateMap = Maps.newConcurrentMap();
   }
 
   @Override
@@ -111,9 +116,8 @@ public class ActionExecutionFunction implements SkyFunction {
       // prints the error in the top-level reporter and also dumps the recorded StdErr for the
       // action. Label can be null in the case of, e.g., the SystemActionOwner (for build-info.txt).
       throw new ActionExecutionFunctionException(new AlreadyReportedActionExecutionException(e));
-    } finally {
-      declareAdditionalDependencies(env, action);
     }
+
     if (env.valuesMissing()) {
       return null;
     }
@@ -169,20 +173,23 @@ public class ActionExecutionFunction implements SkyFunction {
       Map<Artifact, FileArtifactValue> inputArtifactData,
       Map<Artifact, Collection<Artifact>> expandedMiddlemen,
       Environment env) throws ActionExecutionException, InterruptedException {
-    // Don't initialize the cache if the result has already been computed and this is just a
-    // rerun.
-    FileAndMetadataCache fileAndMetadataCache = null;
-    MetadataHandler metadataHandler = null;
-    Token token = null;
-    long actionStartTime = System.nanoTime();
-    // inputArtifactData is null exactly when we know that the execution result was already
-    // computed on a prior run of this SkyFunction. If it is null we don't need to initialize
-    // anything -- we will get the result directly from SkyframeActionExecutor's cache.
-    if (inputArtifactData != null) {
-      // Check action cache to see if we need to execute anything. Checking the action cache only
-      // needs to happen on the first run, since a cache hit means we'll return immediately, and
-      // there'll be no second run.
-      fileAndMetadataCache = new FileAndMetadataCache(
+    // If this is the second time we are here (because the action discovers inputs, and we had
+    // to restart the value builder after declaring our dependence on newly discovered inputs),
+    // the result returned here is the already-computed result from the first run.
+    // Similarly, if this is a shared action and the other action is the one that executed, we
+    // must use that other action's value, provided here, since it is populated with metadata
+    // for the outputs.
+    if (inputArtifactData == null) {
+      return skyframeActionExecutor.executeAction(action, null, null, -1, null);
+    }
+    ContinuationState state;
+    if (action.discoversInputs()) {
+      state = getState(action);
+    } else {
+      state = new ContinuationState();
+    }
+    // This may be recreated if we discover inputs.
+    FileAndMetadataCache fileAndMetadataCache = new FileAndMetadataCache(
           inputArtifactData,
           expandedMiddlemen,
           skyframeActionExecutor.getExecRoot(),
@@ -192,49 +199,112 @@ public class ActionExecutionFunction implements SkyFunction {
           // present in Skyframe, so we can save a stat by looking them up directly.
           action.discoversInputs() ? env : null,
           tsgm);
-      metadataHandler =
+    MetadataHandler metadataHandler =
           skyframeActionExecutor.constructMetadataHandler(fileAndMetadataCache);
-      token = skyframeActionExecutor.checkActionCache(action, metadataHandler,  
+    long actionStartTime = System.nanoTime();
+    // We only need to check the action cache if we haven't done it on a previous run.
+    if (!state.hasDiscoveredInputs()) {
+      Token token = skyframeActionExecutor.checkActionCache(action, metadataHandler,
           new PackageRootResolverWithEnvironment(env), actionStartTime);
       if (token == Token.NEED_TO_RERUN) {
+        // Sadly, there is no state that we can preserve here across this restart.
         return null;
       }
+      state.token = token;
     }
-    if (token == null && inputArtifactData != null) {
+
+    if (state.token == null) {
       // We got a hit from the action cache -- no need to execute.
       return new ActionExecutionValue(
           fileAndMetadataCache.getOutputData(),
           fileAndMetadataCache.getAdditionalOutputData());
     }
 
-    ActionExecutionContext actionExecutionContext = null;
+    // This may be recreated if we discover inputs.
+    ActionExecutionContext actionExecutionContext =
+        skyframeActionExecutor.constructActionExecutionContext(fileAndMetadataCache,
+            metadataHandler);
+    boolean inputsDiscoveredDuringActionExecution = false;
     try {
-      if (inputArtifactData != null) {
-        actionExecutionContext = skyframeActionExecutor.constructActionExecutionContext(
-            fileAndMetadataCache, metadataHandler);
-        if (action.discoversInputs()) {
-          skyframeActionExecutor.discoverInputs(action, actionExecutionContext);
+      if (action.discoversInputs()) {
+        if (!state.hasDiscoveredInputs()) {
+          state.discoveredInputs =
+              skyframeActionExecutor.discoverInputs(action, actionExecutionContext);
+          if (state.discoveredInputs == null) {
+            // Action had nothing to tell us about discovered inputs before execution. We'll have to
+            // add them afterwards.
+            inputsDiscoveredDuringActionExecution = true;
+          }
         }
+        if (state.discoveredInputs != null
+            && !inputArtifactData.keySet().containsAll(state.discoveredInputs)) {
+          inputArtifactData = addDiscoveredInputs(inputArtifactData, state.discoveredInputs,
+              env);
+          if (env.valuesMissing()) {
+            // This is the only place that we actually preserve meaningful state across restarts.
+            return null;
+          }
+          fileAndMetadataCache = new FileAndMetadataCache(
+              inputArtifactData,
+              expandedMiddlemen,
+              skyframeActionExecutor.getExecRoot(),
+              action.getOutputs(),
+              null,
+              tsgm
+          );
+          actionExecutionContext = skyframeActionExecutor.constructActionExecutionContext(
+              fileAndMetadataCache, fileAndMetadataCache);
+        }
+      }
+      // Clear state before actual execution of action. It will never be needed again because
+      // skyframeActionExecutor is guaranteed to have a result after this.
+      Token token = state.token;
+      if (action.discoversInputs()) {
+        removeState(action);
       }
-      // If this is the second time we are here (because the action discovers inputs, and we had
-      // to restart the value builder after declaring our dependence on newly discovered inputs),
-      // the result returned here is the already-computed result from the first run.
-      // Similarly, if this is a shared action and the other action is the one that executed, we
-      // must use that other action's value, provided here, since it is populated with metadata
-      // for the outputs.
-      // If this action was not shared and this is the first run of the action, this returned
-      // result was computed during the call.
+      state = null;
       return skyframeActionExecutor.executeAction(action, fileAndMetadataCache, token,
           actionStartTime, actionExecutionContext);
     } finally {
       try {
-        if (actionExecutionContext != null) {
-          actionExecutionContext.getFileOutErr().close();
-        }
+        actionExecutionContext.getFileOutErr().close();
       } catch (IOException e) {
         // Nothing we can do here.
       }
+      if (inputsDiscoveredDuringActionExecution) {
+        declareAdditionalDependencies(env, action);
+      }
+    }
+  }
+
+  private static Map<Artifact, FileArtifactValue> addDiscoveredInputs(
+      Map<Artifact, FileArtifactValue> originalInputData, Collection<Artifact> discoveredInputs,
+      Environment env) {
+    Map<Artifact, FileArtifactValue> result = new HashMap<>(originalInputData);
+    Set<SkyKey> keys = new HashSet<>();
+    for (Artifact artifact : discoveredInputs) {
+      if (!result.containsKey(artifact)) {
+        // Note that if the artifact is derived, the mandatory flag is ignored.
+        keys.add(ArtifactValue.key(artifact, /*mandatory=*/false));
+      }
+    }
+    // We do not do a getValuesOrThrow() call for the following reasons:
+    // 1. No exceptions can be thrown for non-mandatory inputs;
+    // 2. Any derived inputs must be in the transitive closure of this action's inputs. Therefore,
+    // if there was an error building one of them, then that exception would have percolated up to
+    // this action already, through one of its declared inputs, and we would not have reached input
+    // discovery.
+    // Therefore there is no need to catch and rethrow exceptions as there is with #checkInputs.
+    Map<SkyKey, SkyValue> data = env.getValues(keys);
+    if (env.valuesMissing()) {
+      return null;
     }
+    for (Map.Entry<SkyKey, SkyValue> depsEntry : data.entrySet()) {
+      Artifact input = ArtifactValue.artifact(depsEntry.getKey());
+      result.put(input,
+          Preconditions.checkNotNull((FileArtifactValue) depsEntry.getValue(), input));
+    }
+    return result;
   }
 
   private static Iterable<SkyKey> toKeys(Iterable<Artifact> inputs,
@@ -291,7 +361,7 @@ public class ActionExecutionFunction implements SkyFunction {
     // Only populate input data if we have the input values, otherwise they'll just go unused.
     // We still want to loop through the inputs to collect missing deps errors. During the
     // evaluator "error bubbling", we may get one last chance at reporting errors even though
-    // some deps are stilling missing.
+    // some deps are still missing.
     boolean populateInputData = !env.valuesMissing();
     NestedSetBuilder<Label> rootCauses = NestedSetBuilder.stableOrder();
     Map<Artifact, FileArtifactValue> inputArtifactData =
@@ -375,6 +445,52 @@ public class ActionExecutionFunction implements SkyFunction {
   }
 
   /**
+   * Should be called once execution is over, and the intra-build cache of in-progress computations
+   * should be discarded. If the cache is non-empty (due to an interrupted/failed build), failure to
+   * call complete() can both cause a memory leak and incorrect results on the subsequent build.
+   */
+  @Override
+  public void complete() {
+    // Discard all remaining state (there should be none after a successful execution).
+    stateMap = Maps.newConcurrentMap();
+  }
+
+  private ContinuationState getState(Action action) {
+    ContinuationState state = stateMap.get(action);
+    if (state == null) {
+      state = new ContinuationState();
+      Preconditions.checkState(stateMap.put(action, state) == null, action);
+    }
+    return state;
+  }
+
+  private void removeState(Action action) {
+    Preconditions.checkNotNull(stateMap.remove(action), action);
+  }
+
+  /**
+   * State to save work across restarts of ActionExecutionFunction due to missing discovered inputs.
+   */
+  private static class ContinuationState {
+    Token token = null;
+    Collection<Artifact> discoveredInputs = null;
+
+    // This will always be false for actions that don't discover their inputs, but we never restart
+    // those actions in any case. For actions that do discover their inputs, they either discover
+    // them before execution, in which case discoveredInputs will be non-null if that has already
+    // happened, or after execution, in which case they returned null when Action#discoverInputs()
+    // was called, and won't restart due to missing dependencies before execution.
+    boolean hasDiscoveredInputs() {
+      return discoveredInputs != null;
+    }
+
+    @Override
+    public String toString() {
+      return token + ", " + discoveredInputs;
+    }
+  }
+
+  /**
    * Used to declare all the exception types that can be wrapped in the exception thrown by
    * {@link ActionExecutionFunction#compute}.
    */
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionReceiver.java b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionReceiver.java
new file mode 100644
index 0000000000..ec4b6f7432
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionReceiver.java
@@ -0,0 +1,19 @@
+// Copyright 2015 Google Inc. 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;
+
+/** Interface for objects that wish to be notified when a skyframe evaluation has finished. */
+public interface CompletionReceiver {
+  void complete();
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 5bf35799b5..2e76c14927 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -206,6 +206,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
   private final AtomicReference<ActionExecutionStatusReporter> statusReporterRef =
       new AtomicReference<>();
   private final SkyframeActionExecutor skyframeActionExecutor;
+  private CompletionReceiver actionExecutionFunction;
   protected SkyframeProgressReceiver progressReceiver;
   private final AtomicReference<CyclesReporter> cyclesReporter = new AtomicReference<>();
 
@@ -320,8 +321,10 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
         buildDataDirectory));
     map.put(SkyFunctions.BUILD_INFO, new WorkspaceStatusFunction());
     map.put(SkyFunctions.COVERAGE_REPORT, new CoverageReportFunction());
-    map.put(SkyFunctions.ACTION_EXECUTION,
-        new ActionExecutionFunction(skyframeActionExecutor, tsgm));
+    ActionExecutionFunction actionExecutionFunction =
+        new ActionExecutionFunction(skyframeActionExecutor, tsgm);
+    map.put(SkyFunctions.ACTION_EXECUTION, actionExecutionFunction);
+    this.actionExecutionFunction = actionExecutionFunction;
     map.put(SkyFunctions.RECURSIVE_FILESYSTEM_TRAVERSAL,
         new RecursiveFilesystemTraversalFunction());
     map.put(SkyFunctions.FILESET_ENTRY, new FilesetEntryFunction());
@@ -955,6 +958,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
       // Also releases thread locks.
       resourceManager.resetResourceUsage();
       skyframeActionExecutor.executionOver();
+      actionExecutionFunction.complete();
     }
   }