Prune modules when building modules themselves to reduce build times shorten

critical paths.

When the inputs of a module M are reduced to a set S, then that same set S also
needs to be supplied to compile something that uses M. To do this, input
discovery is divided into two stages. For each CppCompileAction, the first
stage discovers the necessary modules (M above). These are then added as inputs
to ensure that they are built. The second stage then finds all the modules (S
above) that are required to use those and also adds them as inputs.

For now, the new behavior is guarded by a new flag
--experimental_prune_more_modules.

This is currently implemented by reading the .d files of used modules add
adding all their module dependencies. There are two noteworthy alternatives:
1. Hack up input discovery to understand modules, e.g. if a modular header is
   hit, continue scanning from all it's headers. However, this seems very
   brittle and a lot of additional information would have to be passed to the
   input discovery.
2. Directly pass the results from input discovery of one CppCompileAction to
   another one. However, this seems to tightly couple the execution of
   different CppCompileActions and might lead to a mess of different states,
   more memory consumption, etc.
With the current implementation, there is a bit of runtime overhead of
reading the .d files (many times). This could potentially be improved by
caching the results. However, even without this caching, the runtime overhead
is limited (<10%) for all builds I have tried (I have tried with builds where
all the compile results are already in the executor's cache.

--
MOS_MIGRATED_REVID=140793217

8 files changed, 130 insertions, 43 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 dc1857826b..3dc3e4be40 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
@@ -40,6 +40,7 @@ import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 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;
@@ -179,6 +180,12 @@ public abstract class AbstractAction implements Action, SkylarkValue {
         + " since it does not discover inputs");
   }
 
+  @Override
+  public Iterable<Artifact> discoverInputsStage2(SkyFunction.Environment env)
+      throws ActionExecutionException, InterruptedException {
+    return null;
+  }
+
   @Nullable
   @Override
   public Iterable<Artifact> getInputsWhenSkippingInputDiscovery() {
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 ecf8b66182..12316b7264 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
@@ -19,6 +19,7 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.ConditionallyThread
 import com.google.devtools.build.lib.profiler.Describable;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.skyframe.SkyFunction;
 import java.io.IOException;
 import java.util.Collection;
 import javax.annotation.Nullable;
@@ -152,6 +153,18 @@ public interface Action extends ActionExecutionMetadata, Describable {
       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;
+
+  /**
    * If an action does not know the exact set of inputs ahead of time, it will usually either do:
    * <ul>
    * <li> Execution time pruning: The action provides a superset set of inputs at action creation
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompilationContext.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompilationContext.java
index ae66887989..9ecde15051 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompilationContext.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompilationContext.java
@@ -30,8 +30,10 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -210,7 +212,8 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
     return usePic ? picModuleInfo.transitiveModules : moduleInfo.transitiveModules;
   }
 
-  public Collection<Artifact> getUsedModules(boolean usePic, Set<Artifact> usedHeaders) {
+  public Collection<TransitiveModuleHeaders> getUsedModules(
+      boolean usePic, Set<Artifact> usedHeaders) {
     return usePic
         ? picModuleInfo.getUsedModules(usedHeaders)
         : moduleInfo.getUsedModules(usedHeaders);
@@ -654,7 +657,7 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
       this.picModuleInfo.setHeaderModule(picHeaderModule);
       return this;
     }
-    
+
     /**
      * Sets that the context will be used by a compilation that needs transitive module maps.
      */
@@ -662,7 +665,7 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
       this.provideTransitiveModuleMaps = provideTransitiveModuleMaps;
       return this;
     }
-    
+
     /**
      * Builds the {@link CppCompilationContext}.
      */
@@ -767,13 +770,13 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
      * context.
      */
     private final Artifact headerModule;
-    
+
     /** All header files that are compiled into this module. */
     private final ImmutableSet<Artifact> modularHeaders;
 
     /** All header files that are contained in this module. */
     private final ImmutableSet<Artifact> textualHeaders;
-    
+
     /**
      * All transitive modules that this context depends on, excluding headerModule.
      */
@@ -797,12 +800,13 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
       this.transitiveModuleHeaders = transitiveModuleHeaders;
     }
 
-    public Collection<Artifact> getUsedModules(Set<Artifact> usedHeaders) {
-      Set<Artifact> result = new LinkedHashSet<>();
+    public Collection<TransitiveModuleHeaders> getUsedModules(Set<Artifact> usedHeaders) {
+      List<TransitiveModuleHeaders> result = new ArrayList<>();
       for (TransitiveModuleHeaders transitiveModule : transitiveModuleHeaders) {
-        if (result.contains(transitiveModule.module)) {
-          // If result already contains this module, we will have added it and all its
-          // transitive dependencies to result already. No need to check whether to add it again.
+        if (transitiveModule.module.equals(headerModule)) {
+          // Do not add the module of the current rule for both:
+          // 1. the module compile itself
+          // 2. compiles of other translation units of the same rule.
           continue;
         }
         boolean providesUsedHeader = false;
@@ -813,15 +817,7 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
           }
         }
         if (providesUsedHeader) {
-          result.addAll(transitiveModule.transitiveModules.toCollection());
-          if (!transitiveModule.module.equals(headerModule)) {
-            // Only add the module itself if it is not the main headerModule. This is done because
-            // we don't want to pass the header module when compiling the translation unit that
-            // defines it. Due to how modules are implemented, #includes from a translation unit to
-            // the headers that it implements remain textual. Thus, adding this header module as a
-            // dependency would only prolong the build without any possible benefit.
-            result.add(transitiveModule.module);
-          }
+          result.add(transitiveModule);
         }
       }
       return result;
@@ -834,8 +830,8 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
       private Artifact headerModule = null;
       private final Set<Artifact> modularHeaders = new LinkedHashSet<>();
       private final Set<Artifact> textualHeaders = new LinkedHashSet<>();
-      private NestedSetBuilder<Artifact> transitiveModules = NestedSetBuilder.stableOrder();
-      private NestedSetBuilder<TransitiveModuleHeaders> transitiveModuleHeaders =
+      private final NestedSetBuilder<Artifact> transitiveModules = NestedSetBuilder.stableOrder();
+      private final NestedSetBuilder<TransitiveModuleHeaders> transitiveModuleHeaders =
           NestedSetBuilder.stableOrder();
 
       public Builder setHeaderModule(Artifact headerModule) {
@@ -927,5 +923,13 @@ public final class CppCompilationContext implements TransitiveInfoProvider {
       this.headers = headers;
       this.transitiveModules = transitiveModules;
     }
+
+    public Artifact getModule() {
+      return module;
+    }
+
+    public Collection<Artifact> getTransitiveModules() {
+      return transitiveModules.toCollection();
+    }
   }
 }
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 e8683fca58..0ef064b77a 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
@@ -53,6 +53,8 @@ import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.rules.cpp.CcToolchainFeatures.FeatureConfiguration;
 import com.google.devtools.build.lib.rules.cpp.CppCompileActionContext.Reply;
 import com.google.devtools.build.lib.rules.cpp.CppConfiguration.Tool;
+import com.google.devtools.build.lib.skyframe.ActionLookupValue;
+import com.google.devtools.build.lib.skyframe.ActionLookupValue.ActionLookupKey;
 import com.google.devtools.build.lib.util.DependencySet;
 import com.google.devtools.build.lib.util.FileType;
 import com.google.devtools.build.lib.util.Fingerprint;
@@ -63,6 +65,9 @@ import com.google.devtools.build.lib.util.ShellEscaper;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -207,6 +212,9 @@ public class CppCompileAction extends AbstractAction
    */
   private Collection<Artifact> additionalInputs = null;
 
+  /** Set when a two-stage input discovery is used. */
+  private Collection<Artifact> usedModules = null;
+
   private CcToolchainFeatures.Variables overwrittenVariables = null;
 
   private ImmutableList<Artifact> resolvedInputs = ImmutableList.<Artifact>of();
@@ -504,8 +512,18 @@ public class CppCompileAction extends AbstractAction
 
     if (shouldPruneModules) {
       Set<Artifact> initialResultSet = Sets.newLinkedHashSet(initialResult);
-      Collection<Artifact> usedModules = context.getUsedModules(usePic, initialResultSet);
+      Set<Artifact> usedModules = Sets.newLinkedHashSet();
+      for (CppCompilationContext.TransitiveModuleHeaders usedModule :
+          context.getUsedModules(usePic, initialResultSet)) {
+        usedModules.add(usedModule.getModule());
+        if (!cppConfiguration.getPruneMoreModules()) {
+          usedModules.addAll(usedModule.getTransitiveModules());
+        }
+      }
       initialResultSet.addAll(usedModules);
+      if (cppConfiguration.getPruneMoreModules()) {
+        this.usedModules = usedModules;
+      }
       initialResult = initialResultSet;
     }
 
@@ -533,6 +551,39 @@ public class CppCompileAction extends AbstractAction
   }
 
   @Override
+  public Iterable<Artifact> discoverInputsStage2(SkyFunction.Environment env)
+      throws ActionExecutionException, InterruptedException {
+    if (this.usedModules == null) {
+      return null;
+    }
+    Map<Artifact, SkyKey> skyKeys = new HashMap<>();
+    for (Artifact artifact : this.usedModules) {
+      skyKeys.put(artifact, ActionLookupValue.key((ActionLookupKey) artifact.getArtifactOwner()));
+    }
+    Map<SkyKey, SkyValue> skyValues = env.getValues(skyKeys.values());
+    Set<Artifact> additionalModules = Sets.newLinkedHashSet();
+    for (Artifact artifact : this.usedModules) {
+      SkyKey skyKey = skyKeys.get(artifact);
+      ActionLookupValue value = (ActionLookupValue) skyValues.get(skyKey);
+      Preconditions.checkNotNull(
+          value, "Owner %s of %s not in graph %s", artifact.getArtifactOwner(), artifact, skyKey);
+      CppCompileAction action = (CppCompileAction) value.getGeneratingAction(artifact);
+      for (Artifact input : action.getInputs()) {
+        if (CppFileTypes.CPP_MODULE.matches(input.getFilename())) {
+          additionalModules.add(input);
+        }
+      }
+    }
+    this.additionalInputs =
+        new ImmutableList.Builder<Artifact>()
+            .addAll(this.additionalInputs)
+            .addAll(additionalModules)
+            .build();
+    this.usedModules = null;
+    return additionalModules;
+  }
+
+  @Override
   public Iterable<Artifact> getInputsWhenSkippingInputDiscovery() {
     if (useHeaderModules && cppConfiguration.getSkipUnusedModules()) {
       this.additionalInputs = context.getTransitiveModules(usePic).toCollection();
@@ -1169,7 +1220,6 @@ public class CppCompileAction extends AbstractAction
   public void execute(
       ActionExecutionContext actionExecutionContext)
           throws ActionExecutionException, InterruptedException {
- 
     if (useHeaderModules) {
       // If modules pruning is used, modules will be supplied via additionalInputs, otherwise they
       // are regular inputs.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionBuilder.java
index d1ee7f3967..9a3780d07c 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionBuilder.java
@@ -263,17 +263,7 @@ public class CppCompileActionBuilder {
     if (!fake && !shouldScanIncludes) {
       realMandatoryInputsBuilder.addTransitive(context.getDeclaredIncludeSrcs());
     }
-    // We disable pruning header modules in CPP_MODULE_COMPILEs as that would lead to
-    // module-out-of-date errors. The problem surfaces if a module A depends on a module B, but the
-    // headers of module A don't actually use any of B's headers.  Then we would not need to rebuild
-    // A when B changes although we are storing a dependency in it. If B changes and we then build
-    // something that uses A (a header of it), we mark A and all of its transitive deps as inputs.
-    // We still don't need to rebuild A, as none of its inputs have changed, but we do rebuild B
-    // now and then the two modules are out of sync.
-    boolean shouldPruneModules =
-        shouldScanIncludes
-            && useHeaderModules
-            && !getActionName().equals(CppCompileAction.CPP_MODULE_COMPILE);
+    boolean shouldPruneModules = shouldScanIncludes && useHeaderModules;
     if (useHeaderModules && !shouldPruneModules) {
       realMandatoryInputsBuilder.addTransitive(context.getTransitiveModules(usePic));
     }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java
index 91ef450c2c..f0b6e4d61b 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppConfiguration.java
@@ -1777,6 +1777,10 @@ public class CppConfiguration extends BuildConfiguration.Fragment {
     return cppOptions.skipUnusedModules;
   }
 
+  public boolean getPruneMoreModules() {
+    return cppOptions.pruneMoreModules;
+  }
+
   public LibcTop getLibcTop() {
     return cppOptions.libcTop;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppOptions.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppOptions.java
index 0a1f16b089..0af42a5aee 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppOptions.java
@@ -546,6 +546,14 @@ public class CppOptions extends FragmentOptions {
   public boolean skipUnusedModules;
 
   @Option(
+    name = "experimental_prune_more_modules",
+    defaultValue = "false",
+    category = "experimental",
+    help = "If enabled, modules pruning is used when building modules themselves."
+  )
+  public boolean pruneMoreModules;
+
+  @Option(
     name = "experimental_omitfp",
     defaultValue = "false",
     category = "semantics",
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 7e07b3c4da..aea3226b1b 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
@@ -384,7 +384,7 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
     ActionExecutionContext actionExecutionContext = null;
     try {
       if (action.discoversInputs()) {
-        if (!state.hasDiscoveredInputs()) {
+        if (state.discoveredInputs == null) {
           try {
             state.discoveredInputs = skyframeActionExecutor.discoverInputs(action,
                 perActionFileCache, metadataHandler, env);
@@ -403,6 +403,24 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
           perActionFileCache = new PerActionFileCache(state.inputArtifactData);
           metadataHandler =
               new ActionMetadataHandler(state.inputArtifactData, action.getOutputs(), tsgm.get());
+
+          // 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 (state.discoveredInputsStage2 != null) {
+            addDiscoveredInputs(state.inputArtifactData, state.discoveredInputsStage2, env);
+            if (env.valuesMissing()) {
+              return null;
+            }
+            perActionFileCache = new PerActionFileCache(state.inputArtifactData);
+            metadataHandler =
+                new ActionMetadataHandler(state.inputArtifactData, action.getOutputs(), tsgm.get());
+          }
         } else {
           // The action generally tries to discover its inputs during execution. If there are any
           // additional inputs necessary to execute the action, make sure they are available now.
@@ -736,6 +754,7 @@ 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;
 
     boolean hasCollectedInputs() {
@@ -748,14 +767,6 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
       return result;
     }
 
-    // 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 we set discoveredInputs then.
-    boolean hasDiscoveredInputs() {
-      return discoveredInputs != null;
-    }
-
     boolean hasCheckedActionCache() {
       // If token is null because there was an action cache hit, this method is never called again
       // because we return immediately.