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
ted-xie authored and Copybara-Service committed Aug 17, 2023
1 parent 3abbf20 commit 74bc046
Show file tree
Hide file tree
Showing 11 changed files with 219 additions and 35 deletions.
95 changes: 71 additions & 24 deletions src/main/java/com/google/devtools/build/lib/rules/android/AarImport.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

/**
Expand All @@ -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;

Expand Down Expand Up @@ -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))
Expand All @@ -290,7 +326,8 @@ private static SpawnAction createAarEmbeddedProguardExtractorActions(
CustomCommandLine.builder()
.addExecPath("--input_aar", aar)
.addExecPath("--output_proguard_file", proguardSpecArtifact)
.build())
.build(),
paramFileInfo)
.build(ruleContext);
}

Expand Down Expand Up @@ -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))
Expand All @@ -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);
}

Expand All @@ -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))
Expand All @@ -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);
}

Expand All @@ -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 =
Expand Down
18 changes: 18 additions & 0 deletions src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -1481,6 +1494,11 @@ public boolean useDataBindingAndroidX() {
return dataBindingAndroidX;
}

@Override
public boolean persistentAarExtractor() {
return persistentAarExtractor;
}

@Override
public boolean persistentBusyboxTools() {
return persistentBusyboxTools;
Expand Down
7 changes: 7 additions & 0 deletions .../java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
34 changes: 34 additions & 0 deletions src/test/shell/bazel/android/aar_integration_test.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ py_binary(
"aar_native_libs_zip_creator.py",
],
deps = [
":json_worker_wrapper",
":junction_lib",
"//third_party/py/abseil",
],
Expand Down Expand Up @@ -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",
],
Expand All @@ -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",
],
Expand All @@ -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",
],
Expand Down Expand Up @@ -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({
Expand Down
11 changes: 11 additions & 0 deletions tools/android/BUILD.tools
View file
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,7 @@ py_binary(
deps = [
":junction_lib",
"//third_party/py/abseil",
":json_worker_wrapper",
],
)

Expand All @@ -371,6 +372,7 @@ py_binary(
deps = [
":junction_lib",
"//third_party/py/abseil",
":json_worker_wrapper",
],
)

Expand All @@ -381,6 +383,7 @@ py_binary(
deps = [
":junction_lib",
"//third_party/py/abseil",
":json_worker_wrapper",
],
)

Expand All @@ -391,6 +394,7 @@ py_binary(
deps = [
":junction_lib",
"//third_party/py/abseil",
":json_worker_wrapper",
],
)

Expand All @@ -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",
Expand Down
Loading

0 comments on commit 74bc046

Please sign in to comment.