Skip to content

Commit

Permalink
support persistent worker for aar extractors
Browse files Browse the repository at this point in the history
This PR supports persistent worker for Android aar extractors, it'll significantly improve the aar extractors performance especially in the first time compilation without cache.
It's not enable by default, can use flag `build --experimental_persistent_aar_extractor=true` to enable it.

using worker:
![image](https://github.com/bazelbuild/bazel/assets/10076965/804f8167-4a52-4321-ba4a-9d1f1dd57644)
not using worker:
![image](https://github.com/bazelbuild/bazel/assets/10076965/26ef4047-4ff7-45f3-9767-ce48418aa089)

Closes #18496.

PiperOrigin-RevId: 557923249
Change-Id: Ica6b6e0a6e3f1dcb68bb9aba941d538888ca013d
  • Loading branch information
ted-xie authored and copybara-github committed Aug 17, 2023
1 parent 3abbf20 commit 74bc046
Showing 11 changed files with 219 additions and 35 deletions.
Original file line number Diff line number Diff line change
@@ -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;
@@ -49,6 +53,7 @@
import com.google.devtools.build.lib.rules.java.ProguardSpecProvider;
import com.google.devtools.build.lib.starlarkbuildapi.android.DataBindingV2ProviderApi;
import com.google.devtools.build.lib.vfs.PathFragment;
import java.util.List;
import javax.annotation.Nullable;

/**
@@ -65,6 +70,9 @@ 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;

@@ -272,13 +280,41 @@ 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();
}

/**
* Returns either the worker param file or null depending on whether persistent worker mode is on.
*/
@Nullable
private static ParamFileInfo getParamFileInfo(RuleContext ruleContext) {
return isPersistentAarExtractor(ruleContext) ? WORKERS_FORCED_PARAM_FILE_INFO : null;
}

/** Adds worker-related execution requirements if worker mode enabled. */
private static void modifyExecutionInfo(
RuleContext ruleContext, SpawnAction.Builder actionBuilder) {
if (isPersistentAarExtractor(ruleContext)) {
ImmutableMap.Builder<String, String> executionInfo = ImmutableMap.builder();
executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
executionInfo.put(ExecutionRequirements.REQUIRES_WORKER_PROTOCOL, "json");
actionBuilder.setExecutionInfo(executionInfo.buildOrThrow());
}
}

/**
* 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();
ParamFileInfo paramFileInfo = getParamFileInfo(ruleContext);
modifyExecutionInfo(ruleContext, actionBuilder);
return actionBuilder
.useDefaultShellEnvironment()
.setExecutable(
ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_EMBEDDED_PROGUARD_EXTACTOR))
@@ -290,7 +326,8 @@ private static SpawnAction createAarEmbeddedProguardExtractorActions(
CustomCommandLine.builder()
.addExecPath("--input_aar", aar)
.addExecPath("--output_proguard_file", proguardSpecArtifact)
.build())
.build(),
paramFileInfo)
.build(ruleContext);
}

@@ -352,7 +389,10 @@ private static SpawnAction createAarResourcesExtractorActions(
Artifact databindingBrFiles,
Artifact databindingSetterStoreFiles) {

return new SpawnAction.Builder()
SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
ParamFileInfo paramFileInfo = getParamFileInfo(ruleContext);
modifyExecutionInfo(ruleContext, actionBuilder);
return actionBuilder
.useDefaultShellEnvironment()
.setExecutable(
ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_RESOURCES_EXTRACTOR))
@@ -370,7 +410,8 @@ private static SpawnAction createAarResourcesExtractorActions(
.addExecPath("--output_assets_dir", assetsDir)
.addExecPath("--output_databinding_br_dir", databindingBrFiles)
.addExecPath("--output_databinding_setter_store_dir", databindingSetterStoreFiles)
.build())
.build(),
paramFileInfo)
.build(ruleContext);
}

@@ -379,7 +420,10 @@ private static SpawnAction createAarEmbeddedJarsExtractorActions(
Artifact aar,
Artifact jarsTreeArtifact,
Artifact singleJarParamFile) {
return new SpawnAction.Builder()
SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
ParamFileInfo paramFileInfo = getParamFileInfo(ruleContext);
modifyExecutionInfo(ruleContext, actionBuilder);
return actionBuilder
.useDefaultShellEnvironment()
.setExecutable(
ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_EMBEDDED_JARS_EXTACTOR))
@@ -393,7 +437,8 @@ private static SpawnAction createAarEmbeddedJarsExtractorActions(
.addExecPath("--input_aar", aar)
.addExecPath("--output_dir", jarsTreeArtifact)
.addExecPath("--output_singlejar_param_file", singleJarParamFile)
.build())
.build(),
paramFileInfo)
.build(ruleContext);
}

@@ -420,31 +465,33 @@ private static SpawnAction createAarJarsMergingActions(

private static SpawnAction createAarNativeLibsFilterActions(
RuleContext ruleContext, Artifact aar, Artifact outputZip) {
SpawnAction.Builder actionBuilder =
new SpawnAction.Builder()
.useDefaultShellEnvironment()
.setExecutable(
ruleContext.getExecutablePrerequisite(
AarImportBaseRule.AAR_NATIVE_LIBS_ZIP_CREATOR))
.setMnemonic("AarNativeLibsFilter")
.setProgressMessage("Filtering AAR native libs by architecture for %{label}")
.addInput(aar)
.addOutput(outputZip)
.addCommandLine(
CustomCommandLine.builder()
.addExecPath("--input_aar", aar)
.add("--cpu", ruleContext.getConfiguration().getCpu())
.addExecPath("--output_zip", outputZip)
.build());
return actionBuilder.build(ruleContext);
SpawnAction.Builder actionBuilder = new SpawnAction.Builder();
ParamFileInfo paramFileInfo = getParamFileInfo(ruleContext);
modifyExecutionInfo(ruleContext, actionBuilder);
return actionBuilder
.useDefaultShellEnvironment()
.setExecutable(
ruleContext.getExecutablePrerequisite(AarImportBaseRule.AAR_NATIVE_LIBS_ZIP_CREATOR))
.setMnemonic("AarNativeLibsFilter")
.setProgressMessage("Filtering AAR native libs by architecture for %{label}")
.addInput(aar)
.addOutput(outputZip)
.addCommandLine(
CustomCommandLine.builder()
.addExecPath("--input_aar", aar)
.add("--cpu", ruleContext.getConfiguration().getCpu())
.addExecPath("--output_zip", outputZip)
.build(),
paramFileInfo)
.build(ruleContext);
}

private static DataBindingV2Provider createDatabindingProvider(
RuleContext ruleContext,
SpecialArtifact databindingBrFiles,
SpecialArtifact databindingSetterStoreFiles) {

Iterable<? extends DataBindingV2ProviderApi<Artifact>> databindingProvidersFromDeps =
List<? extends DataBindingV2ProviderApi<Artifact>> databindingProvidersFromDeps =
ruleContext.getPrerequisites("deps", DataBindingV2Provider.PROVIDER);

Iterable<? extends DataBindingV2ProviderApi<Artifact>> databindingProvidersFromExports =
Original file line number Diff line number Diff line change
@@ -825,6 +825,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",
@@ -1132,6 +1142,7 @@ public FragmentOptions getExec() {
exec.manifestMergerOrder = manifestMergerOrder;
exec.oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest =
oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest;
exec.persistentAarExtractor = persistentAarExtractor;
exec.persistentBusyboxTools = persistentBusyboxTools;
exec.persistentMultiplexBusyboxTools = persistentMultiplexBusyboxTools;
exec.disableNativeAndroidRules = disableNativeAndroidRules;
@@ -1177,6 +1188,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;
@@ -1239,6 +1251,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;
Original file line number Diff line number Diff line change
@@ -231,6 +231,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,
34 changes: 34 additions & 0 deletions src/test/shell/bazel/android/aar_integration_test.sh
Original file line number Diff line number Diff line change
@@ -147,4 +147,38 @@ EOF
assert_one_of $apk_contents "lib/armeabi-v7a/libapp.so"
}

function test_aar_extractor_worker() {
create_new_workspace
setup_android_sdk_support
cat > AndroidManifest.xml <<EOF
<manifest package="com.example"/>
EOF
mkdir -p res/layout
cat > res/layout/mylayout.xml <<EOF
<?xml version="1.0" encoding="utf-8"?>
<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android" />
EOF
mkdir assets
echo "some asset" > assets/a
zip example.aar AndroidManifest.xml res/layout/mylayout.xml assets/a
cat > BUILD <<EOF
aar_import(
name = "example",
aar = "example.aar",
)
android_binary(
name = "app",
custom_package = "com.example",
manifest = "AndroidManifest.xml",
deps = [":example"],
)
EOF

bazel clean
bazel build --experimental_persistent_aar_extractor :app || fail "build failed"
apk_contents="$(zipinfo -1 bazel-bin/app.apk)"
assert_one_of $apk_contents "assets/a"
assert_one_of $apk_contents "res/layout/mylayout.xml"
}

run_suite "aar_import integration tests"
11 changes: 11 additions & 0 deletions tools/android/BUILD
Original file line number Diff line number Diff line change
@@ -86,6 +86,7 @@ py_binary(
"aar_native_libs_zip_creator.py",
],
deps = [
":json_worker_wrapper",
":junction_lib",
"//third_party/py/abseil",
],
@@ -121,6 +122,7 @@ py_binary(
name = "aar_embedded_proguard_extractor",
srcs = ["aar_embedded_proguard_extractor.py"],
deps = [
":json_worker_wrapper",
":junction_lib",
"//third_party/py/abseil",
],
@@ -136,6 +138,7 @@ py_binary(
name = "aar_embedded_jars_extractor",
srcs = ["aar_embedded_jars_extractor.py"],
deps = [
":json_worker_wrapper",
":junction_lib",
"//third_party/py/abseil",
],
@@ -151,6 +154,7 @@ py_binary(
name = "aar_resources_extractor",
srcs = ["aar_resources_extractor.py"],
deps = [
":json_worker_wrapper",
":junction_lib",
"//third_party/py/abseil",
],
@@ -193,6 +197,13 @@ py_library(
visibility = ["//visibility:private"],
)

py_library(
name = "json_worker_wrapper",
srcs = ["json_worker_wrapper.py"],
srcs_version = "PY3",
visibility = ["//visibility:private"],
)

py_test(
name = "junction_test",
srcs = select({
11 changes: 11 additions & 0 deletions tools/android/BUILD.tools
Original file line number Diff line number Diff line change
@@ -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,13 @@ py_library(
visibility = ["//visibility:private"],
)

py_library(
name = "json_worker_wrapper",
srcs = ["json_worker_wrapper.py"],
visibility = ["//visibility:private"],
srcs_version = "PY3",
)

alias(
name = "android_runtest",
actual = "fail.sh",
Loading

0 comments on commit 74bc046

Please sign in to comment.