Partition ObjC header scanning actions to better handle large targets.

PiperOrigin-RevId: 153140874

3 files changed, 73 insertions, 37 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java b/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java
index e999d9ee09..5631ca82c9 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/CompilationSupport.java
@@ -1330,6 +1330,13 @@ public abstract class CompilationSupport {
   /**
    * Creates and registers ObjcHeaderScanning {@link SpawnAction}. Groups all the actions by their
    * compilation command line arguments and creates a ObjcHeaderScanning action for each unique one.
+   *
+   * <p>The number of sources to scan per actions are bounded so that targets with a high number of
+   * sources are not penalized. A large number of sources may require a lot of processing
+   * particularly when the headers required for different sources vary greatly and the caching
+   * mechanism in the tool is largely useless. In these instances these actions would benefit by
+   * being distributed so they don't contribute to the critical path. The partition size is
+   * configurable so that it can be tuned.
    */
   protected void registerHeaderScanningActions(
       ImmutableList<ObjcHeaderThinningInfo> headerThinningInfo,
@@ -1339,51 +1346,65 @@ public abstract class CompilationSupport {
       return;
     }
 
-    FilesToRunProvider headerScannerTool = getHeaderThinningToolExecutable();
-    PrerequisiteArtifacts appleSdks =
-        ruleContext.getPrerequisiteArtifacts(ObjcRuleClasses.APPLE_SDK_ATTRIBUTE, Mode.TARGET);
     ListMultimap<ImmutableList<String>, ObjcHeaderThinningInfo>
         objcHeaderThinningInfoByCommandLine = groupActionsByCommandLine(headerThinningInfo);
     // Register a header scanning spawn action for each unique set of command line arguments
     for (ImmutableList<String> args : objcHeaderThinningInfoByCommandLine.keySet()) {
-      SpawnAction.Builder builder =
-          new SpawnAction.Builder()
-              .setMnemonic("ObjcHeaderScanning")
-              .setExecutable(headerScannerTool)
-              .addInputs(appleSdks.list());
-      CustomCommandLine.Builder cmdLine =
-          CustomCommandLine.builder()
-              .add("--arch")
-              .add(appleConfiguration.getSingleArchitecture().toLowerCase())
-              .add("--platform")
-              .add(appleConfiguration.getSingleArchPlatform().getLowerCaseNameInPlist())
-              .add("--sdk_version")
-              .add(
-                  appleConfiguration
-                      .getSdkVersionForPlatform(appleConfiguration.getSingleArchPlatform())
-                      .toStringWithMinimumComponents(2))
-              .add("--xcode_version")
-              .add(appleConfiguration.getXcodeVersion().toStringWithMinimumComponents(2))
-              .add("--");
-      for (ObjcHeaderThinningInfo info : objcHeaderThinningInfoByCommandLine.get(args)) {
-        cmdLine.addJoinPaths(
-            ":",
-            Lists.newArrayList(info.sourceFile.getExecPath(), info.headersListFile.getExecPath()));
-        builder.addInput(info.sourceFile).addOutput(info.headersListFile);
+      // As infos is in insertion order we should reliably get the same sublists below
+      for (List<ObjcHeaderThinningInfo> partition :
+          Lists.partition(
+              objcHeaderThinningInfoByCommandLine.get(args),
+              objcConfiguration.objcHeaderThinningPartitionSize())) {
+        registerHeaderScanningAction(objcProvider, compilationArtifacts, args, partition);
       }
-      ruleContext.registerAction(
-          builder
-              .setCommandLine(cmdLine.add("--").add(args).build())
-              .addInputs(compilationArtifacts.getPrivateHdrs())
-              .addTransitiveInputs(attributes.hdrs())
-              .addTransitiveInputs(objcProvider.get(ObjcProvider.HEADER))
-              .addInputs(compilationArtifacts.getPchFile().asSet())
-              .addTransitiveInputs(objcProvider.get(ObjcProvider.STATIC_FRAMEWORK_FILE))
-              .addTransitiveInputs(objcProvider.get(ObjcProvider.DYNAMIC_FRAMEWORK_FILE))
-              .build(ruleContext));
     }
   }
 
+  private void registerHeaderScanningAction(
+      ObjcProvider objcProvider,
+      CompilationArtifacts compilationArtifacts,
+      ImmutableList<String> args,
+      List<ObjcHeaderThinningInfo> infos) {
+    SpawnAction.Builder builder =
+        new SpawnAction.Builder()
+            .setMnemonic("ObjcHeaderScanning")
+            .setExecutable(getHeaderThinningToolExecutable())
+            .addInputs(
+                ruleContext
+                    .getPrerequisiteArtifacts(ObjcRuleClasses.APPLE_SDK_ATTRIBUTE, Mode.TARGET)
+                    .list());
+    CustomCommandLine.Builder cmdLine =
+        CustomCommandLine.builder()
+            .add("--arch")
+            .add(appleConfiguration.getSingleArchitecture().toLowerCase())
+            .add("--platform")
+            .add(appleConfiguration.getSingleArchPlatform().getLowerCaseNameInPlist())
+            .add("--sdk_version")
+            .add(
+                appleConfiguration
+                    .getSdkVersionForPlatform(appleConfiguration.getSingleArchPlatform())
+                    .toStringWithMinimumComponents(2))
+            .add("--xcode_version")
+            .add(appleConfiguration.getXcodeVersion().toStringWithMinimumComponents(2))
+            .add("--");
+    for (ObjcHeaderThinningInfo info : infos) {
+      cmdLine.addJoinPaths(
+          ":",
+          Lists.newArrayList(info.sourceFile.getExecPath(), info.headersListFile.getExecPath()));
+      builder.addInput(info.sourceFile).addOutput(info.headersListFile);
+    }
+    ruleContext.registerAction(
+        builder
+            .setCommandLine(cmdLine.add("--").add(args).build())
+            .addInputs(compilationArtifacts.getPrivateHdrs())
+            .addTransitiveInputs(attributes.hdrs())
+            .addTransitiveInputs(objcProvider.get(ObjcProvider.HEADER))
+            .addInputs(compilationArtifacts.getPchFile().asSet())
+            .addTransitiveInputs(objcProvider.get(ObjcProvider.STATIC_FRAMEWORK_FILE))
+            .addTransitiveInputs(objcProvider.get(ObjcProvider.DYNAMIC_FRAMEWORK_FILE))
+            .build(ruleContext));
+  }
+
   /**
    * Groups {@link ObjcHeaderThinningInfo} objects based on the command line arguments of the
    * ObjcCompile action.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcCommandLineOptions.java b/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcCommandLineOptions.java
index d771cc0c86..514e2a0c9c 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcCommandLineOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcCommandLineOptions.java
@@ -293,6 +293,14 @@ public class ObjcCommandLineOptions extends FragmentOptions {
   public boolean experimentalObjcHeaderThinning;
 
   @Option(
+    name = "objc_header_thinning_partition_size",
+    defaultValue = "120",
+    category = "undocumented",
+    help = "The maximum number of source files to process within in each header scanning action."
+  )
+  public int objcHeaderThinningPartitionSize;
+
+  @Option(
     name = "objc_header_scanner_tool",
     defaultValue = "@bazel_tools//tools/objc:header_scanner",
     category = "undocumented",
diff --git a/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcConfiguration.java
index da9fd874b7..6f3925e6c6 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/objc/ObjcConfiguration.java
@@ -81,6 +81,7 @@ public class ObjcConfiguration extends BuildConfiguration.Fragment {
   private final boolean enableAppleBinaryNativeProtos;
   private final HeaderDiscovery.DotdPruningMode dotdPruningPlan;
   private final boolean experimentalHeaderThinning;
+  private final int objcHeaderThinningPartitionSize;
   private final Label objcHeaderScannerTool;
   private final Label appleSdk;
   private final boolean generateXcodeProject;
@@ -123,6 +124,7 @@ public class ObjcConfiguration extends BuildConfiguration.Fragment {
             ? HeaderDiscovery.DotdPruningMode.USE
             : HeaderDiscovery.DotdPruningMode.DO_NOT_USE;
     this.experimentalHeaderThinning = objcOptions.experimentalObjcHeaderThinning;
+    this.objcHeaderThinningPartitionSize = objcOptions.objcHeaderThinningPartitionSize;
     this.objcHeaderScannerTool = objcOptions.objcHeaderScannerTool;
     this.appleSdk = objcOptions.appleSdk;
     this.generateXcodeProject = objcOptions.generateXcodeProject;
@@ -362,6 +364,11 @@ public class ObjcConfiguration extends BuildConfiguration.Fragment {
     return experimentalHeaderThinning;
   }
 
+  /** Returns the max number of source files to add to each header scanning action. */
+  public int objcHeaderThinningPartitionSize() {
+    return objcHeaderThinningPartitionSize;
+  }
+
   /** Returns the label for the ObjC header scanner tool. */
   public Label getObjcHeaderScannerTool() {
     return objcHeaderScannerTool;