support persistent worker for aar extractors

bazelbuild · May 31, 2023 · aa06c30 · aa06c30
1 parent 9181f32
commit aa06c30
Show file tree

Hide file tree

Showing 11 changed files with 173 additions and 23 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java b/src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
@@ -14,10 +14,14 @@
 package com.google.devtools.build.lib.rules.android;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
+import com.google.devtools.build.lib.actions.ExecutionRequirements;
 import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
+import com.google.devtools.build.lib.actions.ParamFileInfo;
+import com.google.devtools.build.lib.actions.ParameterFile;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.FileProvider;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
@@ -65,6 +69,11 @@ public class AarImport implements RuleConfiguredTargetFactory {
   private static final String MERGED_JAR = "classes_and_libs_merged.jar";
   private static final String PROGUARD_SPEC = "proguard.txt";
 
+  private static final ParamFileInfo WORKERS_FORCED_PARAM_FILE_INFO =
+      ParamFileInfo.builder(ParameterFile.ParameterFileType.UNQUOTED)
+          .setUseAlways(true)
+              .build();
+
   private final JavaSemantics javaSemantics;
   private final AndroidSemantics androidSemantics;
 
@@ -274,13 +283,27 @@ private NestedSet<Artifact> extractProguardSpecs(RuleContext ruleContext, Artifa
     return builder.addTransitive(proguardSpecs).add(proguardSpecArtifact).build();
   }
 
+  private static boolean isPersistentAarExtractor(RuleContext ruleContext) {
+    AndroidConfiguration androidConfig =
+        ruleContext.getConfiguration().getFragment(AndroidConfiguration.class);
+    return androidConfig.persistentAarExtractor();
+  }
+
   /**
    * Creates action to extract embedded Proguard.txt from an AAR. If the file is not found, an empty
    * file will be created
    */
   private static SpawnAction createAarEmbeddedProguardExtractorActions(
       RuleContext ruleContext, Artifact aar, Artifact proguardSpecArtifact) {
-    return new SpawnAction.Builder()
+    SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
+    boolean isPersistentAarExtractor = isPersistentAarExtractor(ruleContext);
+    if (isPersistentAarExtractor) {
+      ImmutableMap.Builder<String, String> executionInfo = ImmutableMap.builder();
+      executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+      executionInfo.put(ExecutionRequirements.REQUIRES_WORKER_PROTOCOL, "json");
+        actionBuilder.setExecutionInfo(executionInfo.build());
+    }
+    return actionBuilder
         .useDefaultShellEnvironment()
         .setExecutable(
             ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_EMBEDDED_PROGUARD_EXTACTOR))
@@ -292,7 +315,7 @@ private static SpawnAction createAarEmbeddedProguardExtractorActions(
             CustomCommandLine.builder()
                 .addExecPath("--input_aar", aar)
                 .addExecPath("--output_proguard_file", proguardSpecArtifact)
-                .build())
+                .build(), isPersistentAarExtractor ? WORKERS_FORCED_PARAM_FILE_INFO : null)
         .build(ruleContext);
   }
 
@@ -354,7 +377,15 @@ private static SpawnAction createAarResourcesExtractorActions(
       Artifact databindingBrFiles,
       Artifact databindingSetterStoreFiles) {
 
-    return new SpawnAction.Builder()
+    SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
+    boolean isPersistentAarExtractor = isPersistentAarExtractor(ruleContext);
+    if (isPersistentAarExtractor) {
+      ImmutableMap.Builder<String, String> executionInfo = ImmutableMap.builder();
+      executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+      executionInfo.put(ExecutionRequirements.REQUIRES_WORKER_PROTOCOL, "json");
+        actionBuilder.setExecutionInfo(executionInfo.build());
+    }
+    return actionBuilder
         .useDefaultShellEnvironment()
         .setExecutable(
             ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_RESOURCES_EXTRACTOR))
@@ -372,7 +403,7 @@ private static SpawnAction createAarResourcesExtractorActions(
                 .addExecPath("--output_assets_dir", assetsDir)
                 .addExecPath("--output_databinding_br_dir", databindingBrFiles)
                 .addExecPath("--output_databinding_setter_store_dir", databindingSetterStoreFiles)
-                .build())
+                .build(), isPersistentAarExtractor ? WORKERS_FORCED_PARAM_FILE_INFO : null)
         .build(ruleContext);
   }
 
@@ -381,7 +412,15 @@ private static SpawnAction createAarEmbeddedJarsExtractorActions(
       Artifact aar,
       Artifact jarsTreeArtifact,
       Artifact singleJarParamFile) {
-    return new SpawnAction.Builder()
+    SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
+    boolean isPersistentAarExtractor = isPersistentAarExtractor(ruleContext);
+    if (isPersistentAarExtractor) {
+      ImmutableMap.Builder<String, String> executionInfo = ImmutableMap.builder();
+      executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+      executionInfo.put(ExecutionRequirements.REQUIRES_WORKER_PROTOCOL, "json");
+        actionBuilder.setExecutionInfo(executionInfo.build());
+    }
+    return actionBuilder
         .useDefaultShellEnvironment()
         .setExecutable(
             ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_EMBEDDED_JARS_EXTACTOR))
@@ -395,7 +434,7 @@ private static SpawnAction createAarEmbeddedJarsExtractorActions(
                 .addExecPath("--input_aar", aar)
                 .addExecPath("--output_dir", jarsTreeArtifact)
                 .addExecPath("--output_singlejar_param_file", singleJarParamFile)
-                .build())
+                    .build(), isPersistentAarExtractor ? WORKERS_FORCED_PARAM_FILE_INFO : null)
         .build(ruleContext);
   }
 
@@ -422,8 +461,16 @@ private static SpawnAction createAarJarsMergingActions(
 
   private static SpawnAction createAarNativeLibsFilterActions(
       RuleContext ruleContext, Artifact aar, Artifact outputZip) {
-    SpawnAction.Builder actionBuilder =
-        new SpawnAction.Builder()
+    SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
+    boolean isPersistentAarExtractor = isPersistentAarExtractor(ruleContext);
+    if (isPersistentAarExtractor) {
+      ImmutableMap.Builder<String, String> executionInfo = ImmutableMap.builder();
+      executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+      executionInfo.put(ExecutionRequirements.REQUIRES_WORKER_PROTOCOL, "json");
+        actionBuilder.setExecutionInfo(executionInfo.build());
+    }
+    return
+        actionBuilder
             .useDefaultShellEnvironment()
             .setExecutable(
                 ruleContext.getExecutablePrerequisite(
@@ -437,8 +484,8 @@ private static SpawnAction createAarNativeLibsFilterActions(
                     .addExecPath("--input_aar", aar)
                     .add("--cpu", ruleContext.getConfiguration().getCpu())
                     .addExecPath("--output_zip", outputZip)
-                    .build());
-    return actionBuilder.build(ruleContext);
+                    .build(), isPersistentAarExtractor ? WORKERS_FORCED_PARAM_FILE_INFO : null)
+            .build(ruleContext);
   }
 
   private static DataBindingV2Provider createDatabindingProvider(

diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
@@ -831,6 +831,16 @@ public static class Options extends FragmentOptions {
                 + "transitive classpath, otherwise it uses the deploy jar")
     public boolean oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest;
 
+    @Option(
+        name = "experimental_persistent_aar_extractor",
+        defaultValue = "false",
+        documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
+        effectTags = {
+          OptionEffectTag.EXECUTION,
+        },
+        help = "Enable persistent aar extractor by using workers.")
+    public boolean persistentAarExtractor;
+
     @Option(
         name = "persistent_android_resource_processor",
         defaultValue = "null",
@@ -1128,6 +1138,7 @@ public FragmentOptions getExec() {
       exec.allowAndroidLibraryDepsWithoutSrcs = allowAndroidLibraryDepsWithoutSrcs;
       exec.oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest =
           oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest;
+      exec.persistentAarExtractor = persistentAarExtractor;
       exec.persistentBusyboxTools = persistentBusyboxTools;
       exec.persistentMultiplexBusyboxTools = persistentMultiplexBusyboxTools;
       exec.disableNativeAndroidRules = disableNativeAndroidRules;
@@ -1174,6 +1185,7 @@ public FragmentOptions getExec() {
   private final boolean dataBindingV2;
   private final boolean dataBindingUpdatedArgs;
   private final boolean dataBindingAndroidX;
+  private final boolean persistentAarExtractor;
   private final boolean persistentBusyboxTools;
   private final boolean persistentMultiplexBusyboxTools;
   private final boolean persistentDexDesugar;
@@ -1236,6 +1248,7 @@ public AndroidConfiguration(BuildOptions buildOptions) throws InvalidConfigurati
     this.dataBindingV2 = options.dataBindingV2;
     this.dataBindingUpdatedArgs = options.dataBindingUpdatedArgs;
     this.dataBindingAndroidX = options.dataBindingAndroidX;
+    this.persistentAarExtractor = options.persistentAarExtractor;
     this.persistentBusyboxTools = options.persistentBusyboxTools;
     this.persistentMultiplexBusyboxTools = options.persistentMultiplexBusyboxTools;
     this.persistentDexDesugar = options.persistentDexDesugar;
@@ -1481,6 +1494,11 @@ public boolean useDataBindingAndroidX() {
     return dataBindingAndroidX;
   }
 
+  @Override
+  public boolean persistentAarExtractor() {
+    return persistentAarExtractor;
+  }
+
   @Override
   public boolean persistentBusyboxTools() {
     return persistentBusyboxTools;

diff --git a/.../java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java b/.../java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java
@@ -224,6 +224,13 @@ public interface AndroidConfigurationApi extends StarlarkValue {
       documented = false)
   boolean useDataBindingAndroidX();
 
+  @StarlarkMethod(
+     name = "persistent_aar_extractor",
+     structField = true,
+     doc = "",
+     documented = false)
+  boolean persistentAarExtractor();
+
   @StarlarkMethod(
       name = "persistent_busybox_tools",
       structField = true,

diff --git a/third_party/py/abseil/absl/app.py b/third_party/py/abseil/absl/app.py
@@ -204,7 +204,7 @@ def register_and_parse_flags_with_usage(argv=None):
   return argv
 
 
-def _run_main(main, argv):
+def _run_main(main, argv, exit=True):
   """Calls main, optionally with pdb or profiler."""
   if FLAGS.run_with_pdb:
     sys.exit(pdb.runcall(main, argv))
@@ -222,9 +222,12 @@ def _run_main(main, argv):
     else:
       atexit.register(profiler.print_stats)
     retval = profiler.runcall(main, argv)
-    sys.exit(retval)
+    if exit:
+      sys.exit(retval)
   else:
-    sys.exit(main(argv))
+    retval = main(argv)
+    if exit:
+      sys.exit(retval)
 
 
 def _call_exception_handlers(exception):
@@ -243,7 +246,7 @@ def _call_exception_handlers(exception):
         pass
 
 
-def run(main, argv=None):
+def run(main, argv=None, exit=True):
   """Begins executing the program.
 
   Args:
@@ -259,7 +262,7 @@ def run(main, argv=None):
   try:
     argv = _run_init(sys.argv if argv is None else argv)
     try:
-      _run_main(main, argv)
+      _run_main(main, argv, exit=exit)
     except UsageError as error:
       usage(shorthelp=True, detailed_error=error, exitcode=error.exitcode)
     except:

diff --git a/tools/android/BUILD b/tools/android/BUILD
@@ -88,6 +88,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -123,6 +124,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -138,6 +140,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -153,6 +156,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -191,6 +195,12 @@ py_library(
     visibility = ["//visibility:private"],
 )
 
+py_library(
+    name = "json_worker_wrapper",
+    srcs = ["json_worker_wrapper.py"],
+    visibility = ["//visibility:private"],
+)
+
 py_test(
     name = "junction_test",
     srcs = select({

diff --git a/tools/android/BUILD.tools b/tools/android/BUILD.tools
@@ -352,6 +352,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -371,6 +372,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -381,6 +383,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -391,6 +394,7 @@ py_binary(
     deps = [
         ":junction_lib",
         "//third_party/py/abseil",
+        ":json_worker_wrapper",
     ],
 )
 
@@ -408,6 +412,12 @@ py_library(
     visibility = ["//visibility:private"],
 )
 
+py_library(
+    name = "json_worker_wrapper",
+    srcs = ["json_worker_wrapper.py"],
+    visibility = ["//visibility:private"],
+)
+
 alias(
     name = "android_runtest",
     actual = "fail.sh",

diff --git a/tools/android/aar_embedded_jars_extractor.py b/tools/android/aar_embedded_jars_extractor.py
@@ -29,6 +29,7 @@
 from absl import flags
 
 from tools.android import junction
+from tools.android import json_worker_wrapper
 
 FLAGS = flags.FLAGS
 
@@ -92,5 +93,4 @@ def main(unused_argv):
 
 
 if __name__ == "__main__":
-  FLAGS(sys.argv)
-  app.run(main)
+  json_worker_wrapper.wrap_worker(FLAGS, main, app.run)
diff --git a/tools/android/aar_embedded_proguard_extractor.py b/tools/android/aar_embedded_proguard_extractor.py
@@ -27,6 +27,7 @@
 from absl import flags
 
 from tools.android import junction
+from tools.android import json_worker_wrapper
 
 FLAGS = flags.FLAGS
 
@@ -71,5 +72,4 @@ def main(unused_argv):
 
 
 if __name__ == "__main__":
-  FLAGS(sys.argv)
-  app.run(main)
+  json_worker_wrapper.wrap_worker(FLAGS, main, app.run)
diff --git a/tools/android/aar_native_libs_zip_creator.py b/tools/android/aar_native_libs_zip_creator.py
@@ -30,6 +30,7 @@
 from absl import flags
 
 from tools.android import junction
+from tools.android import json_worker_wrapper
 
 FLAGS = flags.FLAGS
 
@@ -94,5 +95,4 @@ def main(unused_argv):
 
 
 if __name__ == "__main__":
-  FLAGS(sys.argv)
-  app.run(main)
+  json_worker_wrapper.wrap_worker(FLAGS, main, app.run)